从Numpy或Tensorflow中的线性数组对角对角正方形数组进行向量化创建

Question

我有一个形状为[batch_size, N]的数组，例如：

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

，我需要创建一个形状为[batch_size, N, N]的3个索引数组，其中每个batch都有一个N x N对角矩阵，对角线由相应的batch取元素，例如在这种情况下，在这种简单情况下，我正在寻找的结果是：

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

如何在没有for循环和展开矢量化的情况下进行此操作？我想这是维度的延伸，但是我找不到正确的函数来做到这一点。 （我在使用tensorflow和使用numpy进行原型设计时需要它。）

Answer 1

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

Answer 2

方法1

这是向量数组，其中np.einsum用于输入数组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

方法2

另一个基于切片的分配-

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

Answer 3

您可以使用numpy.diag

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

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

---

编辑下图显示了上述解决方案的平均执行时间（实线），并将其与@Divakar的解决方案（虚线）进行了比较（针对不同的批处理大小和不同的矩阵大小）

我不相信您会得到很多改进，但这只是基于此简单指标

Answer 4

将np.expand_dims与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.]]])

说明

np.expand_dims(a, axis=1)向a添加了一个新轴，该轴现在将是(3, 1, 2) ndarray：

array([[[1, 2]],

       [[3, 4]],

       [[5, 6]]])

您现在可以将此数组与大小为N的单位矩阵相乘，可以使用np.eye生成该矩阵：

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

这将产生所需的输出：

a*np.eye(N)

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

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

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