本文介绍一种无需显式循环、利用 pytorch 的 `gather` 和广播机制，批量从二维张量各行中按动态起始索引提取固定长度子序列的方法。
在深度学习和张量处理中，常需对二维张量（如 batch × seq_len）的每一行独立切片——例如按每条样本的动态起始位置截取固定长度的子序列。若使用 Python 循环或列表推导式，不仅低效，还破坏了张量计算的并行性与可微性。幸运的是，PyTorch 提供了完全向量化（vectorized）的解决方案：结合 torch.arange 构造索引矩阵 + torch.gather 沿指定维度收集元素。
核心思想是：为每行生成一个长度为 chunksize 的连续索引序列（即 start[i], start[i]+1, ..., start[i]+chunksize-1），堆叠成二维索引张量；再用 gather(dim=1) 沿列方向依据该索引矩阵从原数据中“拾取”对应列值。
以下为完整实现：
import torch
def batch_slice_rows(data: torch.Tensor, start_indices: torch.Tensor, chunksize: int, dim: int = 1) -> torch.Tensor:
"""
向量化地从 data 的每一行（dim=1）中提取长度为 chunksize 的子序列。
Args:
data: 输入二维张量，形状为 (N, M)
start_indices: 起始索引张量，形状为 (N,)，dtype=torch.long
chunksize: 每个子序列的固定长度（必须为正整数）
dim: 切片维度（默认为 1，即按行切列）
Returns:
输出张量，形状为 (N, chunksize)
"""
if dim != 1:
raise NotImplementedError("仅支持 dim=1（按行切片）的向量化实现")
# 验证输入类型与范围
assert start_indices.dtype in (torch.long, torch.int64), "start_indices 必须为 long 类型"
assert (start_indices >= 0).all() and (start_indices + chunksize <= data.size(1)).all(), \
f"索引越界：start_indices={start_indices}, data.shape={data.shape}, chunksize={chunksize}"
# 为每行构造 [start[i], start[i]+1, ..., start[i]+chunksize-1]
# 使用广播：(N, 1) + (1, chunksize) → (N, chunksize)
offsets = torch.arange(chunksize, device=data.device, dtype=start_indices.dtype)
index_tensor = start_indices.unsqueeze(1) + offsets.unsqueeze(0) # shape: (N, chunksize)
# 沿 dim=1 收集数据（等价于按行索引列）
return torch.gather(data, dim=1, index=index_tensor)
# 示例使用
data = torch.tensor([[ 1., 2., 3., 4., 5.],
[ 6., 7., 8., 9., 10.],
[11., 12., 13., 14., 15.]])
start_idx = torch.tensor([0, 3, 1], dtype=torch.long)
result = batch_slice_rows(data, start_idx, chunksize=2)
print(result)
# 输出：
# tensor([[ 1., 2.],
# [ 9., 10.],
# [12., 13.]])✅ 关键优势：
- 完全免循环，GPU 友好，支持梯度回传（gather 是可微操作）；
- 时间复杂度 O(N×chunksize)，远优于循环版的 O(N×chunksize×Python开销)；
- 索引张量通过广播生成，内存高效，避免 stack + 列表推导的 Python 层开销（原答案中 torch.stack([...]) 在大数据量下易成瓶颈）。
⚠️ 注意事项：
- start_indices 必须为 torch.long 类型（浮点型索引不被 gather 接受）；
- 所有 start[i] + chunksize 必须 ≤ data.size(1)，否则触发越界错误；
- 当 chunksize 不固定时，需改用 torch.nested（PyTorch 2.0+）或 padding + mask 方案，无法纯向量化。
总结：该方法以简洁、高效、可扩展的方式解决了“按行动态切片”的典型需求，是构建高性能数据预处理流水线与自定义层的重要技巧。
