Minor cleanup of onnxscript converter (#2748)

* Add some missing documentation and type annotation.
* Cleanup some variable names
* Cleanup some ir.Value creation

---------

Signed-off-by: Ganesan Ramalingam <grama@microsoft.com>
Co-authored-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: gramalingam

onnxscript/_internal/analysis.py

@@ -144,7 +144,7 @@ def assigned_in_block(block: Sequence[ast.stmt]) -> Set[str]:
         error_message = self._formatter(stmt, f"Unsupported statement type {type(stmt)!r}.")
         raise ValueError(error_message)
 
-    def do_liveness_analysis(self, fun: ast.FunctionDef):
+    def do_liveness_analysis(self, fun: ast.FunctionDef) -> None:
         """Perform liveness analysis of the given function-ast."""
 
         def visit(stmt: ast.stmt, live_out: Set[str]) -> Set[str]:
@@ -212,7 +212,7 @@ def visitBlock(block: Sequence[ast.stmt], live_out: Set[str]) -> Set[str]:
         for s in reversed(fun.body):
             live = visit(s, live)
 
-    def exposed_uses(self, stmts: Sequence[ast.stmt]):
+    def exposed_uses(self, stmts: Sequence[ast.stmt]) -> set[str]:
         """Return the set of variables that are used before being defined by given block.
         In essence, this identifies the "inputs" to a given code-block.
         For example, consider the following code-block:
@@ -284,7 +284,7 @@ def visit(stmt: ast.stmt, live_out: Set[str]) -> Set[str]:
 
         return visitBlock(stmts, set())
 
-    def outer_scope_variables(self, fun: ast.FunctionDef):
+    def outer_scope_variables(self, fun: ast.FunctionDef) -> set[str]:
         """Return the set of outer-scope variables used in a nested function.
 
         Args:

onnxscript/_internal/converter.py

@@ -297,6 +297,12 @@ def _bind(self, name: str, val: LocalSymValue) -> None:
     def _lookup(
         self, name: str, info: sourceinfo.SourceInfo, raise_exception: bool = True
     ) -> SymValue:
+        """Maps a python variable name to the corresponding value used during translation.
+
+        Typically, a python variable X will correspond to an ONNX value Y. But other special
+        cases include: constant values or functions (mapped to Graph attributes), etc.
+        """
+
         for scope in self._locals:
             if name in scope:
                 return scope[name]
@@ -1158,13 +1164,9 @@ def _translate_if_stmt(self, stmt: ast.If) -> None:
         live_defs = list(live_def_set)
         test = self._translate_expr(stmt.test, "cond")
         lineno = self._source_of(stmt).lineno
-        thenGraph = self._translate_block(
-            stmt.body, f"thenGraph_{lineno}", live_defs, parent_stmt=stmt
-        )
+        thenGraph = self._translate_block(stmt.body, f"thenGraph_{lineno}", live_defs)
         thenAttr = self._make_onnx_attr("then_branch", thenGraph)
-        elseGraph = self._translate_block(
-            stmt.orelse, f"elseGraph_{lineno}", live_defs, parent_stmt=stmt
-        )
+        elseGraph = self._translate_block(stmt.orelse, f"elseGraph_{lineno}", live_defs)
         elseAttr = self._make_onnx_attr("else_branch", elseGraph)
 
         def rename(x):
@@ -1196,7 +1198,7 @@ def _translate_loop_stmt(self, loop_stmt: Union[ast.For, ast.While]) -> None:
         if isinstance(loop_stmt, ast.For):
             if not isinstance(loop_stmt.target, ast.Name):
                 self.fail(loop_stmt, "For loop target must be a single variable.")
-            p_loop_var = loop_stmt.target.id
+            python_loop_var_name = loop_stmt.target.id
             # iter
             iter = loop_stmt.iter
             assert isinstance(iter, ast.Call), "Loop bound not a call."
@@ -1210,8 +1212,8 @@ def _translate_loop_stmt(self, loop_stmt: Union[ast.For, ast.While]) -> None:
                 self.fail(loop_stmt, "Unsupported loop bound, it should be 'range(?)'.")
             assert not iter.keywords, "Unsupported loop bound."
             o_loop_bound = self._translate_expr(iter.args[0], "loop_bound")
-            o_cond_var = ir.Value(name=self.generate_unique_name("cond_in"))  # TODO(Rama)
-            i_cond_var = o_cond_var
+            onnx_cond_var = ir.Value(name=self.generate_unique_name("cond_in"))  # TODO(Rama)
+            i_cond_var = onnx_cond_var
             cond_while = None
             o_loop_condition = None  # No condition for a for loop.
         elif isinstance(loop_stmt, ast.While):
@@ -1222,11 +1224,11 @@ def _translate_loop_stmt(self, loop_stmt: Union[ast.For, ast.While]) -> None:
                     "Unexpected condition type {type(loop_stmt)!r} for a while loop, "
                     "it should be 'while <condition_name>:'.",
                 )
-            p_loop_var = "infinite_loop"
+            python_loop_var_name = "infinite_loop"
             o_loop_bound = None
             i_cond_var = ir.Value(name=test.id)  # TODO(Rama)
             cond_while = ir.Value(name=test.id)  # TODO(Rama)
-            o_cond_var = None
+            onnx_cond_var = None
             o_loop_condition = self._translate_name_expr(test)
             # we need to go through all the instructions to see
             # which instruction defines the condition test.id
@@ -1246,19 +1248,16 @@ def _translate_loop_stmt(self, loop_stmt: Union[ast.For, ast.While]) -> None:
 
         # build loop_body
         self._enter_scope("loop_body", loop_stmt)
-        o_loop_var = self.generate_unique_name(p_loop_var)
-        self._current_fn.append_parameter(
-            make_value(
-                o_loop_var,
-                onnx_types.INT64,
-                self._source_of(loop_stmt),
-            )
+        onnx_loop_var_name = self.generate_unique_name(python_loop_var_name)
+        onnx_loop_var = make_value(
+            onnx_loop_var_name,
+            onnx_types.INT64,
+            self._source_of(loop_stmt),
         )
+        self._current_fn.append_parameter(onnx_loop_var)
         self._bind(
-            p_loop_var,
-            values.Dynamic(
-                ir.Value(name=o_loop_var), values.DynamicKind.Loop, self._source_of(loop_stmt)
-            ),
+            python_loop_var_name,
+            values.Dynamic(onnx_loop_var, values.DynamicKind.Loop, self._source_of(loop_stmt)),
         )
 
         self._current_fn.append_parameter(
@@ -1270,17 +1269,19 @@ def _translate_loop_stmt(self, loop_stmt: Union[ast.For, ast.While]) -> None:
         )
 
         for pv in loop_state_vars:
-            ov = self.generate_unique_name(pv)
+            onnx_var_name = self.generate_unique_name(pv)
             # TODO: retrieve the annotation for variable pv is any is specified.
             # typeinfo = self._eval_constant_expr(pv.annotation)
             typeinfo = None
             self._current_fn.append_parameter(
-                make_value(ov, typeinfo, self._source_of(loop_stmt))
+                make_value(onnx_var_name, typeinfo, self._source_of(loop_stmt))
             )
             self._bind(
                 pv,
                 values.Dynamic(
-                    ir.Value(name=ov), values.DynamicKind.Loop, self._source_of(loop_stmt)
+                    ir.Value(name=onnx_var_name),
+                    values.DynamicKind.Loop,
+                    self._source_of(loop_stmt),
                 ),
             )
 
@@ -1313,7 +1314,7 @@ def _translate_loop_stmt(self, loop_stmt: Union[ast.For, ast.While]) -> None:
                 continue
             self._translate_stmt(s)
 
-        o_cond_out = self.generate_unique_name("cond_out")
+        onnx_cond_out_name = self.generate_unique_name("cond_out")
 
         if cond_while is not None:
             # Loop while
@@ -1324,35 +1325,35 @@ def _translate_loop_stmt(self, loop_stmt: Union[ast.For, ast.While]) -> None:
                     f"Unable to find condition variable {cond_while.name} in known "
                     f"variables {list(current_scope)!r}.",
                 )
-            o_cond_var = current_scope[cond_while.name].value
+            onnx_cond_var = current_scope[cond_while.name].value
 
         self.emit(
-            [o_cond_out],
+            [onnx_cond_out_name],
             values.Op(self.default_opset, operator_name),
-            [condition_name or o_cond_var],
+            [condition_name or onnx_cond_var],
             [],
         )
 
         self._current_fn.outputs.append(
             make_value(
-                o_cond_out,
+                onnx_cond_out_name,
                 onnx_types.BOOL,
                 self._source_of(loop_stmt),
             )
         )
         for pv in loop_state_vars:
-            ov = self._py_var_to_onnx_var(pv, self._source_of(loop_stmt))
-            if ov.name not in self._current_fn.assigned_names:
+            onnx_var = self._py_var_to_onnx_var(pv, self._source_of(loop_stmt))
+            if onnx_var.name not in self._current_fn.assigned_names:
                 # When converting the loop-body into a graph, we need to handle
                 # identity assignments of the form "x = y" inside the loop body
                 # specially if y represents a value computed outside the loop body.
                 # In this case, we create a copy of y, treating the statement as
                 # shorthand for "x = op.Identity(y)".
-                ov = self._emit_copy(ov, pv)
+                onnx_var = self._emit_copy(onnx_var, pv)
             # TODO: retrieve variable type for the annotation if any.
             typeinfo = None
             self._current_fn.outputs.append(
-                make_value(ov.name, typeinfo, self._source_of(loop_stmt))
+                make_value(onnx_var.name, typeinfo, self._source_of(loop_stmt))
             )
         body = self._exit_scope()
         inputs = [o_loop_bound, o_loop_condition] + [
@@ -1382,49 +1383,44 @@ def _translate_block(
         stmts: Sequence[ast.stmt],
         name: str,
         live_defs: Sequence[str],
-        parent_stmt: ast.stmt,
-    ):
-        """Translation of a statement-block to GraphProto attribute."""
-        info_stmt = stmts[0] if len(stmts) > 0 else parent_stmt
-        source = self._source_of(info_stmt)
+    ) -> ir.Graph:
+        """Translation of a then/else statement-block to an ir.Graph."""
         self._enter_scope(name, None)
         for s in stmts:
             self._translate_stmt(s)
-        for pvar in live_defs:
-            if pvar in self._current_scope():
-                pv_val = self._current_scope()[pvar]
-                output = self._to_onnx_var(pv_val, pvar)
+        for python_var in live_defs:
+            if python_var in self._current_scope():
+                python_var_value = self._current_scope()[python_var]
+                output = self._to_onnx_var(python_var_value, python_var)
                 if output.name not in self._current_fn.assigned_names:
+                    # TODO (Rama): Unclear how this can happen. If python_var is in current_scope,
+                    # then it should have been assigned a value in the current graph.
+                    #
                     # To return an outer-scope variable, an ONNX Graph has to
                     # use an explicit copy via Identity.
-                    output = self._emit_copy(output, pvar)
-                self._current_fn.outputs.append(
-                    make_value(
-                        output.name,
-                        pv_val.typeinfo,
-                        source,
-                    )
-                )
+                    output = self._emit_copy(output, python_var)
+                self._current_fn.outputs.append(output)
             else:
-                pv_val = None
+                python_var_value = None
                 for scope in self._locals:  # TODO: skip _current_scope
-                    if pvar in scope:
-                        pv_val = scope[pvar]
+                    if python_var in scope:
+                        python_var_value = scope[python_var]
                         break
-                if pv_val is None:
+                if python_var_value is None:
                     self.fail(
                         stmts[0],
-                        f"ir.Value {pvar} is not assigned a value along a conditional "
+                        f"ir.Value {python_var} is not assigned a value along a conditional "
                         f"branch, known variables: {list(self._locals)}.",
                     )
                 # introduce a copy
-                ovar = self._emit_copy(self._to_onnx_var(pv_val, pvar), pvar)
+                output = self._emit_copy(
+                    self._to_onnx_var(python_var_value, python_var), python_var
+                )
 
                 # TODO: retrieve the annotation if any.
-                typeinfo = None
-                self._current_fn.outputs.append(make_value(ovar.name, typeinfo, source))
-        graph = self._exit_scope()
-        return graph.graph
+                self._current_fn.outputs.append(output)
+        function_ir = self._exit_scope()
+        return function_ir.graph
 
     def _translate_nested_function_def(self, fn: ast.FunctionDef) -> None:
         """Translate a nested function definition."""
@@ -1465,14 +1461,13 @@ def _translate_function_signature_common(
                 self._current_fn.append_parameter(attr)
                 self._bind(x.arg, values.AttrRef(x.arg, typeinfo, self._source_of(x)))
             else:
-                self._current_fn.append_parameter(
-                    make_value(x.arg, typeinfo, self._source_of(x))
-                )
+                onnx_parameter = make_value(x.arg, typeinfo, self._source_of(x))
+                self._current_fn.append_parameter(onnx_parameter)
                 self._used_vars.add(x.arg)
                 self._bind(
                     x.arg,
                     values.Dynamic(
-                        ir.Value(name=x.arg), values.DynamicKind.Input, self._source_of(x)
+                        onnx_parameter, values.DynamicKind.Input, self._source_of(x)
                     ),
                 )
         if fn.returns:

onnxscript/_internal/irbuilder.py

@@ -24,8 +24,18 @@ def __init__(self, name: str, domain: str = "") -> None:
         super().__init__(domain, name, graph=graph, attributes=[])
         self.ordered_inputs_and_attrs: list[Union[ir.Value, ir.Attr]] = []
 
-        # a dictionary of nested function-definitions
+        # A dictionary of nested function-definitions: when an onnxscript function outer_f
+        # is translated, and it contains a nested function inner_f, then the inner function
+        # is translated and stored here. It will be used in any subsequent concrete execution
+        # of outer_f. Such nested functions are used in two different ways: it can be converted
+        # into a GraphProto to be stored as a graph-valued attribute of a node; alternatively,
+        # in a python-based execution mode, it can be called as a python function. It serves
+        # to enable a python-based debugging experience for higher-order functions such as Scan
+        # and SequenceMap.
         self.nested_functions: dict[str, IRFunction] = {}
+
+        # For nested functions, this dictionary maps outer-scope (python) variable names
+        # to their corresponding translated values.
         self.outer_scope_variables: dict[Any, Any] = {}
 
     @property

onnxscript/_internal/values.py

@@ -13,6 +13,7 @@
 import typing
 from enum import IntFlag
 from typing import (  # type: ignore[attr-defined]
+    TYPE_CHECKING,
     Any,
     Callable,
     ClassVar,
@@ -35,6 +36,9 @@
 from onnxscript.ir import _schemas
 from onnxscript.onnx_types import ONNXType
 
+if TYPE_CHECKING:
+    from onnxscript._internal.type_annotation import TypeAnnotationValue
+
 _R = TypeVar("_R")
 _P = ParamSpec("_P")
 
@@ -886,9 +890,13 @@ class DynamicKind(IntFlag):
 
 class Dynamic(SymbolValue):
     def __init__(
-        self, onnx_var: ir.Value, kind: DynamicKind, info: sourceinfo.SourceInfo, typeinfo=None
+        self,
+        onnx_var: ir.Value,
+        kind: DynamicKind,
+        info: sourceinfo.SourceInfo,
+        typeinfo: TypeAnnotationValue | None = None,
     ) -> None:
-        """Initializes Dynamic.
+        """Represents an ir.Value with some extra information.
 
         Arguments:
             onnx_var: the name of the ONNX variable used to represent this value