如何用TensorFlow算子按指定次数重复张量的各行？

嘿，这个按指定次数重复张量行的需求，在TensorFlow里其实有挺简洁的实现方式，我给你两种靠谱的方案，适配不同版本的TensorFlow：

这是最直观高效的方式，从TensorFlow 2.4开始，tf.repeat支持对指定轴的元素按自定义次数重复，完美匹配你的场景：

import tensorflow as tf

# 定义你的原始张量和重复次数列表
input_tensor = tf.constant([[0, 1, 2], [1, 2, 3], [4, 5, 6]])
repeat_counts = tf.constant([3, 2, 1])

# 沿着行维度（axis=0）按指定次数重复每行
result = tf.repeat(input_tensor, repeat_counts, axis=0)

# 输出结果
print(result.numpy())

运行后你会得到期望的输出：

[[0 1 2]
 [0 1 2]
 [0 1 2]
 [1 2 3]
 [1 2 3]
 [4 5 6]]

简单解释下：axis=0指定我们要操作的是行维度，repeat_counts里的每个数值对应原始张量中对应行的重复次数，一步到位就能得到目标结果。

如果你的TensorFlow版本低于2.4，没法直接用上面的方法，可以通过生成重复索引的方式间接实现：

import tensorflow as tf

input_tensor = tf.constant([[0, 1, 2], [1, 2, 3], [4, 5, 6]])
repeat_counts = tf.constant([3, 2, 1])

# 先生成原始张量的行索引：[0, 1, 2]
row_indices = tf.range(tf.shape(input_tensor)[0])
# 把索引按指定次数重复，得到：[0, 0, 0, 1, 1, 2]
repeated_indices = tf.repeat(row_indices, repeat_counts)

# 根据重复后的索引提取对应行
result = tf.gather(input_tensor, repeated_indices)

print(result.numpy())

这个思路的核心是先构造出重复的行索引，再用tf.gather从原始张量中批量提取对应行，最终结果和第一种方法完全一致，而且兼容更早的TensorFlow版本。

这两种方案都是纯TensorFlow算子实现，没有用Python循环，完全适配张量流的计算逻辑，不管是在 eager 模式还是图模式下都能稳定运行，不会有性能瓶颈。

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