Search code examples
pythonarraysnumpytensorflowvectorization

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).

Solution

  • 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]]]