[Feature] Indexing specs #1105
[Feature] Indexing specs #1105
Conversation
…orSpec, UnboundedContinuousTensorSpec & UnboundedDiscreteTensorSpec
…orSpec, UnboundedContinuousTensorSpec & UnboundedDiscreteTensorSpec
facebook-github-bot
added
the
CLA Signed
| @@ -135,7 +135,12 @@ def _slice_indexing(shape: list[int], idx: slice): | |||
| return [n_items] + shape[1:] | |||
|
|
|||
|
|
|||
| def _shape_indexing(shape: list[int], idx: SHAPE_INDEX_TYPING): | |||
| def _shape_indexing( | |||
There was a problem hiding this comment.
this seems pretty crucial and we may refactor that at some point (e.g. using fake tensors)
Can we add a short docstring to say what it is about?
There was a problem hiding this comment.
we have a similar function in tensordict called _getitem_batch_size, is it inspired by that?
There was a problem hiding this comment.
Docstring added! _shape_indexing is definitely similar to _getitem_batch_size. I couldn't use the latter though as it doesn't perform some indexing checks, which are only executed when actually indexing the tensor. I assumed we wanted those checks, hence the reimplementation
from tensordict.utils import _getitem_batch_size
from torchrl.data.tensor_specs import _shape_indexing
_getitem_batch_size(torch.Size((3, 2)), 5) # torch.Size([2])
_shape_indexing(torch.Size((3, 2)), 5) # IndexError: index 5 is out of bounds for axis 0 with size 3
If we can work with fake_mode to have both fast shape indexing and such checks, that would definitely be the best of both worlds!
Description
Add indexing support to remaining specs:
Already supported by previous PR #1081 :
Note: although indexing & tests have been implemented for BoundedTensorSpec & MultiDiscreteTensorSpec, a NotImplementedError is currently set to prevent a different behavior with the other specs until pytorch/pytorch#100080 is addressed.
Motivation and Context
Address feature request of adding indexing to specs: #1051.
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!
cc @matteobettini