tf.diag of a 2D matrix

I have a 2D tensor with various arrays defined as:

x = tf.constant([[0,1,2],[-1,0,1],[-1,-2,0]])

and I want to convert each array to a diagonal matrix as:

diag_x =
[[[ 0,  0,  0],
  [ 0,  1,  0],
  [ 0,  0,  2]],
 [[-1,  0,  0],
  [ 0,  0,  0],
  [ 0,  0,  1]],
 [[-2,  0,  0],
  [ 0, -1,  0],
  [ 0,  0,  0]]]

but if I use the operation tf.diag(x) the output is not this.

I finally found the solution:

tf.matrix_diag(x)