本文介绍一种无需显式循环即可从 pytorch 二维张量各行中按指定起始索引和统一长度批量切片的方法，核心是利用 `torch.arange` 构造索引张量，并通过 `gather` 实现向量化索引选取。
在深度学习与科学计算中，常需对批量数据（如序列、特征图）进行行级动态切片——即每行按各自起始位置截取固定长度的子序列。若使用 Python 循环或列表推导式逐行处理，不仅代码冗长，还会严重损害 GPU 加速优势。幸运的是，PyTorch 提供了完全向量化的解决方案：torch.gather + 动态索引张量构造。
其关键思想是：将每行的切片需求（起始索引 + 长度）转化为一个与输出形状一致的列索引张量，再通过 gather(dim=1, index=...) 沿列维度（dim=1）收集对应元素。
以下是完整实现：
import torch
def batch_row_slice(data: torch.Tensor, start_indices: torch.Tensor, length: int, dim: int = 1) -> torch.Tensor:
"""
从2D张量每行中提取固定长度子序列（向量化实现）
Args:
data: 输入2D张量，shape = (N, D)
start_indices: 每行起始列索引，shape = (N,)，dtype=torch.long
length: 每个切片的固定长度（所有行相同）
dim: 沿哪个维度切片（默认为列维度，即 dim=1）
Returns:
切片结果，shape = (N, length)
"""
if dim != 1:
raise NotImplementedError("仅支持沿列维度（dim=1）切片")
# 构造每行的索引序列：[[s0, s0+1], [s1, s1+1], ..., [s_{N-1}, s_{N-1}+1]]
# 使用 broadcasting 和 arange 实现高效生成
arange = torch.arange(length, device=data.device, dtype=start_indices.dtype)
# start_indices: (N,) → (N, 1); arange: (length,) → (1, length)
# 广播后得到 (N, length) 的起始偏移矩阵
index_tensor = start_indices.unsqueeze(1) + arange.unsqueeze(0)
# 使用 gather 沿 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) # 注意：必须为 long 类型！
result = batch_row_slice(data, start_idx, length=2)
print(result)
# 输出：
# tensor([[ 1., 2.],
# [ 9., 10.],
# [12., 13.]])✅ 关键要点说明：
- start_indices 必须是 torch.long 类型（浮点型索引不被 gather 接受）；
- length 必须对所有行保持一致，这是向量化前提；
- 索引越界行为由 gather 自动处理（超出范围时返回 0 填充），生产环境建议提前校验 start_idx + length
- 该方法天然支持 GPU 张量，全程无 CPU-GPU 数据拷贝，性能远超循环；
- 可轻松扩展至更高维（如 (B, T, D) 中按 B 维切时间步），只需调整 dim 和索引广播逻辑。
通过此方法，你可以在保持代码简洁性的同时，充分发挥 PyTorch 的张量并行能力，适用于 Transformer attention mask 构建、时序窗口采样、动态 padding 等典型场景。
