Support sum() forward and backward for NJT

[jbschlosser]

Stack from ghstack (oldest at bottom):

• -> Support sum() forward and backward for NJT #128031

[pytorch-bot]

🔗 Helpful Links

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

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

❌ 10 New Failures, 4 Unrelated Failures

As of commit 1b29546 with merge base 31d5753 ():

NEW FAILURES - The following jobs have failed:

• pull / linux-focal-py3.11-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh)
  'test/test_nestedtensor.py::TestNestedTensorSubclassCPU::test_unbind_backward_cpu_float16'
• pull / linux-focal-py3.11-clang10 / test (default, 2, 3, linux.2xlarge) (gh)
  'test/test_nestedtensor.py::TestNestedTensorSubclassCPU::test_unbind_backward_cpu_float16'
• pull / linux-focal-py3.11-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  test_nestedtensor.py::TestNestedTensorSubclassCPU::test_sum_requires_grad_True_cpu_float16
• pull / linux-focal-py3.12-clang10 / test (default, 2, 3, linux.2xlarge) (gh)
  'test/test_nestedtensor.py::TestNestedTensorSubclassCPU::test_unbind_backward_cpu_float16'
• pull / linux-focal-py3.12-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  test_nestedtensor.py::TestNestedTensorSubclassCPU::test_sum_requires_grad_True_cpu_float16
• pull / linux-focal-py3.8-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh)
  'test/test_nestedtensor.py::TestNestedTensorSubclassCPU::test_unbind_backward_cpu_float16'
• pull / linux-focal-py3.8-clang10 / test (default, 2, 3, linux.2xlarge) (gh)
  'test/test_nestedtensor.py::TestNestedTensorSubclassCPU::test_unbind_backward_cpu_float16'
• pull / linux-focal-py3.8-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  'test/test_nestedtensor.py::TestNestedTensorSubclassCPU::test_unbind_backward_cpu_float16'
• pull / linux-jammy-py3.10-clang15-asan / test (default, 2, 6, linux.4xlarge) (gh)
  'test/test_nestedtensor.py::TestNestedTensorSubclassCPU::test_unbind_backward_cpu_float16'
• pull / linux-jammy-py3.8-gcc11 / test (default, 2, 3, linux.2xlarge) (gh)
  'test/test_nestedtensor.py::TestNestedTensorSubclassCPU::test_unbind_backward_cpu_float16'

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / linux-jammy-py3.8-gcc11 / test (distributed, 1, 2, linux.2xlarge) (gh) (similar failure)
  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:

• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 2, 5, linux.4xlarge.nvidia.gpu) (gh) (related job)
  'test/test_nestedtensor.py::TestNestedTensorSubclassCUDA::test_unbind_backward_cuda_float16'
• pull / linux-focal-cuda12.1-py3.10-gcc9-experimental-split-build / test (default, 2, 5, linux.4xlarge.nvidia.gpu, unstable) (gh) (#129080)
  'test/test_nestedtensor.py::TestNestedTensorSubclassCUDA::test_unbind_backward_cuda_float16'
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu) (gh) (related job)
  'test/test_nestedtensor.py::TestNestedTensorSubclassCUDA::test_unbind_backward_cuda_float16'

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

[soulitzer]

I don't remember what the conclusion on this was when we discussed lol - but it seems like it is tricky to land something with suboptimal performance when there is already a simple way around this: extract out buffer and then sum.

[jbschlosser]

it seems like it is tricky to land something with suboptimal performance when there is already a simple way around this: extract out buffer and then sum.

yeah, I don't disagree, but it seems like people are very quick to run into "sum() isn't defined for NJT" when they first start using it. It's a noticeable gap and feels really silly not to have it.

Maybe we can do some hackery to make sum() composite implicit for NJT and get a good backwards automatically through _nested_get_values() + values.sum()?

[soulitzer]

Maybe we can do some hackery to make sum() composite implicit for NJT and get a good backwards automatically through _nested_get_values() + values.sum()?

Ooh good idea, this is just torch function right?