tf.linalg.set_diag