Repository: pytorch
Issue: [Tracker] Move nested tensors to beta
Labels: triaged,module: nestedtensor

This issue tracks the status of readying nested tensors to be promoted from prototype status -> beta, targeting the 2.5 release. See also: [op request tracker issue](https://github.com/pytorch/pytorch/issues/118107) for new op coverage requests.

* Public APIs / operator support
    * -[X] `torch.nested.nested_tensor(tensor_list, layout=torch.jagged)` - construct a jagged layout NT without preserving autograd history (#111078)
    * -[X] `torch.nested.as_nested_tensor(tensor_list, layout=torch.jagged)` - construct a jagged layout NT preserving autograd history (#112304)
    * -[x] `torch.nested.as_nested_tensor(t, layout=torch.jagged)` - construct a jagged layout NT with consistent sequence lengths from a dense tensor (in progress by @jbschlosser in #113280)
    * -[ ] Conversion between strided and jagged layouts using `to(layout)` (in progress by @ani300 in #115749)
    * -[x] Non-copying constructor- construct a jagged layout NT from the `(values, offsets)` components directly (@jbschlosser in #121518)
    * -[ ] #108567: Slicing of arbitrary dims for NTs of both layouts
    * -[ ] Resolve rough edges of SDPA support for jagged layout NTs ([internal only: link](https://docs.google.com/document/d/1_86Yq3DY4NK66eQXuJ3pxAjZEds6nnVipYpymdmEjrk))
        * -[x] Implement backend selection in `NestedTensor.__torch_function__()` impl of SDPA (#114164)
        * -[x] Support transpose of jagged layout NTs on the ragged dim (need to track this appropriately via `ragged_idx` and ensure the proper invariants hold for any ops we support) (#114164)
        * -[x] `torch.nested.narrow()` - inverse of `to_padded_tensor()` i.e. return a NT view of a dense buffer for e.g. kv-caching ([internal only: design doc link](https://docs.google.com/document/d/1Yj5Nv25DfSt0Q2AJ_uYnwV8nnZ-Xth6anIRDrRq5gJM)) (#112770)
        * -[x] Working backward support (#115576)
        * -[x] SDPA when ragged on a dim other than sequence length (e.g. for SAM) (#115636)
        * -[ ] Robust validation for all use cases
        * -[x] Fix Proxy error coming from PT2 + `__torch_function__()` subclass issue #120654 (in progress by @soulitzer)
    * -[x] Support FlexAttention (in progress by @jbschlosser in #136792)
    * -[x] Autograd support for broadcasting https://github.com/pytorch/pytorch/pull/112519
    * -[ ] AccumulateGrad support when existing .grad has not been cleared
    * -[x] Ops that are currently `CompositeImplicitAutogradNestedTensor` must be implemented for jagged via `NestedTensor.__torch_function__()` (#115191, #115192)
        * This includes `reshape()` / `reshape_as()`
    * -[ ] Backward formulas specific to C++ NTs must be updated to work for jagged layout NTs as well
        * This includes functions that involve NT <-> dense views: `split_with_size()`, `unbind()`, `values()` (in progress by @soulitzer)
    * -[ ] (Stretch) Revisit vmap + NT integration? ([internal only: link to issues encountered for SAM](https://docs.google.com/document/d/1_yiHOBbaI4qFWqBfebjWPHOhkxKW3v-CHu3lj5apv1Y))
    * [x] support not batch_first for Transformers with NT https://github.com/pytorch/pytorch/pull/106668
    * [x] make FSDP unit tests (and others that expect bit-accurate repeatability between training and inference) pass by disabling Transformer inference with NT (and other fastpath options) using context manager   https://github.com/pytorch/pytorch/pull/107163
    * [x] Padded dense <-> jagged conversions (#125947)
    * [x] Fix integration with `inference_mode()` (in progress by @jbschlosser in #134683)
* General UX
    * -[ ] Unify printed form of NTs with strided / jagged layouts
    * -[x] Don't break PyTorch Profiler when `with_flops=True` (#112767)
    * -[x] [Serialization logic](https://github.com/pytorch/pytorch/blob/0241ed9331e1f56d1aca00bdc753af3e74435642/torch/nested/_internal/nested_tensor.py#L155-L166) is incomplete (in progress by @jbschlosser in #137031)
* PT2 integration
    * -[ ] Goal: End-to-end support for example models: SAM, HSTU (stretch)
    * -[x] Ensure clean graph breaking / fallbacks for C++ NTs + torch.compile https://github.com/pytorch/pytorch/pull/112662
    * -[x] Support broadcasting over ragged dimension (e.g. `(B, *, D)` + `(B, 1, D)`) without specializing on batch size (in progress by @jbschlosser; #133021)
    * -[x] Fix fake-ification of NT view inputs
        * -[x] Patched in #112205
        * -[x] Doing this properly requires us to return actual views. This is also needed for proper autograd integration, in-place op support, kv-caching, etc. (in progress by @jbschlosser in #113279)
    * -[x] Fix automatic dynamic / `mark_dynamic()` support on jagged layout NTs (#114311)
    * -[x] Fix support for constructing jagged NTs in-line within the compiled function rather than being passed as args (graph break patches this but has bad perf) (#133196)
    * -[ ] Support usage of nt.shape within compiled region; currently (when *not following a graph break*) we try to wrap the singleton symint in ConstantVariable and blow up (see [here](https://github.com/pytorch/pytorch/blob/d78fe039eb628873ad81a1ff1d601090e9834fd9/torch/_dynamo/variables/tensor.py#L231-L233))
    * -[ ] Fix guard eval when running a function initially compiled with a jagged NT input with a strided NT input
    * -[x] Dynamic shapes support for min / max sequence length (in progress by @jbschlosser in #122836)
    * -[ ] Data-dependent output shape support without graph breaks (in progress by @jbschlosser in e.g. #142063)
    * -[ ] Factory function support (#119977)
    * -[ ] Composition of NJTs with different metadata (e.g. different offsets tensors) but same conceptual ragged structure
* Documentation 
    * -[ ] Write a tutorial or note describing the strided and jagged layouts, when to use each, and torch.compile integration with the jagged layout
    * -[ ] Ensure existing NT docs are up-to-date (at the very least, indicate prototype -> beta there :) )
        * -[x] OG NT tutorial implementing MHA should use SDPA and `as_nested_tensor()` instead of `nested_tensor()` to ensure gradients flow properly (https://github.com/pytorch/tutorials/pull/2813)
        * -[ ] Update SDPA docs to use NJT instead of NST
    * -[ ] Write a tutorial with an end-to-end example of utilizing a KV-cache with NTs
    * -[x] Write Transformer tutorials utilizing NJT + torch.compile (https://github.com/pytorch/tutorials/pull/3075)
* Performance
    * -[ ] Minimize memory usage by rewriting autograd formulas to save less for backwards. In practice, this has dependencies on a nested int -> {offsets, min / max seq len} registry / factory function support for nested ints (in progress by @soulitzer)
* Testing
    * -[x] More extensive coverage for ops supported by the jagged layout NT subclass, possibly utilizing `OpInfo`s (@jbschlosser in #131704)
    * -[x] #110288: Root cause and fix flaky NT tests with `PYTORCH_TEST_WITH_DYNAMO=1` https://github.com/pytorch/pytorch/issues/110288#issuecomment-1789364941 (should be fixed by #112662)
    * -[x] Forward / backward tests in eager and torch.compile across various permutations of NJT types (#138370)
* Type checking
    * -[ ] Pyre integration

cc @cpuhrsch @bhosmer @drisspg @soulitzer