TensorFlow中tf.vectorized_map与tf.roll异常及批量roll实现问询

问题背景 #

我们需要实现批量版本的tf.roll：对批量张量x中的每个样本x[i]，沿最后一个轴按各自的偏移量shifts[i]进行移位。示例输入：

shifts = tf.constant([3,2,1,0,-1])
x = tf.repeat(tf.range(5)[None], repeats=shifts.shape[0], axis=0)
# x的输出：
# [[0 1 2 3 4]
#  [0 1 2 3 4]
#  [0 1 2 3 4]
#  [0 1 2 3 4]
#  [0 1 2 3 4]]

预期输出：

[[2 3 4 0 1]
  [3 4 0 1 2]
  [4 0 1 2 3]
  [0 1 2 3 4]
  [1 2 3 4 0]]

但在TensorFlow 2.9.2和2.10版本中，使用tf.vectorized_map结合tf.roll得到不符合预期的结果：

y = tf.vectorized_map(
    lambda x: tf.roll(x[0], shift=x[1], axis=-1),
    elems=[x, shifts],
)
# y的输出：
# [[3 4 0 1 2]
#  [4 0 1 2 3]
#  [0 1 2 3 4]
#  [1 2 3 4 0]
#  [0 1 2 3 4]]

替换为tf.map_fn则能得到预期输出，且tf.vectorized_map处理加法等简单运算时结果正常，说明tf.vectorized_map与tf.roll结合存在异常。

问题解答 #

1. 为何tf.vectorized_map与tf.roll结合会出现异常？ #

tf.vectorized_map的核心是将逐样本的函数逻辑转换为向量化的批量运算图，而非逐样本循环执行。在TensorFlow 2.9.2和2.10版本中，tf.vectorized_map对tf.roll的向量化转换存在bug：当shift参数是与批量维度对齐的标量张量时，转换过程中出现了shift值的偏移或维度匹配错误，导致实际应用的偏移量与传入的shifts数组不符。

而加法这类简单运算的向量化逻辑更直接，不会触发该bug；tf.map_fn是逐样本循环执行函数，不存在向量化转换的逻辑，因此能得到正确结果。

2. 批量版本的tf.roll的推荐实现方式？ #

推荐以下三种实现方式，按性能和场景选择：

方式一：使用tf.map_fn（简单可靠）

这是最直接的实现方式，逻辑清晰不易出错，适合中小规模批量数据：

y1 = tf.map_fn(
    lambda x: tf.roll(x[0], shift=x[1], axis=-1),
    elems=[x, shifts],
    fn_output_signature=x.dtype
)

方式二：手动构造索引实现纯向量化操作（性能最优）

通过构造每个样本的目标索引，使用tf.gather实现批量移位，无循环开销，适合大规模批量数据：

batch_size = tf.shape(x)[0]
seq_len = tf.shape(x)[1]

# 生成每个样本的基础索引
base_indices = tf.tile(tf.range(seq_len)[None], [batch_size, 1])
# 计算每个位置的目标索引：(基础索引 - 偏移量) 对序列长度取模
target_indices = tf.math.floormod(base_indices - tf.expand_dims(shifts, 1), seq_len)
# 按索引收集元素得到批量移位结果
batch_roll = tf.gather(x, target_indices, batch_dims=1)

方式三：升级TensorFlow版本

如果项目允许，升级到TensorFlow 2.11及以上版本，该版本已修复tf.vectorized_map与tf.roll结合的bug，可正常使用tf.vectorized_map实现需求。

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