NJT <-> padded dense conversions

[jbschlosser]

Stack from ghstack (oldest at bottom):

• (WIP) to_padded_tensor() triton kernel for NJT #121947
• Use return_and_correct_aliasing() for NJT + compatible storage setting #126552
• -> NJT <-> padded dense conversions #125947

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/125947

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 5 New Failures, 12 Unrelated Failures

As of commit 6033447 with merge base 4e03263 ():

NEW FAILURES - The following jobs have failed:

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
  'Test'
• inductor / cuda12.1-py3.12-gcc9-sm86 / test (inductor, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
  'Test'
• Lint / lintrunner-noclang / linux-job (gh)
  >>> Lint for test/test_nestedtensor.py:
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  test_nestedtensor.py::TestNestedTensorSubclassCUDA::test_to_padded_tensor_compile_nt_dim_2_requires_grad_False_cuda_float16
• pull / linux-jammy-py3.8-gcc11 / test (distributed, 1, 2, linux.2xlarge) (gh)
  Process completed with exit code 1.

UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (aot_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128903)
  ImportError: attempted relative import with no known parent package
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (aot_inductor_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128903)
  ImportError: attempted relative import with no known parent package
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128902)
  ImportError: attempted relative import with no known parent package
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128902)
  ImportError: attempted relative import with no known parent package
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128901)
  ImportError: attempted relative import with no known parent package
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128901)
  ImportError: attempted relative import with no known parent package
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128929)
  ImportError: attempted relative import with no known parent package
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (aot_eager_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128929)
  ImportError: attempted relative import with no known parent package
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamic_aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128931)
  ImportError: attempted relative import with no known parent package
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamic_aot_eager_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128931)
  ImportError: attempted relative import with no known parent package
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamo_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128932)
  ImportError: attempted relative import with no known parent package
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamo_eager_torchbench, 2, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128932)
  ImportError: attempted relative import with no known parent package

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vadimkantorov]

Related old discussion on this being a useful primitive (in the old days for collate_fn of data loading) and deserving more fame :)

• [feature request] Jagged / padding version of torch.stack / torch.cat + some general nested tensor discussion #65156

One useful thing here is also to support "padding multiples" per dimension

[vadimkantorov]

Maybe one way to auto-construct NJT from torch.stack([...]) call in default collate could be:

1. from dataset.__getitem__ return dense tensors but wrapped in NJT (but the internal representation should be just regular dense tensor)
2. support that torch.stack([..., ]) returns a NJT if elements in the input list are NJT (even if inside they are just dense tensors)

like this the collate_fn code could be kept unchanged, but if the inputs are wrapped as NJT, it would start to produce a NJT...