考虑关于map_fn的tf文档：

"map_fn会将fn使用的运算应用于elems的每个元素，从而导致O（elems. shape [0]）个总运算。由于map_fn可以并行处理元素，这在一定程度上有所缓解。但是，使用map_fn表示的转换通常仍低于使用矢量化运算表示的等效转换。"

在给定输入Tensor的情况下，可以尝试以下方法：

input_tensor = tf.range([10])

# <tf.Tensor: shape=(10,), dtype=int32, numpy=array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)>

转换为方阵：

res = tf.repeat(tf.expand_dims(input_tensor, 0), input_tensor.shape[0], axis = 0)

  # array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]], dtype=int32)>

然后将map_fn应用于该Tensor，包括输入中的具有负值的范围向量：

elements = tf.range(10, dtype=tf.int32) * -1
w,_ = tf.map_fn(lambda x: (tf.roll(x[0], x[1], axis=0), x[1]), (res, elements), dtype=(tf.int32, tf.int32))

这会将元素排列为（左侧）：

#<tf.Tensor: shape=(10, 10), dtype=int32, numpy=
#array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
#       [1, 2, 3, 4, 5, 6, 7, 8, 9, 0],
#       [2, 3, 4, 5, 6, 7, 8, 9, 0, 1],
#       [3, 4, 5, 6, 7, 8, 9, 0, 1, 2],
#       [4, 5, 6, 7, 8, 9, 0, 1, 2, 3],
#       [5, 6, 7, 8, 9, 0, 1, 2, 3, 4],
#       [6, 7, 8, 9, 0, 1, 2, 3, 4, 5],
#       [7, 8, 9, 0, 1, 2, 3, 4, 5, 6],
#       [8, 9, 0, 1, 2, 3, 4, 5, 6, 7],
#       [9, 0, 1, 2, 3, 4, 5, 6, 7, 8]], dtype=int32)>

最后，使用Tensor切片获取所需的元素，如下所示：

window = 8
 tf.slice(w, [0, 0], [(w.shape[0] - window) + 1, window])

给出：

#<tf.Tensor: shape=(3, 8), dtype=int32, numpy=
#array([[0, 1, 2, 3, 4, 5, 6, 7],
#       [1, 2, 3, 4, 5, 6, 7, 8],
#       [2, 3, 4, 5, 6, 7, 8, 9]], dtype=int32)>

对于窗口= 4

window = 4
tf.slice(w, [0, 0], [(w.shape[0] - window) + 1, window])

给出：

#array([[0, 1, 2, 3],
#   [1, 2, 3, 4],
#   [2, 3, 4, 5],
#   [3, 4, 5, 6],
#   [4, 5, 6, 7],
#   [5, 6, 7, 8],
#   [6, 7, 8, 9]], dtype=int32)>

尝试将其转换为一个tf图，看看它是否比普通的python循环有更好的性能。