Retire get_schema in Op | chore!(api)

ghstack-source-id: 771b415
Pull Request resolved: #698

Signed-off-by: Justin Chu <[email protected]>

Author: justinchuby

onnxscript/converter.py

@@ -135,7 +135,7 @@ def __init__(self, name: Optional[Union[str, List[str]]], kind: ConverterExpress
     def is_const(self):
         return self.kind == ConverterExpressionKind.CONST
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.name
 
 
@@ -777,7 +777,7 @@ def translate_call_expr(self, node):
         else:
             args = [self.translate_opt_expr(x) for x in node.args]
             attrs = [self.translate_attr(x.arg, x.value) for x in node.keywords]
-        args = autocast.static_cast_inputs(self, callee.get_schema(), *args)
+        args = autocast.static_cast_inputs(self, callee.opschema, *args)
 
         # In ONNX, there is no way to explicitly specify a None value for an attribute.
         # Instead, the attribute must be omitted from the attribute list.
@@ -786,7 +786,7 @@ def translate_call_expr(self, node):
         return callee, args, attrs
 
     def _cast_like_binary_expression(self, op, left, right):
-        schema = op.get_schema()
+        schema = op.opschema
         return autocast.static_cast_inputs(self, schema, left, right)
 
     def translate_bool_op_expr(self, node: ast.BoolOp) -> ConverterExpression:

onnxscript/values.py

@@ -53,6 +53,8 @@ class Opset:
     Only a single instance of Opset is created for a given (domain, version) pair.
     """
 
+    domain: str
+    version: int
     cache: dict[tuple[type, str, int], Opset] = {}
 
     def __new__(cls, domain: str, version: int):
@@ -271,16 +273,13 @@ def __call__(self, *args, **kwargs):
         # FIXME(after #225): Move import to the top of the file.
         from onnxscript import evaluator  # pylint: disable=import-outside-toplevel
 
-        schema = self.get_schema()
+        schema = self.opschema
         if schema is None:
             raise RuntimeError(
                 f"Op '{self.name}' does not have an OpSchema and cannot be evaluated."
             )
         return evaluator.default().eval(schema, args, kwargs)
 
-    def is_single_op(self) -> bool:
-        return isinstance(self.opname, str)
-
     @property
     def name(self) -> str:
         return self._name
@@ -293,22 +292,16 @@ def opset(self) -> Opset:
     def opschema(self) -> Optional[onnx.defs.OpSchema]:
         return self._opschema
 
-    def get_schema(self) -> Optional[onnx.defs.OpSchema]:
-        """Returns the ONNX OpSchema for this op."""
-        if self.opschema is not None:
-            return self.opschema
-        return self.opset[self.opname]
-
     def has_schema(self) -> bool:
         """Returns True if this op has an OpSchema."""
-        return self.get_schema() is not None
+        return self.opschema is not None
 
     def param_schemas(self) -> Optional[tuple[ParamSchema, ...]]:
         """Returns the parameter schemas for this op, if it has one."""
         if self._param_schemas is not None:
             return self._param_schemas
 
-        op_schema = self.get_schema()
+        op_schema = self.opschema
         if op_schema is None:
             return None