如何为TensorFlow Ragged Tensor添加基于起始值的前导填充

给定如下Ragged Tensor：

[[1],
 [4, 2],
 [1, 2, 3]]

我们需要将其填充为固定最后维度长度的张量，且前导填充值为对应子张量的第一个元素，期望输出：

[[1 1 1],
 [4 4 2],
 [1 2 3]]

实现方案（基于TensorFlow） #

以下方案可直接扩展到更大规模的Ragged Tensor：

import tensorflow as tf

# 定义输入Ragged Tensor
rt = tf.ragged.constant([[1], [4, 2], [1, 2, 3]])
# 目标最后维度长度
target_len = 3

# 1. 获取每个子张量的第一个元素
first_elements = rt[:, 0]
# 2. 计算每个子张量需要填充的长度
pad_lengths = target_len - rt.row_lengths()
# 3. 生成每个子张量对应的填充序列
pad_sequences = tf.repeat(tf.expand_dims(first_elements, axis=1), pad_lengths, axis=1)
# 4. 将填充序列与原Ragged Tensor拼接，再转换为普通张量
result = tf.concat([pad_sequences, rt], axis=1).to_tensor()

print(result)

代码说明 #

• rt.row_lengths()：获取Ragged Tensor每个子张量的长度，用于计算需要填充的元素个数
• tf.repeat(..., pad_lengths, axis=1)：根据需要填充的长度，将每个起始元素重复对应次数，生成前导填充序列
• tf.concat([pad_sequences, rt], axis=1)：将填充序列拼接到原张量的前面，最后转换为普通张量得到结果

运行后输出：

tf.Tensor(
[[1 1 1]
 [4 4 2]
 [1 2 3]], shape=(3, 3), dtype=int32)

内容的提问来源于stack exchange，提问作者d84_n1nj4