如何在TensorFlow中为矩阵各行设置不同的滚动偏移量？

当然可行，tf.roll确实不支持为不同行设置独立滚动量，但可以通过构造自定义索引的方式实现需求，以下是具体实现步骤：

实现代码 #

import tensorflow as tf

# 定义原始矩阵
mat = tf.constant([[0, 1, 2],
                   [3, 4, 5],
                   [6, 7, 8]])
# 每行对应的滚动量（沿axis=1）
roll_amounts = tf.constant([0, 1, 2])

# 获取矩阵的列数
cols = mat.shape[1]
# 构造每行的基础索引（[[0,1,2], [0,1,2], [0,1,2]]）
base_idx = tf.tile(tf.range(cols)[tf.newaxis, :], [tf.shape(mat)[0], 1])
# 计算偏移后的索引：每行基础索引减去对应滚动量
shifted_idx = base_idx - roll_amounts[:, tf.newaxis]
# 取模得到有效索引（避免负数索引）
idx = tf.math.mod(shifted_idx, cols)
# 根据索引提取元素，得到滚动后的矩阵
result = tf.gather(mat, idx, axis=1)

# 打印结果
print(result)

代码解释 #

1. 构造基础索引：通过tf.tile复制列索引，生成和原矩阵行数一致的基础索引矩阵，确保每行都有完整的列索引序列。
2. 计算偏移索引：将每行的基础索引减去对应滚动量，得到偏移后的索引（可能包含负数）。
3. 修正索引范围：用tf.math.mod对偏移索引取模，将负数索引转换为有效正索引，确保索引在列数范围内。
4. 提取元素：通过tf.gather按修正后的索引从原矩阵中提取元素，完成每行独立滚动的操作。

运行上述代码后，输出结果与预期完全一致：

tf.Tensor(
[[0 1 2]
 [5 3 4]
 [7 8 6]], shape=(3, 3), dtype=int32)

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