如何在TensorFlow 2.0中实现类似PyTorch的切片赋值？

TensorFlow的EagerTensor是不可变对象，不支持PyTorch那样的直接切片赋值操作，以下是几种适用于现代TensorFlow（2.x及以上）的可行方案：

方案1：针对tf.Variable的切片赋值（最接近PyTorch写法） #

如果你的input_seq是tf.Variable类型，可以直接使用切片的assign方法实现原地更新：

import tensorflow as tf
import numpy as np

# 初始化示例Variable
input_seq = tf.Variable(tf.zeros((5, 3)))
# 示例numpy数组
stroke = np.random.randn(4, 3)
# 转换为Tensor
stroke_tensor = tf.convert_to_tensor(stroke[:-1, :])

# 执行切片赋值
input_seq[1:, :].assign(stroke_tensor)

方案2：通过拼接生成新张量（适合连续规则切片） #

如果input_seq是普通EagerTensor（不可变），可以将不需要修改的部分与新值拼接，生成新的张量：

input_seq = tf.zeros((5, 3))
stroke = np.random.randn(4, 3)
stroke_tensor = tf.convert_to_tensor(stroke[:-1, :])

# 拼接保留的前1行和新值
new_input_seq = tf.concat([
    input_seq[:1, :],
    stroke_tensor
], axis=0)

这种方法代码简洁，性能高效，适合处理开头/结尾的连续切片场景。

方案3：使用tf.tensor_scatter_nd_update（通用任意切片） #

对于更复杂的非连续切片，推荐使用tf.tensor_scatter_nd_update，它可以精准指定要更新的位置：

input_seq = tf.zeros((5, 3))
stroke = np.random.randn(4, 3)
stroke_tensor = tf.convert_to_tensor(stroke[:-1, :])

# 构造要更新的位置索引：行范围1到末尾，所有列
rows = tf.range(1, input_seq.shape[0])
cols = tf.range(input_seq.shape[1])
indices = tf.stack(tf.meshgrid(rows, cols, indexing='ij'), axis=-1)

# 生成更新后的新张量
new_input_seq = tf.tensor_scatter_nd_update(input_seq, indices, stroke_tensor)

该方法适用于所有切片场景，无论是连续还是非连续的位置更新。

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