[torchlib] prims::sum (#2778)

Fix pytorch/pytorch#173074

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: justinchuby

onnxscript/function_libs/torch_lib/ops/prims.py

@@ -772,12 +772,18 @@ def prims_sub(self: TTensor, other: TTensor) -> TTensor:
     return op.Sub(self, other)
 
 
+@torch_op("prims::sum", trace_only=True)
 def prims_sum(
     inp: TensorType, dims: Optional[Sequence[int]], output_dtype: Optional[int] = None
 ) -> TensorType:
     """sum(Tensor inp, int[]? dims, *, ScalarType? output_dtype=None) -> Tensor"""
 
-    raise NotImplementedError()
+    result = op.ReduceSum(inp, dims, keepdims=False)
+
+    if output_dtype is not None and output_dtype != -1:
+        result = op.Cast(result, to=output_dtype)
+
+    return result
 
 
 def prims_svd(A: TensorType, full_matrices: bool) -> tuple[TensorType, TensorType, TensorType]:

tests/function_libs/torch_lib/e2e_ops_tests.py

@@ -902,6 +902,18 @@ def forward(self, x, index, update):
         )
         _testing.assert_onnx_program(onnx_program)
 
+    def test_std_mean(self):
+        """Test torch.std_mean which will be decomposed into prims.sum."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.std_mean(x)
+
+        onnx_program = torch.onnx.export(
+            Model(), (torch.rand(10, 10, 10),), dynamo=True, verbose=False
+        )
+        _testing.assert_onnx_program(onnx_program)
+
 
 if __name__ == "__main__":
     unittest.main()