Joel Schlosser 3abbde976d Allow any single non-batch dim to be ragged for NJT (#137125) ...

Fixes #137512

Relaxes the restriction that the ragged dim is immediately next to the batch dim e.g. `(B, *, D_0, ..., D_N)`. This allows for constructing NJTs of shape e.g. `(B, D, j0)` directly. It's possible before this PR to get an NJT of e.g. shape `(B, D, j0)` by constructing an NJT of shape `(B, j0, D)` and transposing it. This PR allows a user to go straight there without the transpose. The standard `torch.nested.nested_tensor(list)` constructor has been updated to support this.

At the very least, this is useful for testing on transposed NJTs. I'm willing to make this functionality private if needed.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/137125
Approved by: https://github.com/cpuhrsch, https://github.com/soulitzer

2024-11-06 18:50:08 +00:00

..

_internal

Allow any single non-batch dim to be ragged for NJT (#137125)

2024-11-06 18:50:08 +00:00

__init__.py

Implement 2D version of masked_select for nestedtensors (#133889)

2024-08-20 21:46:32 +00:00