[Relax][Training] Optimizer API

python/tvm/relax/__init__.py

@@ -86,3 +86,6 @@
     TupleStructInfo,
     FuncStructInfo,
 )
+
+# Training utils
+from .training import optimizer

python/tvm/relax/block_builder.py

@@ -363,7 +363,7 @@ def dataflow(self) -> DataflowScope:
         """
         return DataflowScope(self)
 
-    def emit(self, expr: Expr) -> Var:
+    def emit(self, expr: Expr, name_hint: str = "") -> Var:
         """Emit an expr.
         This infers the shape and type of the expr, create a variable,
         and bind the expr to the variable.
@@ -373,12 +373,15 @@ def emit(self, expr: Expr) -> Var:
         expr : tvm.relax.Expr
             The Expr to be emitted.
 
+        name_hint : str
+            Name hint for the bound variable.
+
         Returns
         -------
         ret : tvm.relax.Var
             A newly created variable that gets bound to the input expr.
         """
-        return _ffi_api.BlockBuilderEmit(self, expr)  # type: ignore
+        return _ffi_api.BlockBuilderEmit(self, expr, name_hint)  # type: ignore
 
     def call_te(self, func: Callable, *args: Any, **kwargs: Any) -> Expr:
         """Generate a call node according to the te function.
@@ -581,7 +584,7 @@ def rx_func(x: Tensor((n,), "float32"), y: Tensor(((n + 1),), "float32"))
         """
         return self.emit(self.call_te(func, *args, **kwargs))
 
-    def match_cast(self, value: Expr, struct_info: StructInfo) -> Var:
+    def match_cast(self, value: Expr, struct_info: StructInfo, name_hint: str = "") -> Var:
         """Emit a MatchCast.
 
         Parameters
@@ -592,29 +595,37 @@ def match_cast(self, value: Expr, struct_info: StructInfo) -> Var:
         struct_info : StructInfo
             The struct info to be matched.
 
+        name_hint : str
+            Name hint for the bound variable.
+
         Returns
         -------
         ret : tvm.relax.Var
             A newly created variable that get bounds to be the casted result.
         """
-        return _ffi_api.BlockBuilderEmitMatchCast(self, value, struct_info)  # type: ignore
+        return _ffi_api.BlockBuilderEmitMatchCast(  # type: ignore
+            self, value, struct_info, name_hint
+        )
 
-    def emit_output(self, output: Union[Expr, Tuple, List[Expr]]) -> Var:
+    def emit_output(self, output: Union[Expr, Tuple, List[Expr]], name_hint: str = "") -> Var:
         """Emit output for the current dataflow block or function.
 
         Parameters
         ----------
         output : Expr | Tuple | List[Expr]
             The output of the current block/function.
 
+        name_hint : str
+            Name hint for the bound variable.
+
         Returns
         -------
         ret : tvm.relax.Var
             The return variable which gets bound to the output.
         """
         if isinstance(output, (list, tuple)):
             output = Tuple(output)
-        return _ffi_api.BlockBuilderEmitOutput(self, output)  # type: ignore
+        return _ffi_api.BlockBuilderEmitOutput(self, output, name_hint)  # type: ignore
 
     def emit_func_output(
         self,

python/tvm/relax/training/__init__.py

@@ -0,0 +1,19 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""The Relax training APIs."""
+
+from . import optimizer