Tensorflow 将张量（表示位置）转换为掩码的好方法是什么？

我有一个loc张量，批量大小为2，长度为5。我想基于loc张量生成mask_trg张量。 loc张量表示当前批次的当前位置。 例如，loc[0，2]=2表示掩码_trg[0，2]=[1，1，1，0，0]，因为当前位置是2，所有路径都标记为1

我尝试了多种方法，但不确定生成掩模张量的有效方法（基于向量）是什么

loc = tf.constant(
    [
        [1, 1, 2, 2, 4],
        [1, 2, 4, 4, 4]
    ])

mask_trg = tf.constant(
    [
        [
            [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]
        ],
    ])

我认为y2是你想要的

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)