Refactor

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

Author: justinchuby

onnxscript/function_libs/torch_lib/ops/core.py

@@ -7657,11 +7657,8 @@ def aten_refine_names(self: TensorType, names: Sequence[str]) -> TensorType:
     raise NotImplementedError()
 
 
-@torch_op("aten::remainder.Tensor", trace_only=True)
-def aten_remainder(self: TTensor, other: TTensor) -> TTensor:
-    """remainder.Tensor(Tensor self, Tensor other) -> Tensor"""
-
-    if self.dtype.is_integer():
+def _aten_remainder(self: TTensor, other: TTensor, integer: bool) -> TTensor:
+    if integer:
         return op.Mod(self, other)
 
     # TODO(justinchuby): Improve fp16 precision by following the logic in
@@ -7673,34 +7670,27 @@ def aten_remainder(self: TTensor, other: TTensor) -> TTensor:
     return op.Sub(self, op.Mul(rounded_quotient, other))
 
 
+@torch_op("aten::remainder.Tensor", trace_only=True)
+def aten_remainder(self: TTensor, other: TTensor) -> TTensor:
+    """remainder.Tensor(Tensor self, Tensor other) -> Tensor"""
+
+    return _aten_remainder(self, other, integer=self.dtype.is_integer())
+
+
 @torch_op("aten::remainder.Scalar", trace_only=True)
 def aten_remainder_scalar(self: TTensor, other: float) -> TTensor:
     """remainder.Scalar(Tensor self, Scalar other) -> Tensor"""
 
     other_tensor = ir.tensor(other, dtype=self.dtype)
-
-    if self.dtype.is_integer():
-        return op.Mod(self, other_tensor)
-
-    # a - a.div(b, rounding_mode="floor") * b
-    rounded_quotient = op.Floor(op.Div(self, other_tensor))
-
-    return op.Sub(self, op.Mul(rounded_quotient, other_tensor))
+    return _aten_remainder(self, other_tensor, integer=self.dtype.is_integer())
 
 
 @torch_op("aten::remainder.Scalar_Tensor", trace_only=True)
 def aten_remainder_scalar_tensor(self: float, other: TTensor) -> TTensor:
     """remainder.Scalar_Tensor(Scalar self, Tensor other) -> Tensor"""
 
     self_tensor = ir.tensor(self, dtype=other.dtype)
-
-    if other.dtype.is_integer():
-        return op.Mod(self_tensor, other)
-
-    # a - a.div(b, rounding_mode="floor") * b
-    rounded_quotient = op.Floor(op.Div(self_tensor, other))
-
-    return op.Sub(self_tensor, op.Mul(rounded_quotient, other))
+    return _aten_remainder(self_tensor, other, integer=other.dtype.is_integer())
 
 
 @torch_op("_operator::mod", trace_only=True)