Python 从Numpy或Tensorflow中的线性阵列矢量化创建对角正方形阵列阵列

我有一个shape

[batch\u size，N]

数组，例如：

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

我需要用shape

[batch\u size，N，N]

创建一个3个索引数组，其中对于每个

批

我有一个

nxn

对角矩阵，其中对角线由相应的

批

元素获取，例如，在本例中，在这个简单的情况下，我要寻找的结果是：

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

如果没有for循环和展开矢量化，如何进行此操作？我猜这是维度的扩展，但我找不到正确的函数来实现这一点。 （我需要它，因为我正在使用tensorflow和numpy进行原型制作）。

您可以使用

---

编辑下图显示了上述溶液的平均执行时间（实线），并将其与@Divakar的（虚线）进行了比较，以确定不同批次大小和不同矩阵大小

我不相信您会有多大的改进，但这只是基于这个简单的度量标准，它与一个具有

解释

np.expand_dims（a，axis=1）

向

a

添加一个新轴，该轴现在将成为

（3，1，2）

ndarray：

array([[[1, 2]],

       [[3, 4]],

       [[5, 6]]])

现在，您可以将此数组与大小

N

标识矩阵相乘，您可以使用以下方法生成该矩阵：

这将产生所需的输出：

a*np.eye(N)

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

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

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

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

方法#1

这是一个矢量化的输入数组，

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

---

我同意这是最简单的解决方案，但我想知道是否有人会通过像

einsum

这样的函数将其高度矢量化并放在一行中。另一个区别是，您的输出是列表，而不是数组，因此不适合进一步的矢量化计算。

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

# 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

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