Joel Schlosser 41f315417c Fix NJT linear_backward() memory usage (#141163) ... Fixes #141112 The formula we're using for `linear_backward()` is inefficient for higher dim input sizes, even if the input is trivially higher dim (e.g. via use of `unsqueeze()`). This PR updates the formula to match the more efficient version employed by NST. Specifically, note the leading dim collapse for `grad_output`'s values before we compute the various matmuls. d5ee1d1b58/aten/src/ATen/native/nested/NestedTensorBackward.cpp (L37-L70) Testing for correctness is done via existing gradcheck tests (e.g. `test_backward_nn_functional_linear`). I added a memory usage test but I think it's likely there's a better way to do this. Pull Request resolved: https://github.com/pytorch/pytorch/pull/141163 Approved by: https://github.com/Skylion007, https://github.com/cpuhrsch, https://github.com/soulitzer		2024-11-21 15:22:45 +00:00
..
_internal	Fix NJT linear_backward() memory usage (#141163)	2024-11-21 15:22:45 +00:00
__init__.py	Allow NJT by default for weights_only torch.load (take 2) (#140739)	2024-11-19 02:44:53 +00:00