从Numpy或Tensorflow中的线性数组矢量化创建对角正方形数组
问题内容：
我有一个形状数组[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进行原型制作）。
问题答案：
在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]]]
