如何在TensorFlow 2中实现张量行提取拼接与补零对齐？

问题概述 #

给定2D张量data和不规则嵌套行索引数组row_ids，需要按每个子索引列表提取对应行、拼接成一维张量，再对较短的结果补零使所有行长度一致，且操作需支持反向传播。

解决方案代码 #

import tensorflow as tf

data = tf.constant([
            [300, 301, 302],
            [100, 101, 102],
            [200, 201, 202],
            [120, 121, 122],
            [210, 211, 212],
            [410, 411, 412],
            [110, 111, 112],
            [400, 401, 402],
        ], dtype=tf.float32)

row_ids = [ [ 1, 6, 3 ], [ 2, 4 ], [ 0 ], [ 7, 5] ]

# 1. 将不规则索引转换为RaggedTensor
ragged_row_ids = tf.ragged.constant(row_ids)

# 2. 提取对应行，得到RaggedTensor形式的结果
extracted_ragged = tf.gather(data, ragged_row_ids)

# 3. 计算目标拼接长度：最大子列表行数 × 原数据每行特征数
max_row_count = ragged_row_ids.bounding_shape()[1]
feature_len = tf.shape(data)[1]
target_len = max_row_count * feature_len

# 拼接每行并补零到统一长度
result = extracted_ragged.merge_dims(1, 2).to_tensor(default_value=0.0, shape=(None, target_len))

# 验证结果匹配度
tf.debugging.assert_equal(result, tf.constant([
        [ 100, 101, 102, 110, 111, 112, 120, 121, 122],
        [ 200, 201, 202, 210, 211, 212,   0,   0,   0],
        [ 300, 301, 302,   0,   0,   0,   0,   0,   0],
        [ 400, 401, 402, 410, 411, 412,   0,   0,   0]
    ], dtype=tf.float32))

print("结果匹配：", result.numpy())

步骤详解 #

1. 转换为RaggedTensor：
   使用tf.ragged.constant()将不规则嵌套的row_ids转为RaggedTensor，TensorFlow会自动适配不同长度的子列表，这是处理不规则索引结构的核心。

2. 提取对应行：
   tf.gather()支持直接传入RaggedTensor作为索引，返回3阶RaggedTensor（形状[4, None, 3]），每个样本对应一组可变长度的行数据。

3. 拼接与补零：

   • merge_dims(1, 2)将每个样本的多行（形状[None, 3]）拼接为一维张量（形状[None]），得到2阶RaggedTensor。
   • to_tensor()将RaggedTensor转为普通张量，对较短行补零到target_len，保证所有行长度统一。

反向传播支持 #

所有用到的操作（tf.ragged.constant、tf.gather、merge_dims、to_tensor）均为TensorFlow 2原生可微分操作，完全支持反向传播，不会影响模型梯度计算。

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