Skip to content

[Indexing] Incoherent Tensor indexing for nested lists #100080

Open
@remidomingues

Description

🐛 Describe the bug

The behavior of pytorch.tensor.__getitem__() is incoherent when using nested lists:

import torch
t = torch.tensor([1, 2])

t[0].shape  # torch.Size([])
t[[0]].shape  # torch.Size([1])
t[[[0]]].shape  # torch.Size([1]) <--- No dimension added here
t[[[[0]]]].shape  # torch.Size([1, 1]) <--- Subsequent shapes are thus off by 1 dimension

That behavior is not observed with tensordict or numpy, where dimensions are added as expected.

import numpy as np
nd = np.array([1, 2])

nd[0].shape  # ()
nd[[0]].shape  # (1,)
nd[[[0]]].shape  # (1, 1)
nd[[[[0]]]].shape  # (1, 1, 1)

Versions

PyTorch version: 2.0.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.3.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: Could not collect
Libc version: N/A

Python version: 3.10.8 (main, Nov 24 2022, 08:08:27) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.3.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Versions of relevant libraries:
[pip3] numpy==1.24.2
[pip3] torch==2.0.0
[pip3] torchrl==0.1.0+9a85f45
[conda] numpy 1.24.2 pypi_0 pypi
[conda] torch 2.0.0 pypi_0 pypi
[conda] torchrl 0.1.0+9a85f45 dev_0

cc @mruberry @rgommers @ezyang @gchanan @zou3519

Metadata

Assignees

No one assigned

    Labels

    module: advanced indexingRelated to x[i] = y, index functionsmodule: edge casesAdversarial inputs unlikely to occur in practicemodule: numpyRelated to numpy support, and also numpy compatibility of our operatorstriagedThis issue has been looked at a team member, and triaged and prioritized into an appropriate module

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions