You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

求助:Python 3.10中Tensor索引避开*运算符的简洁实现方法

在Python 3.10中实现PyTorch Tensor可变维度的简洁索引方法

Python的索引语法不支持在方括号内直接用*解包列表(比如t[*index]这种写法是非法的),但你可以通过以下简洁方法实现可变维度的Tensor索引:

核心方案:将索引列表转换为元组

PyTorch完全支持用元组作为Tensor的索引,而列表转元组的操作非常简便。只要你的index列表长度和Tensor的维度数匹配,直接用t[tuple(index)]就能达到你想要的效果。

示例代码:

import torch
t = torch.zeros(5, 6, 8)
index = [3, 2, 7]
# 列表转元组后执行索引
print(t[tuple(index)])

这个方法不受维度数量限制——不管你的Tensor是8维还是10维,只要index列表的元素个数和Tensor维度数一致,就能直接套用,完全替代硬编码的索引写法。

扩展场景:混合索引(可选)

如果需要在索引中混合使用具体位置和切片(比如部分维度用索引、其余维度取全部),可以通过元组拼接实现。例如,若Tensor是3维,index只包含前2维的索引,要取第3维全部元素:

index = [3, 2]
print(t[tuple(index) + (slice(None),)])  # 等价于t[3, 2, :]

内容的提问来源于stack exchange,提问作者Wentinn Liao

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.07.05 01:17:11