python - 从 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 制作原型(prototype))。

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