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

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

Solution 2:

Approach #1

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

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

# Get diagonal viewand assign into it input arrayvalues
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

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

Solution 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

Solution 5:

You basically want a function that does the opposite of/reverses np.block(..)

I needed the same thing, so I wrote this little function:

defsplit_blocks(x, m=2, n=2):
    """
    Reverse the action of np.block(..)

    >>> x = np.random.uniform(-1, 1, (2, 18, 20))
    >>> assert (np.block(split_blocks(x, 3, 4)) == x).all()

    :param x: (.., M, N) input matrix to split into blocks
    :param m: number of row splits
    :param n: number of column, splits
    :return:
    """
    x = np.array(x, copy=False)
    nd = x.ndim
    *shape, nr, nc = x.shape
    returnlist(map(list, x.reshape((*shape, m, nr//m, n, nc//n)).transpose(nd-2, nd, *range(nd-2), nd-1, nd+1)))