刚好遇到过类似的动态张量抽样需求，结合PyTorch的动态图特性，给你两种常用场景的实现方案，完全适配N随时变化的情况：

1. 无放回抽样（不重复选取k个元素） #

对应np.random.choice(..., replace=False)的行为，每个元素最多被选中一次，且所有元素被选中的概率均等。

实现思路很简单：给每个元素生成一个均匀分布的随机值，取随机值最大的前k个元素的索引，再用这些索引从原张量中提取元素——因为随机值是均匀分布的，topk的结果就等价于随机无放回抽样。

import torch

def random_choice_no_replace(tensor, k):
    # 输入tensor形状为(N, 2)，N支持动态变化
    current_n = tensor.size(0)
    if k > current_n:
        raise ValueError(f"无放回抽样时k={k}不能大于当前张量的行数N={current_n}")
    
    # 生成和张量行数一致的均匀随机数，获取topk的索引
    rand_values = torch.rand(current_n, device=tensor.device)
    _, selected_indices = rand_values.topk(k)
    
    # 根据索引抽取目标元素
    return tensor[selected_indices]

2. 有放回抽样（允许重复选取k个元素） #

对应np.random.choice(..., replace=True)的行为，元素可以被多次选中，每次抽样时每个元素被选中的概率都是1/N。

直接生成k个范围在[0, N-1]的均匀随机整数索引，再用索引提取元素即可，PyTorch的randint天然支持动态的范围参数，完美适配N变化的场景。

import torch

def random_choice_with_replace(tensor, k):
    # 输入tensor形状为(N, 2)，N支持动态变化
    current_n = tensor.size(0)
    
    # 生成k个符合要求的随机索引
    selected_indices = torch.randint(low=0, high=current_n, size=(k,), device=tensor.device)
    
    # 根据索引抽取目标元素
    return tensor[selected_indices]

额外注意点 #

• 设备一致性：代码里特意指定了device=tensor.device，如果你的张量在GPU上运行，随机索引也会生成在同一块GPU上，避免跨设备数据迁移的错误。
• 动态适配：所有操作都依赖tensor.size(0)获取当前的N值，完全没有硬编码的固定尺寸，不管N怎么变都能正常工作。
• 概率等价性：两种实现的概率分布和np.random.choice完全一致，可以放心替换使用。

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