Update core.py

microsoft · Jul 8, 2024 · c21f6ac · c21f6ac
1 parent 4455dd3
commit c21f6ac
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -6657,10 +6657,18 @@ def aten_prelu_backward(
     raise NotImplementedError()
 
 
-def aten_prod(self: TensorType, dtype: Optional[int] = None) -> TensorType:
+@torch_op(("aten::prod"), trace_only=True)
+def aten_prod(self: TReal, dtype: Optional[int] = None) -> TReal:
     """prod(Tensor self, *, ScalarType? dtype=None) -> Tensor"""
 
-    raise NotImplementedError()
+    return op.ReduceProd(self)
+
+
+@torch_op("aten::prod.dim_int", trace_only=True)
+def aten_prod_dim(self: TReal, dim: int, keepdim: bool = False, dtype: Optional[int] = None) -> TReal:
+    """prod(Tensor self, *, ScalarType? dtype=None) -> Tensor"""
+
+    return op.ReduceProd(self, axes=dim, keepdims=keepdim)
 
 
 def aten_promote_types(type1: int, type2: int) -> int: