如何在TensorFlow 1.x中按不同次数重复张量行？

给定形状为[?, dim]的张量x，以及指定每行重复次数的形状为[?, 1]的张量rep_nums，以下是兼容TensorFlow 1.x的实现方案：

方法一：基于索引掩码与tf.gather #

import tensorflow as tf

# 示例输入张量
x = tf.constant([[ 0,  1,  2,  3,  4],
                 [ 5,  6,  7,  8,  9],
                 [10, 11, 12, 13, 14],
                 [15, 16, 17, 18, 19],
                 [20, 21, 22, 23, 24]])
rep_nums = tf.constant([[1],[2],[1],[3],[1]])

# 1. 将重复次数张量转为一维
rep_nums_flat = tf.squeeze(rep_nums, axis=1)
# 2. 生成原始行索引序列
row_indices = tf.range(tf.shape(x)[0])
# 3. 生成重复后的索引
max_rep = tf.reduce_max(rep_nums_flat)
# 创建掩码标记每个索引的重复位置
mask = tf.sequence_mask(rep_nums_flat, max_rep)
# 展平掩码与索引，过滤得到重复索引
expanded_indices = tf.tile(tf.expand_dims(row_indices, 1), [1, max_rep])
flat_indices = tf.reshape(expanded_indices, [-1])
repeated_indices = tf.boolean_mask(flat_indices, tf.reshape(mask, [-1]))
# 4. 根据索引提取行
result = tf.gather(x, repeated_indices)

# 运行验证
with tf.Session() as sess:
    print(sess.run(result))

方法二：基于tf.map_fn的逐行处理 #

import tensorflow as tf

# 示例输入张量
x = tf.constant([[ 0,  1,  2,  3,  4],
                 [ 5,  6,  7,  8,  9],
                 [10, 11, 12, 13, 14],
                 [15, 16, 17, 18, 19],
                 [20, 21, 22, 23, 24]])
rep_nums = tf.constant([[1],[2],[1],[3],[1]])

# 逐行执行tile操作，再拼接结果
tiled_rows = tf.map_fn(
    lambda args: tf.tile(tf.expand_dims(args[0], 0), [args[1][0], 1]),
    (x, rep_nums),
    dtype=tf.int32
)
result = tf.concat(tiled_rows, axis=0)

# 运行验证
with tf.Session() as sess:
    print(sess.run(result))

两种方法都能输出你期望的结果：

[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]
 [15 16 17 18 19]
 [15 16 17 18 19]
 [20 21 22 23 24]]

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