construct 3d diagonal tensor using 2d tensor

Question

Given A = [[1,2],[3,4],[5,6]]. How to use tf.diag() to construct a 3d tensor where each stack is a 2d diagonal matrix using the values from A? So the output should be B = [[[1,0],[0,2]],[[3,0],[0,4]],[[5,0],[0,6]]]. I want to use this as my Gaussian covariance matries.

Answer 1

Ok I figure it out. tf.matrix_diag() does the trick...