Как правильно перенести тензор (с указанием местоположения) в маску?

y = tf.one_hot(ptr, 5)
print(y)
y2 = tf.cumsum(y, axis=-1, reverse=True)
print(y2)

tf.Tensor(
[[[0. 1. 0. 0. 0.]
  [0. 1. 0. 0. 0.]
  [0. 0. 1. 0. 0.]
  [0. 0. 1. 0. 0.]
  [0. 0. 0. 0. 1.]]

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

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

Я думаю, что y2 - это то, что вы хотите.