[UX] Highlight TVMScript with Pygments (#185)

python/gen_requirements.py

@@ -64,6 +64,7 @@
         (
             "Base requirements needed to install tvm",
             [
+                "Pygments",
                 "attrs",
                 "cloudpickle",
                 "decorator",

python/tvm/ir/module.py

@@ -292,6 +292,7 @@ def script(self, tir_prefix: str = "T", show_meta: bool = False) -> str:
     def show(self, style: Optional[str] = None) -> None:
         """
         A sugar for print highlighted TVM script.
+
         Parameters
         ----------
         style : str, optional

python/tvm/relax/expr.py

@@ -223,6 +223,35 @@ def __call__(self, *args):
         """
         return Call(self, args, None, None)
 
+    def script(self, show_meta: bool = False) -> str:
+        """Print relax.Function into TVMScript
+
+        Parameters
+        ----------
+        show_meta : bool
+            Whether to show meta information
+
+        Returns
+        -------
+        script : str
+            The TVM Script of the relax.Function
+        """
+        return tvm._ffi.get_global_func("script.AsRelaxScript")(self, show_meta)  # type: ignore
+
+    def show(self, style: str = "light") -> None:
+        """
+        A sugar for print highlighted TVM script.
+
+        Parameters
+        ----------
+        style : str, optional
+            Pygments styles extended by "light" (default) and "dark", by default "light"
+        """
+        from tvm.script.highlight import cprint  # pylint: disable=import-outside-toplevel
+
+        # Use deferred import to avoid circular import while keeping cprint under tvm/script
+        cprint(self, style=style)
+
 
 @tvm._ffi.register_object("relax.expr.ExternFunc")
 class ExternFunc(BaseFunc):

python/tvm/script/highlight.py

@@ -25,6 +25,7 @@
 from tvm.tir import PrimFunc
 
 
+
 def cprint(printable: Union[IRModule, PrimFunc], style: Optional[str] = None) -> None:
     """
     Print highlighted TVM script string with Pygments

python/tvm/tir/function.py

@@ -198,6 +198,7 @@ def script(self, tir_prefix: str = "T", show_meta: bool = False) -> str:
     def show(self, style: Optional[str] = None) -> None:
         """
         A sugar for print highlighted TVM script.
+
         Parameters
         ----------
         style : str, optional