What is the simplest way to create an upper triangle tensor with Tensorflow?

Similar to this question, I want to convert this tensor

tensor = tf.ones((5, 5))
tf.Tensor(
[[1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1.]], shape=(5, 5), dtype=float32)

to an upper triangular tensor with ones and zeros everywhere else:

tf.Tensor(
[[1. 1. 1. 1. 1.]
 [0. 1. 1. 1. 1.]
 [0. 0. 1. 1. 1.]
 [0. 0. 0. 1. 1.]
 [0. 0. 0. 0. 1.]], shape=(5, 5), dtype=float32)

any ideas?

>Solution :

You can use tf.linalg.band_part:

>>> tensor = tf.ones((5, 5))
>>> tf.linalg.band_part(tensor, 0, -1)
<tf.Tensor: shape=(5, 5), dtype=float32, numpy=
array([[1., 1., 1., 1., 1.],
       [0., 1., 1., 1., 1.],
       [0., 0., 1., 1., 1.],
       [0., 0., 0., 1., 1.],
       [0., 0., 0., 0., 1.]], dtype=float32)>

Note that you can use that function to extract an arbitrary number of diagonals from the tensor. Useful cases include:

tf.linalg.band_part(input, 0, -1) ==> Upper triangular part.
tf.linalg.band_part(input, -1, 0) ==> Lower triangular part.
tf.linalg.band_part(input, 0, 0) ==> Diagonal.