Vectorized creation of an array of diagonal square arrays from a liner array in Numpy or Tensorflow

问题

I have an array of shape [batch_size, N], for example:

[[1  2]
 [3  4]
 [5  6]]

and I need to create a 3 indices array with shape [batch_size, N, N] where for every batch I have a N x N diagonal matrix, where diagonals are taken by the corresponding batch element, for example in this case, In this simple case, the result I am looking for is:

[
  [[1,0],[0,2]],
  [[3,0],[0,4]],
  [[5,0],[0,6]],
]

How can I make this operation without for loops and exploting vectorization? I guess it is an extension of dimension, but I cannot find the correct function to do this. (I need it as I am working with tensorflow and prototyping with numpy).

回答1:

Try it in tensorflow:

import tensorflow as tf
A = [[1,2],[3 ,4],[5,6]]
B = tf.matrix_diag(A)
print(B.eval(session=tf.Session()))
[[[1 0]
  [0 2]]

 [[3 0]
  [0 4]]

 [[5 0]
  [0 6]]]

回答2:

Approach #1

Here's a vectorized one with np.einsum for input array, a -

# Initialize o/p array
out = np.zeros(a.shape + (a.shape[1],),dtype=a.dtype)

# Get diagonal view and assign into it input array values
diag = np.einsum('ijj->ij',out)
diag[:] = a

Approach #2

Another based on slicing for assignment -

m,n = a.shape
out = np.zeros((m,n,n),dtype=a.dtype)
out.reshape(-1,n**2)[...,::n+1] = a

回答3:

Using np.expand_dims with an element-wise product with np.eye

a = np.array([[1,  2],
              [3,  4],
              [5, 6]])
N = a.shape[1]
a = np.expand_dims(a, axis=1)
a*np.eye(N)

array([[[1., 0.],
       [0., 2.]],

      [[3., 0.],
       [0., 4.]],

      [[5., 0.],
       [0., 6.]]])

Explanation

np.expand_dims(a, axis=1) adds a new axis to a, which will now be a (3, 1, 2) ndarray:

array([[[1, 2]],

       [[3, 4]],

       [[5, 6]]])

You can now multiply this array with a size N identity matrix, which you can generate with np.eye:

np.eye(N)
array([[1., 0.],
       [0., 1.]])

Which will yield the desired output:

a*np.eye(N)

array([[[1., 0.],
        [0., 2.]],

       [[3., 0.],
        [0., 4.]],

       [[5., 0.],
        [0., 6.]]])

回答4:

Yu can use numpy.diag

m = [[1, 2],
 [3, 4],
 [5, 6]]

[np.diag(b) for b in m]

EDIT The following plot shows the average execution time for the solution above (solid line), and compared it against @Divakar's (dashed line) for different batch-sizes and different matrix sizes

I don't believe you get much of an improvement, but this is just based on this simple metric

来源：https://stackoverflow.com/questions/53741481/vectorized-creation-of-an-array-of-diagonal-square-arrays-from-a-liner-array-in

标签

python

arrays

numpy

tensorflow

vectorization