Merge remote-tracking branch 'upstream/main' into xiaowu/addOp(BatchNorm)

Author: xiaowuhu

onnxscript/evaluator.py

@@ -14,7 +14,7 @@
 import onnx
 import onnx.defs
 import onnx.helper
-from typing_extensions import TypeAlias
+from typing_extensions import Protocol, TypeAlias, runtime_checkable
 
 from onnxscript import autocast, irbuilder, onnx_opset, tensor, utils, values
 from onnxscript._internal import param_manipulation
@@ -117,7 +117,40 @@ def _unwrap_tensors_in_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]:
     return new_kwargs
 
 
-class Evaluator(abc.ABC):
+@runtime_checkable
+class Evaluator(Protocol):
+    """Protocol for evaluating ONNX ops."""
+
+    def eval(
+        self,
+        schema: onnx.defs.OpSchema,
+        inputs: Sequence[ExtendedModeValue],
+        attributes: Mapping[str, Any],
+    ):
+        """Evaluates an ONNX op.
+
+        Args:
+            schema: The OpSchema of the operator to evaluate.
+            inputs: The ONNX inputs to the op.
+            attributes: The ONNX attributes to the op.
+        """
+
+    def eval_function(
+        self,
+        function: values.OnnxFunction,
+        args: Sequence[ExtendedModeValue],
+        kwargs: Mapping[str, ExtendedModeValue],
+    ):
+        """Evaluates an OnnxFunction.
+
+        Args:
+            function: The OnnxFunction to evaluate.
+            args: The positional arguments to the function.
+            kwargs: The keyword arguments to the function.
+        """
+
+
+class BaseEvaluator(Evaluator, abc.ABC):
     """Base class for evaluation of ONNX ops.
 
     The execution of onnxscript functions in eager-mode is dispatched to an Evaluator
@@ -400,7 +433,7 @@ def _schema_id(schema: onnx.defs.OpSchema) -> tuple[str, str, int]:
     return schema.name, schema.domain, schema.since_version
 
 
-class ORTEvaluator(Evaluator):
+class ORTEvaluator(BaseEvaluator):
     """Evaluates ONNX ops using ONNX Runtime."""
 
     def _eval(self, schema, inputs, attributes, closure):

onnxscript/function_libs/torch_aten/graph_building.py

@@ -204,6 +204,9 @@ def __init__(self, graph: TorchScriptGraph):
     def graph(self) -> TorchScriptGraph:
         return self._graph
 
+    def eval(self, schema, inputs, attributes):
+        return self._graph.add_op_call(schema, inputs, attributes)
+
     @beartype
     def eval_function(  # type: ignore[override]
         self,
@@ -225,14 +228,6 @@ def eval_function(  # type: ignore[override]
                 attributes[name] = float(value)
         return self._graph.add_function_call(function, inputs, attributes)
 
-    def _eval(self, schema: onnx.defs.OpSchema, inputs, attributes, closure: Any):
-        del closure  # Unused
-
-        return self._graph.add_op_call(schema, inputs, attributes)
-
-    def eval(self, schema, inputs, attributes):
-        return self._eval(schema, inputs, attributes, closure=None)
-
 
 @beartype
 def _add_attribute_to_torchscript_node(

onnxscript/function_libs/torch_aten/ops/core.py

@@ -769,6 +769,12 @@ def aten_bitwise_not(self: TInt) -> TInt:
     return op.BitwiseNot(self)
 
 
+@torch_op("aten::bitwise_not", overload=True)
+def aten_bitwise_not_bool(self: BOOL) -> BOOL:
+    # bitwise_not(Tensor self) -> Tensor
+    return op.Not(self)
+
+
 @torch_op("aten::bitwise_or")
 def aten_bitwise_or(self: TInt, other: TInt) -> TInt:
     # bitwise_or.Tensor(Tensor self, Tensor other) -> Tensor
@@ -3162,9 +3168,10 @@ def aten_margin_ranking_loss(
 @torch_op("aten::masked_fill")
 def aten_masked_fill(self: TTensor, mask: BOOL, value: TTensor) -> TTensor:
     # masked_fill.Tensor(Tensor self, Tensor mask, Tensor value) -> Tensor
-    mask_cast = op.Cast(mask, to=BOOL.dtype)
+    # NOTE: Do not attempt to cast `mask` to BOOL because mask should not take any other types.
+    # `mask` coming in as other types is often an error and should fail the model.
     value_cast = op.CastLike(value, self)
-    return op.Where(mask_cast, value_cast, self)
+    return op.Where(mask, value_cast, self)
 
 
 def aten_masked_scatter(self: TensorType, mask: TensorType, source: TensorType) -> TensorType: