Python 使用二维张量构造三维对角张量

给定A=[[1,2]，[3,4]，[5,6]]。如何使用tf.diag（）构造一个3d张量，其中每个堆栈都是一个2d对角矩阵，使用来自？所以输出应该是B=[[1,0]，[0,2]，[[3,0]，[0,4]，[[5,0]，[0,6]]。我想用它作为我的高斯协方差矩阵。

好的，我知道了。tf.matrix_diag（）实现了这个技巧