本文介绍一种无需显式循环即可从 pytorch 二维张量每行中按指定起始索引和固定长度提取子张量的方法，利用 `torch.arange` 与 `torch.gather` 实现全向量化索引。

在深度学习与科学计算中，常需对批量数据（如 N×D 的特征矩阵）按行进行变起点、定长度的切片操作。例如：给定一个形状为 (N, D) 的张量 data，以及长度为 N 的起始索引张量 start_idx，要求对第 i 行提取 data[i, start_idx[i]:start_idx[i] + L]，其中 L 为统一子序列长度。若使用 Python 循环或列表推导，不仅低效，还破坏了张量计算的并行性。

PyTorch 提供了高效的向量化方案：构造索引张量 + gather 沿指定维度收集。核心思路是：

1. 对每个起始索引 start_idx[i]，生成对应行的连续索引范围 start_idx[i], start_idx[i]+1, ..., start_idx[i]+L−1；
2. 将这些范围堆叠成形状为 (N, L) 的二维索引张量；
3. 调用 data.gather(dim=1, index=index_tensor)，沿列维度（dim=1）按行采集指定列索引的值。

注意：start_idx 必须为整数类型（如 torch.long），浮点型索引不被支持；且所有子序列长度必须一致（L 固定），否则无法构成规则索引张量。

以下是完整可运行示例：

import torch

def gather_rows_by_range(data: torch.Tensor, start_idx: torch.Tensor, length: int, dim: int = 1) -> torch.Tensor:
    """
    从 data 的每行（若 dim=1）或每列（若 dim=0）中提取长度为 length 的连续子序列，
    起始位置由 start_idx 指定（按行/列对齐）。

    Args:
        data: 输入张量，形状 (N, D)
        start_idx: 起始索引，形状 (N,)，dtype=torch.long
        length: 子序列固定长度（标量）
        dim: 沿哪一维采样（默认 1，即按行取列）

    Returns:
        输出张量，形状 (N, length)
    """
    # 为每行生成 [s, s+1, ..., s+length-1]
    ranges = torch.stack([
        torch.arange(s, s + length, device=data.device, dtype=torch.long)
        for s in start_idx
    ])
    return data.gather(dim, ranges)

# 示例数据
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 = gather_rows_by_range(data, start_idx, length=2, dim=1)

print(result)
# 输出：
# tensor([[ 1.,  2.],
#         [ 9., 10.],
#         [12., 13.]])

✅ 优势总结：

• 完全向量化，GPU 友好，避免 Python 循环开销；
• 支持自动求导（requires_grad=True 时梯度可正确回传）；
• 易扩展至更高维（如 batched 3D 张量，只需调整 dim 和索引构造逻辑）。

⚠️ 注意事项：

• 确保所有 start_idx[i] + length ≤ data.size(dim)，否则将触发 IndexError（PyTorch 不做边界检查）；
• 若需动态长度，需改用 torch.nested（v2.0+）或分组 padding + mask，无法直接用 gather；
• torch.stack 构造索引张量时，若 N 很大，可考虑用广播技巧（如 start_idx.unsqueeze(1) + torch.arange(length)）进一步优化内存与速度：

# 更高效的索引张量构造（推荐用于大数据量）
index_tensor = start_idx.unsqueeze(1) + torch.arange(length, device=data.device)
result = data.gather(1, index_tensor)

该方法是 PyTorch 中实现“行级动态切片”的标准实践，在 Transformer 的 sliding window attention、时序模型的 patching 等场景中广泛应用。