Add support for type annotation

Signed-off-by: Ganesan Ramalingam <[email protected]>

Author: gramalingam

onnxscript/converter.py

@@ -464,6 +464,16 @@ def _eval_constant_expr(self, expr: ast.AST) -> PyValue:
                 )
             ) from e
 
+    def _get_type_annotation(self, annotation: ast.Expr) -> Optional[ta.TypeAnnotationValue]:
+        typeinfo = self._eval_constant_expr(annotation)
+        if not ta.is_valid_type(typeinfo):
+            self.warn(
+                annotation,
+                "Unsupported type annotation.",
+            )
+            typeinfo = None
+        return typeinfo
+
     def _translate_attr(
         self,
         attr_name: str,
@@ -985,9 +995,11 @@ def assign(lhs: ast.AST, rhs: ast.AST) -> None:
                 lhs = lhs.id
                 t = self._translate_expr(rhs, lhs)
                 if isinstance(stmt, ast.AnnAssign):
-                    typeinfo = self._eval_constant_expr(stmt.annotation)
+                    typeinfo = self._get_type_annotation(stmt.annotation)
                 else:
                     typeinfo = None
+                if typeinfo is not None:
+                    irbuilder.set_type_info(t, typeinfo)
                 var = values.Dynamic(t, values.DynamicKind.Intermediate, info, typeinfo)
                 self._bind(lhs, var)
             elif isinstance(lhs, ast.Tuple):
@@ -1400,13 +1412,7 @@ def _translate_function_signature_common(
             else:
                 default_value = None
             if x.annotation:
-                typeinfo = self._eval_constant_expr(x.annotation)
-                if not ta.is_valid_type(typeinfo):
-                    self.warn(
-                        x.annotation,
-                        f"Unsupported type annotation for argument {x.arg}.",
-                    )
-                    typeinfo = None
+                typeinfo = self._get_type_annotation(x.annotation)
             else:
                 # The code can only be exported as a function.
                 typeinfo = None

onnxscript/converter_test.py

@@ -740,6 +740,37 @@ def model(x: FLOAT[10]) -> FLOAT[10]:
         model_false = make_model(False)
         onnxscript.testing.assert_isomorphic(model_false, sub_model.to_model_proto())
 
+    def test_type_annotation(self):
+        """Test that type annotations are processed correctly."""
+
+        @script()
+        def model(x: FLOAT[10]) -> FLOAT[10]:
+            temp: FLOAT[10] = op.Add(x, x)
+            y = op.Mul(temp, temp)
+            return y
+
+        model_proto = model.to_model_proto()
+        input_type = model_proto.graph.input[0].type.tensor_type
+        output_type = model_proto.graph.output[0].type.tensor_type
+        temp_value_info = None
+        for value_info in model_proto.graph.value_info:
+            if value_info.name == "temp":
+                temp_value_info = value_info
+                break
+        self.assertIsNotNone(temp_value_info, "ValueInfo for 'temp' not found in graph.")
+        temp_type = temp_value_info.type.tensor_type
+        self.assertEqual(temp_type.elem_type, onnx.TensorProto.FLOAT)
+        self.assertEqual(len(temp_type.shape.dim), 1)
+        self.assertEqual(temp_type.shape.dim[0].dim_value, 10)
+
+        self.assertEqual(input_type.elem_type, onnx.TensorProto.FLOAT)
+        self.assertEqual(len(input_type.shape.dim), 1)
+        self.assertEqual(input_type.shape.dim[0].dim_value, 10)
+
+        self.assertEqual(output_type.elem_type, onnx.TensorProto.FLOAT)
+        self.assertEqual(len(output_type.shape.dim), 1)
+        self.assertEqual(output_type.shape.dim[0].dim_value, 10)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnxscript/irbuilder.py

@@ -121,21 +121,24 @@ def to_function_proto(self) -> onnx.FunctionProto:
 # IRBuilder: abstracts out details of the IR in the python-to-IR converter
 
 
+def set_type_info(value: ir.Value, typeinfo: TypeAnnotationValue) -> None:
+    """Sets the type information on an IR value."""
+    try:
+        type_and_shape = ir.from_proto(typeinfo.to_type_proto())
+        value.type = type_and_shape.type
+        value.shape = type_and_shape.shape
+    except AttributeError:
+        pass
+    value.meta["typeinfo"] = typeinfo
+
+
 def _make_value(
     varname: str, typeinfo: TypeAnnotationValue, sourceinfo: SourceInfo
 ) -> ir.Value:
-    if typeinfo is None:
-        value = ir.Value(name=varname)
-    else:
-        try:
-            type_and_shape = ir.from_proto(typeinfo.to_type_proto())
-            value = ir.Value(
-                name=varname, type=type_and_shape.type, shape=type_and_shape.shape
-            )
-        except AttributeError:
-            value = ir.Value(name=varname)
+    value = ir.Value(name=varname)
     value.meta.setdefault("sourceinfo", sourceinfo)
-    value.meta.setdefault("typeinfo", typeinfo)
+    if typeinfo is not None:
+        set_type_info(value, typeinfo)
     return value