[torchlib] Fix aten_div rounding_mode (microsoft#2147)

Fix microsoft#2144

Author: justinchuby

onnxscript/function_libs/torch_lib/ops/core.py

@@ -2742,10 +2742,6 @@ def aten_dist(self: TensorType, other: TensorType, p: float = 2.0) -> TensorType
     (
         "aten::div.Tensor",
         "aten::div.Scalar",
-        # When rounding_mode is None, performs a true division
-        # https://pytorch.org/docs/stable/generated/torch.div.html
-        "aten::div.Tensor_mode",
-        "aten::div.Scalar_mode",
         "aten::divide.Tensor",
         "aten::divide.Scalar",
         "aten::true_divide.Tensor",
@@ -2799,41 +2795,45 @@ def aten_div_complex(self: TFloat, other: TFloat) -> TFloat:
 
 
 @torch_op(("aten::div.Tensor_mode", "aten::div.Scalar_mode"), trace_only=True)
-def aten_div_mode(self: TFloat, other: TFloat, rounding_mode: str) -> TFloat:
+def aten_div_mode(self: TFloat, other: TFloat, rounding_mode: Optional[str] = None) -> TFloat:
     """div.Tensor_mode(Tensor self, Tensor other, *, str? rounding_mode) -> Tensor"""
 
-    # TODO(justinchuby): trace_only=False when we use opset19 which supports string comparison
-    assert rounding_mode in {"trunc", "floor"}
+    assert rounding_mode in {"trunc", "floor", None}
 
     if rounding_mode == "trunc":
         # Rounds the results of the division towards zero.
         # Equivalent to C-style integer division
-        result = aten_trunc(op.Div(self, other))
-    else:  # rounding_mode == "floor"
-        result = op.Floor(op.Div(self, other))
+        return aten_trunc(op.Div(self, other))
+    if rounding_mode == "floor":
+        return op.Floor(op.Div(self, other))
 
-    return result
+    return op.Div(self, other)
 
 
 @torch_op(("aten::div.Tensor_mode", "aten::div.Scalar_mode"), trace_only=True)
-def aten_div_mode_int(self: TInt, other: TInt, rounding_mode: str) -> TInt:
+def aten_div_mode_int(
+    self: TInt, other: TInt, rounding_mode: Optional[str] = None
+) -> TensorType:
     """div.Tensor_mode(Tensor self, Tensor other, *, str? rounding_mode) -> Tensor
 
     Variant for integer inputs.
     """
-    # TODO(justinchuby): trace_only=False when we use opset19 which supports string comparison
-    assert rounding_mode in {"trunc", "floor"}
+    assert rounding_mode in {"trunc", "floor", None}
 
     quotient = op.Div(op.Cast(self, to=FLOAT.dtype), op.Cast(other, to=FLOAT.dtype))
 
     if rounding_mode == "trunc":
         # Rounds the results of the division towards zero.
         # Equivalent to C-style integer division
         result = aten_trunc(quotient)
-    else:  # rounding_mode == "floor"
+        return op.CastLike(result, self)
+    if rounding_mode == "floor":
         result = op.Floor(quotient)
+        return op.CastLike(result, self)
 
-    return op.CastLike(result, self)
+    assert rounding_mode is None
+    # When rounding_mode is None, the return type is float32
+    return quotient
 
 
 @torch_op("aten::dot")
@@ -8465,7 +8465,7 @@ def aten_triu_indices(row: int, col: int, offset: int = 0) -> TensorType:
     raise NotImplementedError()
 
 
-@torch_op("aten::trunc")
+@torch_op("aten::trunc", trace_only=True)
 def aten_trunc(self: TFloat) -> TFloat:
     """trunc(Tensor self) -> Tensor"""
 

tests/function_libs/torch_lib/ops_test_common.py

@@ -34,6 +34,7 @@
 
 import onnxscript
 import onnxscript.evaluator
+import onnxscript.ir.passes.common.unused_removal
 from onnxscript import ir
 from onnxscript.function_libs.torch_lib import graph_building
 from tests.function_libs.torch_lib import error_reproduction
@@ -389,6 +390,19 @@ def _format_model_and_input_information(onnx_model, inputs):
 }
 
 
+def add_torchlib_common_imports(model: ir.Model) -> None:
+    """Hack to add torchlib common imports to the model."""
+
+    model.opset_imports["pkg.onnxscript.torch_lib.common"] = 1
+    rank_func = ir.serde.deserialize_function(common_ops.Rank.to_function_proto())
+    is_scalar_func = ir.serde.deserialize_function(common_ops.IsScalar.to_function_proto())
+    model.functions[rank_func.identifier()] = rank_func
+    model.functions[is_scalar_func.identifier()] = is_scalar_func
+    removal_pass = onnxscript.ir.passes.common.unused_removal.RemoveUnusedFunctionsPass()
+    assert removal_pass.in_place
+    removal_pass(model)
+
+
 def dtype_op_schema_compatible(dtype: torch.dtype, schema: onnx.defs.OpSchema) -> bool:
     """Checks if the dtype is compatible with the schema.
 

tests/function_libs/torch_lib/ops_test_data.py

@@ -778,10 +778,7 @@ def _where_input_wrangler(
         test_class_name="TestOutputConsistencyEager",
         reason="fixme: off-by-one and inverted inf. https://github.com/microsoft/onnxscript/issues/989",
     ),
-    TorchLibOpInfo("div_mode_int", core_ops.aten_div_mode_int).skip(
-        variant_name="no_rounding_mode",
-        reason="this variation requires the rounding_mode argument",
-    ),
+    TorchLibOpInfo("div_mode_int", core_ops.aten_div_mode_int),
     TorchLibOpInfo("dot", core_ops.aten_dot),
     TorchLibOpInfo(
         "empty",