Allow slicing of Nested Tensors along constant dimensions #108567
Description
🚀 The feature, motivation and pitch
In many cases several of the nested tensor dimensions are constant, and therefore in principle should allow slicing
import torch
a = torch.rand(5,5,5)
b = torch.rand(5,5,10)
c = torch.rand(5,5,15)
d= torch.nested.nested_tensor([a,b,c])
d[:2]
Will currently return the following
NotImplementedError: Could not run 'aten::slice.Tensor' with arguments from the 'NestedTensorCPU' backend.
May be useful to be able to mark which dims are and will remain constant?
Alternatives
No response
Additional context
No response