Conversions between strided and jagged layouts for Nested Tensors

[ani300]

This PR does 3 things:

1. Adds a copy-free strided->jagged layout conversion for NT
2. Adds a copy-free jagged->strided layout conversion for NT
3. Modifies and expands the .to() API to support the layout argument for the specific case of NT layout conversion.

Stack from ghstack (oldest at bottom):

• -> Conversions between strided and jagged layouts for Nested Tensors #115749

cc @cpuhrsch @jbschlosser @bhosmer @drisspg @soulitzer

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

❌ 25 New Failures

As of commit dcfc6c0 with merge base 48a5414 ():

NEW FAILURES - The following jobs have failed:

• pull / linux-focal-cuda11.8-py3.10-gcc9 / test (distributed, 1, 3, linux.8xlarge.nvidia.gpu) (gh)
  distributed/tensor/parallel/test_tp_examples.py::DistTensorParallelExampleTest::test_loss_parallel
• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 1, 5, linux.4xlarge.nvidia.gpu) (gh)
  test_nn.py::TestNNDeviceTypeCUDA::test_module_to_empty_cuda_float64
• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 2, 5, linux.4xlarge.nvidia.gpu) (gh)
  test_jit.py::TestTracer::test_type_same_device
• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (deploy, 1, 1, linux.4xlarge.nvidia.gpu) (gh)
  [ FAILED ] TorchpyTest.PrintInstruction
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 1, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  test_nn.py::TestNNDeviceTypeCUDA::test_module_to_empty_cuda_float64
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 4, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  dynamo/test_debug_utils.py::TestDebugUtils::test_cast_model_to_fp64_dtype_args
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 5, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  test_jit.py::TestTracer::test_type_same_device
• pull / linux-focal-py3.11-clang10 / test (crossref, 1, 2, linux.2xlarge) (gh)
  test_nn.py::TestNNDeviceTypeCPU::test_module_to_empty_cpu_float64
• pull / linux-focal-py3.11-clang10 / test (default, 1, 3, linux.2xlarge) (gh)
  test_nn.py::TestNNDeviceTypeCPU::test_module_to_empty_cpu_float64
• pull / linux-focal-py3.11-clang10 / test (dynamo, 1, 3, linux.2xlarge) (gh)
  test_nn.py::TestNN::test_module_to_argparse
• pull / linux-focal-py3.11-clang10 / test (dynamo, 2, 3, linux.2xlarge) (gh)
  test_sparse.py::TestSparseCPU::test_sparse_dense_mul_cpu_float16
• pull / linux-focal-py3.11-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  test_nestedtensor.py::TestNestedTensor::test_to_jagged
• pull / linux-focal-py3.12-clang10 / test (default, 1, 3, linux.2xlarge) (gh)
  test_nn.py::TestNNDeviceTypeCPU::test_module_to_empty_cpu_float64
• pull / linux-focal-py3.12-clang10 / test (dynamo, 1, 3, linux.2xlarge) (gh)
  test_nn.py::TestNN::test_module_to_argparse
• pull / linux-focal-py3.12-clang10 / test (dynamo, 2, 3, linux.2xlarge) (gh)
  test_nn.py::TestNN::test_type
• pull / linux-focal-py3.12-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  test_sparse.py::TestSparseCPU::test_sparse_dense_mul_cpu_float16
• pull / linux-focal-py3.8-clang10 / test (crossref, 1, 2, linux.2xlarge) (gh)
  test_nn.py::TestNNDeviceTypeCPU::test_module_to_empty_cpu_float64
• pull / linux-focal-py3.8-clang10 / test (default, 1, 3, linux.2xlarge) (gh)
  test_nn.py::TestNNDeviceTypeCPU::test_module_to_empty_cpu_float64
• pull / linux-focal-py3.8-clang10 / test (dynamo, 2, 3, linux.2xlarge) (gh)
  test_sparse.py::TestSparseCPU::test_sparse_dense_mul_cpu_float16
• pull / linux-focal-py3.8-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  test_nestedtensor.py::TestNestedTensor::test_to_jagged
• pull / linux-focal-py3.8-clang10-onnx / test (default, 1, 2, linux.2xlarge) (gh)
  onnx/test_fx_op_consistency.py::TestOnnxModelOutputConsistency_opset_version_18_model_type_TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAMCPU::test_output_match_nn_functional_huber_loss_cpu_float16
• pull / linux-focal-py3.8-clang10-onnx / test (default, 2, 2, linux.2xlarge) (gh)
  onnx/test_fx_op_consistency.py::TestOnnxModelOutputConsistency_opset_version_18_model_type_TorchModelType.TORCH_NN_MODULECPU::test_output_match_alias_copy_cpu_float16
• pull / linux-jammy-py3.10-clang15-asan / test (default, 1, 6, linux.4xlarge) (gh)
  test_nn.py::TestNNDeviceTypeCPU::test_module_to_empty_cpu_float64
• pull / linux-jammy-py3.8-gcc11 / test (default, 1, 3, linux.2xlarge) (gh)
  test_nn.py::TestNNDeviceTypeCPU::test_module_to_empty_cpu_float64
• pull / linux-jammy-py3.8-gcc11 / test (distributed, 2, 2, linux.2xlarge) (gh)
  distributed/tensor/parallel/test_tp_examples.py::DistTensorParallelExampleTest::test_loss_parallel

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

[soulitzer]

Add a test?

[jbschlosser]

I realize this doesn't work today, but I think we want this sugar to work as the public API for conversion:

nt.to(layout=torch.jagged)

Today this gives:

TypeError: to() received an invalid combination of arguments - got (layout=torch.layout, ), but expected one of:
 * (torch.device device, torch.dtype dtype, bool non_blocking, bool copy, *, torch.memory_format memory_format)
 * (torch.dtype dtype, bool non_blocking, bool copy, *, torch.memory_format memory_format)
 * (Tensor tensor, bool non_blocking, bool copy, *, torch.memory_format memory_format)

I think if we support that in general then we can provide a high-level meta registration for it that is used by FakeTensor to solve the Dynamo issue you're hitting for the tests. A good place to start is in tools/autograd/templates/python_variable_methods.cpp where THPVariable_to() and dispatch_to() are defined.