jananisriram faddb0f30c [NestedTensor] Integrate the mean operator along the jagged dimension into NestedTensor (#131132) ... Summary: Modify the existing `mean` operator in PyTorch, invoked by `torch.mean`, to allow for reductions along the jagged dimension of a nested tensor. The function originally had a basic implementation for reducing along 1 non-ragged dimension. This diff enables PyTorch users to invoke `torch.mean` on a nested tensor when reducing along the ragged dimension, e.g. `*` in a `(B, *, M)` nested tensor. Parametrize unit tests from `sum` to verify the accuracy of the ragged reduction implementation for `torch.mean`. Add unit tests and parametrize `sum` unit tests to verify error handling for unsupported features in `NestedTensor` `torch.mean`. Test Plan: Verify that the new unit test passes via the following command: ``` buck2 run mode/{opt,inplace} //caffe2/test:nested -- --regex test_mean ``` ``` buck2 run mode/{opt,inplace} //caffe2/test:nested -- --regex test_jagged_op ``` Differential Revision: D59654668 Pull Request resolved: https://github.com/pytorch/pytorch/pull/131132 Approved by: https://github.com/davidberard98, https://github.com/jbschlosser		2024-07-23 18:48:34 +00:00
..
__init__.py
nested_tensor.py
ops.py	[NestedTensor] Integrate the mean operator along the jagged dimension into NestedTensor (#131132)	2024-07-23 18:48:34 +00:00
sdpa.py	[NJT+SDPA]Fix flash_attention output when batch_size=1 and seq_len=1 (#130652)	2024-07-15 19:44:04 +00:00