Rename show_inferable_type_annotations to show_all_struct_info

Author: Lunderberg

include/tvm/node/script_printer.h

@@ -72,8 +72,40 @@ class PrinterConfigNode : public Object {
   bool syntax_sugar = true;
   /*! \brief Whether variable names should include the object's address */
   bool show_object_address = false;
-  /*! \brief Whether to show StructInfo that can be inferred from arguments */
-  bool show_inferable_type_annotations = true;
+
+  /*! \brief In Relax, whether to show all StructInfo annotations
+   *
+   * If true (default), all variable bindings will be annotated with
+   * the struct info of the variable being bound.
+   *
+   * If false, the annotations will only be shown when they are
+   * required for correct parsing of the Relax function.  For example,
+   * function parameters must always have struct info annotations, but
+   * the struct info for expressions within a function body may be inferred from their
+   * arguments, and are therefore
+   *
+   * Example:
+   *
+   *     # func.show(show_all_struct_info=True)
+   *     @R.function
+   *     def func(
+   *         A: R.Tensor((10, 20), dtype="float32"),
+   *         B: R.Tensor((10,20), dtype="float32"),
+   *     ) -> R.Tensor((10, 20), dtype="float32"):
+   *         C: R.Tensor((10,20), dtype="float32") = R.add(A, B2)
+   *         return C
+   *
+   *     # func.show(show_all_struct_info=False)
+   *     @R.function
+   *     def func(
+   *         A: R.Tensor((10, 20), dtype="float32"),
+   *         B: R.Tensor((10,20), dtype="float32"),
+   *     ) -> R.Tensor((10, 20), dtype="float32"):
+   *         C = R.add(A, B2)
+   *         return C
+   */
+  bool show_all_struct_info = true;
+
   /* \brief Object path to be underlined */
   Array<ObjectPath> path_to_underline = Array<ObjectPath>();
   /*! \brief Object path to be annotated. */
@@ -99,7 +131,7 @@ class PrinterConfigNode : public Object {
     v->Visit("num_context_lines", &num_context_lines);
     v->Visit("syntax_sugar", &syntax_sugar);
     v->Visit("show_object_address", &show_object_address);
-    v->Visit("show_inferable_type_annotations", &show_inferable_type_annotations);
+    v->Visit("show_all_struct_info", &show_all_struct_info);
     v->Visit("path_to_underline", &path_to_underline);
     v->Visit("path_to_annotate", &path_to_annotate);
     v->Visit("obj_to_underline", &obj_to_underline);

python/tvm/runtime/script_printer.py

@@ -44,7 +44,7 @@ class PrinterConfig(Object):
     num_context_lines: int
     syntax_sugar: bool
     show_object_address: bool
-    show_inferable_type_annotations: bool
+    show_all_struct_info: bool
     path_to_underline: Optional[List[ObjectPath]]
     path_to_annotate: Optional[Dict[ObjectPath, str]]
     obj_to_underline: Optional[List[Object]]
@@ -68,7 +68,7 @@ def __init__(
         num_context_lines: Optional[int] = None,
         syntax_sugar: bool = True,
         show_object_address: bool = False,
-        show_inferable_type_annotations: bool = True,
+        show_all_struct_info: bool = True,
         path_to_underline: Optional[List[ObjectPath]] = None,
         path_to_annotate: Optional[Dict[ObjectPath, str]] = None,
         obj_to_underline: Optional[List[Object]] = None,
@@ -91,7 +91,7 @@ def __init__(
             "num_context_lines": num_context_lines,
             "syntax_sugar": syntax_sugar,
             "show_object_address": show_object_address,
-            "show_inferable_type_annotations": show_inferable_type_annotations,
+            "show_all_struct_info": show_all_struct_info,
             "path_to_underline": path_to_underline,
             "path_to_annotate": path_to_annotate,
             "obj_to_underline": obj_to_underline,
@@ -135,7 +135,7 @@ def script(
         num_context_lines: int = -1,
         syntax_sugar: bool = True,
         show_object_address: bool = False,
-        show_inferable_type_annotations: bool = True,
+        show_all_struct_info: bool = True,
         path_to_underline: Optional[List[ObjectPath]] = None,
         path_to_annotate: Optional[Dict[ObjectPath, str]] = None,
         obj_to_underline: Optional[List[Object]] = None,
@@ -176,9 +176,10 @@ def script(
             Whether to output with syntax sugar, set false for complete printing.
         show_object_address: bool = False
             Whether to include the object's address as part of the TVMScript name
-        show_inferable_type_annotations: bool = True
-            Whether to show type annotations that can be inferred from previous
-            annotations.
+        show_all_struct_info: bool = True
+            If True (default), annotate all variable bindings with the struct
+            info of that variable.  If False, only add annotations where
+            required for unambiguous round-trip of Relax -> TVMScript -> Relax.
         path_to_underline : Optional[List[ObjectPath]] = None
             Object path to be underlined
         path_to_annotate : Optional[Dict[ObjectPath, str]] = None
@@ -192,6 +193,7 @@ def script(
         -------
         script : str
             The TVM Script of the given TVM IR
+
         """
         return _script(
             self,
@@ -211,7 +213,7 @@ def script(
                 num_context_lines=num_context_lines,
                 syntax_sugar=syntax_sugar,
                 show_object_address=show_object_address,
-                show_inferable_type_annotations=show_inferable_type_annotations,
+                show_all_struct_info=show_all_struct_info,
                 path_to_underline=path_to_underline,
                 path_to_annotate=path_to_annotate,
                 obj_to_underline=obj_to_underline,
@@ -287,7 +289,7 @@ def show(
         num_context_lines: int = -1,
         syntax_sugar: bool = True,
         show_object_address: bool = False,
-        show_inferable_type_annotations: bool = True,
+        show_all_struct_info: bool = True,
         path_to_underline: Optional[List[ObjectPath]] = None,
         path_to_annotate: Optional[Dict[ObjectPath, str]] = None,
         obj_to_underline: Optional[List[Object]] = None,
@@ -351,9 +353,10 @@ def show(
             Whether to output with syntax sugar, set false for complete printing.
         show_object_address: bool = False
             Whether to include the object's address as part of the TVMScript name
-        show_inferable_type_annotations: bool = True
-            Whether to show type annotations that can be inferred from previous
-            annotations.
+        show_all_struct_info: bool = True
+            If True (default), annotate all variable bindings with the struct
+            info of that variable.  If False, only add annotations where
+            required for unambiguous round-trip of Relax -> TVMScript -> Relax.
         path_to_underline : Optional[List[ObjectPath]] = None
             Object path to be underlined
         path_to_annotate : Optional[Dict[ObjectPath, str]] = None
@@ -389,7 +392,7 @@ def show(
                 num_context_lines=num_context_lines,
                 syntax_sugar=syntax_sugar,
                 show_object_address=show_object_address,
-                show_inferable_type_annotations=show_inferable_type_annotations,
+                show_all_struct_info=show_all_struct_info,
                 path_to_underline=path_to_underline,
                 path_to_annotate=path_to_annotate,
                 obj_to_underline=obj_to_underline,

src/node/script_printer.cc

@@ -112,8 +112,8 @@ PrinterConfig::PrinterConfig(Map<String, ObjectRef> config_dict) {
   if (auto v = config_dict.Get("show_object_address")) {
     n->show_object_address = Downcast<IntImm>(v)->value;
   }
-  if (auto v = config_dict.Get("show_inferable_type_annotations")) {
-    n->show_inferable_type_annotations = Downcast<IntImm>(v)->value;
+  if (auto v = config_dict.Get("show_all_struct_info")) {
+    n->show_all_struct_info = Downcast<IntImm>(v)->value;
   }
 
   // Checking prefixes if they are valid Python identifiers.

src/script/printer/relax/binding.cc

@@ -45,7 +45,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
           using relax::StructInfo;
           using relax::MatchStructInfo;
           Optional<ExprDoc> ann = NullOpt;
-          if (d->cfg->show_inferable_type_annotations) {
+          if (d->cfg->show_all_struct_info) {
             ann = StructInfoAsAnn(n->var, n_p->Attr("var"), d, n->value);
           }
           ExprDoc rhs = Relax(d, "match_cast")

src/script/printer/relax/utils.h

@@ -91,7 +91,7 @@ inline Optional<ExprDoc> StructInfoAsAnn(const relax::Var& v, const ObjectPath&
       return NullOpt;
     }
   }
-  if (!d->cfg->show_inferable_type_annotations) {
+  if (!d->cfg->show_all_struct_info) {
     Optional<relax::StructInfo> inferred_sinfo = NullOpt;
     if (auto opt = rhs.as<relax::Call>()) {
       auto call = opt.value();

tests/python/relax/test_tvmscript_printer_relax.py

@@ -805,6 +805,12 @@ def func(x: R.Tensor((128, 128), dtype="float32")) -> R.Tensor((128, 128), dtype
 
 
 def test_hide_inferable_struct_info():
+    """Redundant type annotations can be omitted
+
+    When `show_all_struct_info=False`, TVMScript type annotations that
+    provide redundant struct info can be omitted.
+    """
+
     @R.function
     def func(A: R.Tensor([10, 20], "float32"), B: R.Tensor(ndim=2, dtype="float32")):
         # R.match_cast has the struct info as an argument, so it can
@@ -833,7 +839,7 @@ def func(A: R.Tensor([10, 20], "float32"), B: R.Tensor(ndim=2, dtype="float32"))
         return E
 
     _assert_print(
-        func.script(show_inferable_type_annotations=False),
+        func.script(show_all_struct_info=False),
         """
 # from tvm.script import relax as R