text: # 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.
  # pylint: disable=invalid-name, unused-argument, missing-docstring, unused-import
  """
  Relay pass transformation infrastructure.
  """
  import functools
  import inspect
  import types
  import warnings
  
  import tvm.ir
  from tvm import relay, te
  from tvm.runtime import ndarray as _nd
  
  from . import _ffi_api
  from ..backend.utils import mangle_module_name
  
  
  def build_config(opt_level=2, required_pass=None, disabled_pass=None, trace=None):
      """Configure the build behavior by setting config variables. This function
      will be deprecated in TVM v0.7. Instead, we should directly use
      tvm.transform.PassContext.
  
      Parameters
      ----------
      opt_level: int, optional
          Optimization level. The optimization pass name and level are as the
          following:
  
          .. code-block:: python
  
              OPT_PASS_LEVEL = {
                  "SimplifyInference": 0,
                  "OpFusion": 1,
                  "FoldConstant": 2,
                  "FoldScaleAxis": 3,
                  "AlterOpLayout": 3,
                  "CanonicalizeOps": 3,
                  "CanonicalizeCast": 3,
                  "EliminateCommonSubexpr": 3,
                  "CombineParallelConv2D": 4,
                  "CombineParallelDense": 4,
                  "CombineParallelBatchMatmul": 4,
                  "FastMath": 4
              }
  
      required_pass: set of str, optional
          Optimization passes that are required regardless of optimization level.
  
      disabled_pass: set of str, optional
          Optimization passes to be disabled during optimization.
  
      trace: Callable[[IRModule, PassInfo, bool], None]
          A tracing function for debugging or introspection.
  
      Returns
      -------
      pass_context: PassContext
          The pass context for optimizations.
      """
      warnings.warn(
          "relay.build_config will be deprecated. Please use \
                    tvm.transform.PassContext directly",
          DeprecationWarning,
      )
      return tvm.transform.PassContext(opt_level, required_pass, disabled_pass, trace)
  
  
  @tvm._ffi.register_object("relay.FunctionPass")
  class FunctionPass(tvm.ir.transform.Pass):
      """A pass that works on each tvm.relay.Function in a module. A function
      pass class should be created through `function_pass`.
      """
  
  
  def InferType():
      """Infer the type of an expr.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered type inference pass.
      """
      return _ffi_api.InferType()
  
  
  def InferTypeLocal(expr):
      """Infer the type of a single expr, reusing type information to do so.
  
      This populates the checked_type field in expr. We assume existing type information
      in the graph is correct!
  
      Parameters
      ----------
      expr: relay.Expr
          The expression we want to know the type of
  
      Returns
      -------
      type: relay.Type
          The type of the expression
      """
      return _ffi_api.InferTypeLocal(expr)
  
  
  def FoldScaleAxis():
      """Fold the scaling of axis into weights of conv2d/dense. This pass will
      invoke both forward and backward scale folding.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass to fold expressions.
  
      Note
      ----
      Internally, we will call backward_fold_scale_axis before using
      forward_fold_scale_axis as backward folding targets the common conv->bn
      pattern.
      """
      return _ffi_api.FoldScaleAxis()
  
  
  def BackwardFoldScaleAxis():
      """Backward fold axis scaling into weights of conv2d/dense.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass to backward fold expressions.
  
      Note
      ----
      It is recommended to call backward_fold_scale_axis
      before using forward_fold_scale_axis as backward folding targets the common
      conv->bn pattern.
      """
      return _ffi_api.BackwardFoldScaleAxis()
  
  
  def RemoveUnusedFunctions(entry_functions=None):
      """Remove unused global relay functions in a relay module.
  
      Parameters
      ----------
      entry_functions: list[string]
          The set of entry functions to start from.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass to remove unused functions.
      """
      if entry_functions is None:
          entry_functions = ["main"]
      return _ffi_api.RemoveUnusedFunctions(entry_functions)
  
  
  def ForwardFoldScaleAxis():
      """Fold the scaling of axis into weights of conv2d/dense.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass to forward fold expressions.
  
      Note
      ----
      It is recommended to call backward_fold_scale_axis
      before using forward_fold_scale_axis, as backward folding targets the
      common conv->bn pattern.
      """
      return _ffi_api.ForwardFoldScaleAxis()
  
  
  def SimplifyInference():
      """Simplify the data-flow graph for inference phase. An simplified expression
      which is semantically equal to the input expression will be returned.
  
      Note that batch norms will only be simplified if their result is indexed at
      tuple index 0.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass to perform operator simplification.
  
      """
      return _ffi_api.SimplifyInference()
  
  
  def FastMath():
      """Converts the expensive non linear functions to their fast but approximate counterparts.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass to perform fast math operations.
      """
      return _ffi_api.FastMath()
  
  
  def CanonicalizeOps():
      """Canonicalize special operators to basic operators.
      This can simplify followed analysis, e.g. expanding bias_add to
      expand_dims and broadcast_add.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass performing the canonicalization.
      """
      return _ffi_api.CanonicalizeOps()
  
  
  def DeadCodeElimination(inline_once=False, ignore_impurity=False):
      """Remove expressions that do not have any users (dead code).
  
      Parameters
      ----------
      inline_once: Optional[Bool]
          Whether to inline a binding that is referenced exactly once.
      ignore_impurity: Optional[Bool]
          Whether to ignore possible side-effects in let-bound expressions.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass that eliminates the dead code in a Relay program.
      """
      return _ffi_api.DeadCodeElimination(inline_once, ignore_impurity)
  
  
  def LazyGradientInit():
      """Reduces memory usage of gradient tensors
  
      Parameters
      ----------
  
      Returns
      -------
      ret: tvm.transform.Pass
          A pass which delays and/or reduces memory allocation,
          by lazily allocating 0 or one filled tensors.
      """
      return _ffi_api.LazyGradientInit()
  
  
  def FoldConstantExpr(expr, mod):
      """Fold the constant expressions in a Relay program.
      Parameters
      ----------
      expr: Expr
          The expression to fold
      mod: IRModule
          The module the expr lives in (for global calls)
  
      Returns
      -------
      new_expr: Expr
          The expr after Constant Folding
      """
      return _ffi_api.FoldConstantExpr(expr, mod)
  
  
  def FoldConstant():
      """Fold the constant expressions in a Relay program.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass for constant folding.
      """
      return _ffi_api.FoldConstant()
  
  
  def FuseOps(fuse_opt_level=-1):
      """Fuse operators in an expr to a larger operator according to some rules.
  
      Parameters
      ----------
      fuse_opt_level : int
          The level of fuse optimization. -1 indicates that the level will be
          inferred from pass context.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass for operator fusion.
      """
      return _ffi_api.FuseOps(fuse_opt_level)
  
  
  def DefuseOps():
      """The inverse operation of FuseOps. It transforms a fused program returned by FuseOps into the
      program before FuseOps. (i.e., x == DefuseOps(FuseOps(x)))
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass for operator defusion.
      """
      return _ffi_api.DefuseOps()
  
  
  def CombineParallelConv2D(min_num_branches=3):
      """Combine multiple conv2d operators into one.
  
      Parameters
      ----------
      min_num_branches : int
          The minimum number of required parallel branches for performing this
          optimization.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass that combines parallel conv2d operators.
      """
      return _ffi_api.CombineParallelConv2D(min_num_branches)
  
  
  def CombineParallelDense(min_num_branches=3, to_batch=True):
      """Combine multiple dense operators into one. For example:
  
      .. code-block
                      data
              /              \
          dense (2,2)         dense (2,2)
              |                 |
          elemwise/bcast (2,2)  elemwise/bcast (2,2)
  
      Would become:
  
      .. code-block
  
                  data
                  |
              batch_matmul+elemwise/bcast (2,2,2)
  
      or (if to_batch=False)
  
      .. code-block
  
                  data
                  |
              dense+elemwise/bcast (2,2+2)
  
      Parameters
      ----------
      min_num_branches : int
          The minimum number of required parallel branches for performing this
          optimization.
  
      to_batch_matmul : bool
          If True, combine parallel dense ops into batch_matmul op.
          If False, combine parallel dense ops into dense op.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass that combines parallel dense operators.
      """
      return _ffi_api.CombineParallelDense(min_num_branches, to_batch)
  
  
  def CombineParallelBatchMatmul(min_num_branches=3):
      """Combine multiple batch matmul operators into one. For example:
  
      .. code-block
                               data (1, 2, 3)
                           /                  \
          batch_matmul(data, (1, 4, 3))    batch_matmul(data, (1, 5, 3))
              |                                |
          elemwise/bcast (1, 2, 4)         elemwise/bcast (1, 2, 5)
  
      Would become:
  
      .. code-block
  
                  data (1, 2, 3)
                  |
              batch_matmul(data, (1, 4+5, 3))
                  |
              elemwise/bcast (1 ,2, 4+5)
  
      Parameters
      ----------
      min_num_branches : int
          The minimum number of required parallel branches for performing this
          optimization.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass that combines parallel dense operators.
      """
      return _ffi_api.CombineParallelBatchMatmul(min_num_branches)
  
  
  def BatchingOps():
      """Batching parallel operators into one for Conv2D, Dense and BatchMatmul.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The sequential pass which apply batching for different operator types.
      """
      return tvm.transform.Sequential(
          [CombineParallelConv2D(), CombineParallelDense(), CombineParallelBatchMatmul()]
      )
  
  
  def AlterOpLayout():
      """Alternate the layouts of operators or replace primitive operators with
      other expressions.
      This pass can be used for computing convolution in custom layouts or
      other general weight pre-transformation.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass that alters the layout of operators.
      """
      return _ffi_api.AlterOpLayout()
  
  
  class LayoutConfig(object):
      """A structure for customizing the ConvertLayout pass."""
  
      current = None
  
      def __init__(self, skip_layers=None):
          self.skip_counter = 0
          self.skip_layers = skip_layers if skip_layers is not None else []
  
      def check_skip(self):
          skip = self.skip_counter in self.skip_layers
          self.skip_counter += 1
          return skip
  
      def reset(self):
          self.skip_counter = 0
          self.skip_layers = []
  
      def __enter__(self):
          self._old_manager = LayoutConfig.current
          LayoutConfig.current = self
          return self
  
      def __exit__(self, ptype, value, trace):
          LayoutConfig.current = self._old_manager
  
  
  def ConvertLayout(desired_layouts):
      """Given a dest layout, this pass transforms the expr such that most of the ops input data
      layout is changed to the dest layout. In ideal situation, there are only 2 layout transforms,
      one at the start and one at the end.
  
      This pass is not a part of relay.build and is expected to be called between framework-relay
      parser and relay.build call. This is very helpful for hardware backends that support/prefer only
      type of data layout.
  
      RFC - https://discuss.tvm.apache.org/t/layout-conversion-pass/4009
  
      This pass uses most of the AlterOpLayout and InferCorrectLayout infrastructure. We can define
      new layouts for conv2d ops for now. Most of the other operators try to adapt to their input
      layout using the InferCorrectLayout infrastructure.
  
      Parameters
      ----------
      desired_layouts : map of op_name to list of layouts
          Specify a mapping of operator names to a list of layouts to convert to, in the order
          defined by the operator. An example for nn.conv2d could be: {"nn.conv2d", ["NHWC", "OHWI]},
          where the first item in the list specifies the data layout and the second specifies the
          kernel layout.
  
      Returns
      -------
      pass: FunctionPass
        The pass.
      """
      return _ffi_api.ConvertLayout(desired_layouts)
  
  
  def Legalize(legalize_map_attr_name="FTVMLegalize"):
      """Legalizes an expression with another expression.
      This pass can be used to replace an expr with another expr for target
      dependent optimizations. For example, one expr, though semnatically
      equivalent to the other, can have better performance on a target. This pass
      can be used to legalize the expr in a target-dependent manner.
  
      Parameters
      ----------
      legalize_map_attr_name : str
          The Op's attr name which corresponds to the legalize rule function.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass that rewrites an expr.
      """
      return _ffi_api.Legalize(legalize_map_attr_name)
  
  
  def MergeComposite(pattern_table):
      """Merge multiple operators into a single composite relay function.
  
      Parameters
      ----------
      pattern_table : List[Tuple[str, tvm.relay.dataflow_pattern.DFPattern, Function]]
          A list of (pattern_name, pattern, check) tuples.
          The order of the patterns in the list will determine the order
          of priority in which they are matched.
          'check' is a function to check whether an extracted pattern matches.
          It can be implemented by pattern writer but if not specified it will
          always return True.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass that merges operators into a single composite
          relay function.
      """
      pattern_names = []
      patterns = []
      checks = []
      for tup in pattern_table:
          if len(tup) == 2:
              pattern_name, pattern = tup
              check = lambda extract: True
          elif len(tup) == 3:
              pattern_name, pattern, check = tup
  
          pattern_names.append(pattern_name)
          patterns.append(pattern)
          checks.append(check)
  
      return _ffi_api.MergeComposite(pattern_names, patterns, *checks)
  
  
  def MergeCompilerRegions():
      """Merge together compiler regions.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass that merges compiler regions.
      """
      return _ffi_api.MergeCompilerRegions()
  
  
  def ToANormalForm():
      """Turn Graph Normal Form expression into A Normal Form Expression.
      The scope of the root expression is the global scope.
      The scope of any non root expression is the least common ancestor of all it's scope.
      Values are ordered by post-DFS order in each scope.
  
      Returns
      -------
      ret : Union[tvm.transform.Pass, tvm.relay.Expr]
          The registered pass that transforms an expression into A Normal Form.
      """
      return _ffi_api.ToANormalForm()
  
  
  def ToANormalFormExpr(e):
      """ToANormalForm, but on expression level.
  
      Parameters
      ----------
      e : Expr
          The graph expression.
  
      Returns
      -------
      ret : Expr
          The transformed expresion.
      """
      return _ffi_api.ToANormalFormExpr(e)
  
  
  def ToBasicBlockNormalForm():
      """Turn an expression to Basic Block Normal Form.
      We define a block as a group of expressions implied by the scope structure.
      Each graph node can only belong to a single block.
      For any value that is being used in multiple blocks, it has to be referred
      by a Var which is defined in a block, whose scope is the least common ancestor
      of blocks this value is used.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass that transforms an expression into Basic Block Normal Form.
      """
      return _ffi_api.ToBasicBlockNormalForm()
  
  
  def ToCPS(expr, mod=None):
      """
      Turn expression into continuation passing style(CPS).
  
      Every intermediate compute will be passed to a continuation.
  
      Returns
      -------
      result: tvm.transform.Pass
          The registered pass that transforms an expression into CPS.
      """
      return _ffi_api.to_cps(expr, mod)
  
  
  def EtaExpand(expand_constructor=False, expand_global_var=False):
      """Add abstraction over a constructor or global variable bound to a function
  
      Parameters
      ----------
      expand_constructor: bool
          Whether to expand constructors.
  
      expand_global_var: bool
          Whether to expand global variables.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass that eta expands an expression.
      """
      return _ffi_api.EtaExpand(expand_constructor, expand_global_var)
  
  
  def ToGraphNormalForm():
      """Turn a Relay program in A Normal Form into Graph Normal Form
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass that transforms an expression into Graph Normal Form.
      """
      return _ffi_api.ToGraphNormalForm()
  
  
  def EliminateCommonSubexpr(fskip=None):
      """Eliminate common subexpressions.
  
      Parameters
      ----------
      fskip: Callable
          The callback function that decides whether an expression should be
          skipped.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass that eliminates common subexpressions.
      """
      return _ffi_api.EliminateCommonSubexpr(fskip)
  
  
  def PartialEvaluate():
      """Evaluate the static fragment of the code.
  
      Note
      ----
      This transformation could be either `Module -> Module` or `Expr -> Expr`.
      It will directly transform the input expression to a new one if the target
      expression is provided. Otherwise, it will rely on the pass manager to
      carry out transformation.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass that performs partial evaluation on an expression.
      """
      return _ffi_api.PartialEvaluate()
  
  
  def CanonicalizeCast():
      """
      Canonicalize cast expressions to make operator fusion more efficient.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass that canonicalizes cast expression.
      """
      return _ffi_api.CanonicalizeCast()
  
  
  def LambdaLift():
      """
      Lift the closure to global function.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass that lifts the lambda function.
      """
      return _ffi_api.LambdaLift()
  
  
  def PartitionGraph(mod_name="default", bind_constants=True):
      """Partition a Relay program into regions that can be executed on different
      backends.
  
      Parameters
      ----------
      mod_name : string
          Controls the prefix of the name of each partitioned subraph.
          If `mod_name` is None, then `tvmgen_` prefix is used.
          Otherwise, `tvmgen_mod_name_` prefix is used.
  
      bind_constants: bool
          Whether or not to bind constants in partitioned subgraphs. Note that the codegen needs
          to maintain the bound constants; Otherwise the constants will be maintained by
          the metadata module. So it is recommended for C-source based codegens to
          set bind_constants=False to avoid embedding large constants in a C source file.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass that partitions the Relay program.
      """
      mod_name = mangle_module_name(mod_name)
      return _ffi_api.PartitionGraph(mod_name, bind_constants)
  
  
  def AnnotateTarget(targets, include_non_call_ops=True):
      """Annotate ops in an experession with a provied compiler/target and then
      use it for codegen.
  
      Parameters
      ----------
      targets : str or List[str]
          The list of target compilers used for codegen.
      include_non_call_ops : boolean
          If True then non-call ops also will be annotated with targets
          If False then non-call ops will not be processed
  
      Returns
      -------
      ret : tvm.transform.Pass
          The annotated pass that wrapps ops with subgraph_start and
          subgraph_end.
      """
      if isinstance(targets, str):
          targets = [targets]
      return _ffi_api.AnnotateTarget(
          [tvm.runtime.container.String(t) for t in targets], include_non_call_ops
      )
  
  
  def DynamicToStatic():
      """If possible, convert tvm.relay.dynamic* ops to static versions
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass for dynamic->static conversion.
      """
      return _ffi_api.DynamicToStatic()
  
  
  def Inline():
      """Perform inlining on the given Relay IR module. The global functions that
      are marked as `inline` should be always inlined. A cost model will be
      needed in the future to decide if it is profitable to inline the function.
  
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass that performs inlining for a Relay IR module.
      """
      return _ffi_api.Inline()
  
  
  def InlineComposites(target):
      """Perform inlining on the given Relay IR module. The functions originate
      from the MergeComposite pass based on an input pattern table will fold back
      to main. Currently, this is used for the TRT BYOC which expects a single
      primitive function to operate on.
  
      Parameters
      ----------
      target: str
          The byoc target for which ops need to fold back to primitive function.
      Returns
      -------
      ret: tvm.transform.Pass
          The registered pass that performs inlining for a Relay IR module.
      """
      return _ffi_api.InlineComposites(target)
  
  
  def gradient(expr, mod=None, mode="higher_order"):
      """
      Transform the input function,
      returning a function that calculate the original result,
      paired with gradient of the input.
  
      Parameters
      ----------
      expr : tvm.relay.Expr
          The input expression, which is a Function or a GlobalVar.
  
      mod : Optional[tvm.IRModule]
  
      mode : Optional[String]
          The mode of the automatic differentiation algorithm.
          'first_order' only works on first order code, but will not produce
          reference nor closure.
          'higher_order' works on all code using reference and closure.
  
      Returns
      -------
      expr : tvm.relay.Expr
        The transformed expression.
      """
      if mode == "first_order":
          warnings.warn(
              "using transform.gradient for first-order AD is deprecated, please use the"
              "FirstOrderGradient module pass",
              DeprecationWarning,
          )
          if mod is not None:
              raise RuntimeError(
                  "to run first-order AD on a module, please use the FirstOrderGradient module pass."
              )
          return FirstOrderGradient()(tvm.IRModule.from_expr(expr))["main"]
      if mode == "higher_order":
          return _ffi_api.gradient(expr, mod)
      raise Exception("unknown mode")
  
  
  def FirstOrderGradient():
      """
      Transforms all global functions in the module to return the original result, paired with the
      gradients of the inputs. This pass transforms each global function independently and does not
      support interprocedural AD. Additionally, this pass does not support any control-flow or
      references, and should only be used on pure data-flow graphs.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered FirstOrderGradient pass.
      """
      return _ffi_api.FirstOrderGradient()
  
  
  def Defunctionalization(func, mod):
      """
      Performs defunctionalization on func,
      transforming func from a higher-order program to a first-order program.
  
      At each call site, the function is cloned and type parameters are substituted in.
      Function arguments are encoded as datatypes
      and additional apply functions are used for application.
  
      Parameters
      ----------
      func : tvm.relay.Function
          The input function, which should not be polymorphic or be higher-order.
          This is because all types must be known and we can't encode function arguments
          to the program itself.
  
      mod : tvm.IRModule
          The IRModule containing function and type definitions,
          which is also mutated during this pass.
  
      Returns
      -------
      expr : tvm.relay.Function
        The output function.
      """
      return _ffi_api.Defunctionalization(func, mod)
  
  
  def to_cps(func, mod=None):
      """
      Turn expression into CPS expression.
  
      Every intermediate compute will be passed to a continuation.
  
      Parameters
      ----------
      func: tvm.relay.Function
          The input function.
  
      mod: Optional[tvm.IRModule]
          The global module.
  
      Returns
      -------
      result: tvm.relay.Function
        The output function.
      """
      use_mod = mod if mod is not None else tvm.ir.IRModule()
      return _ffi_api.to_cps(func, use_mod)
  
  
  def un_cps(func):
      """
      Turn an cps function into a Function without the continuation argument.
  
      Note that this will not give the exact same interface as before cps:
        If the input/output is higher order, they will still be in cps form.
  
      Parameters
      ----------
      func: tvm.relay.Function
          The input function
  
      Returns
      -------
      result: tvm.relay.Function
          The output function
      """
      return _ffi_api.un_cps(func)
  
  
  def _wrap_class_function_pass(pass_cls, pass_info):
      """Wrap a python class as function pass"""
  
      class PyFunctionPass(FunctionPass):
          """Internal wrapper class to create a class instance."""
  
          def __init__(self, *args, **kwargs):
              # initialize handle in cass pass_cls creation failed.fg
              self.handle = None
              inst = pass_cls(*args, **kwargs)
  
              # it is important not to capture self to
              # avoid a cyclic dependency
              def _pass_func(func, mod, ctx):
                  return inst.transform_function(func, mod, ctx)
  
              self.__init_handle_by_constructor__(_ffi_api.MakeFunctionPass, _pass_func, pass_info)
              self._inst = inst
  
          def __getattr__(self, name):
              # fall back to instance attribute if there is not any
              return self._inst.__getattribute__(name)
  
      functools.update_wrapper(PyFunctionPass.__init__, pass_cls.__init__)
      PyFunctionPass.__name__ = pass_cls.__name__
      PyFunctionPass.__doc__ = pass_cls.__doc__
      PyFunctionPass.__module__ = pass_cls.__module__
      return PyFunctionPass
  
  
  def function_pass(pass_func=None, opt_level=None, name=None, required=None):
      """Decorate a function pass.
  
      This function returns a callback when pass_func
      is provided. Otherwise, it returns the created function pass using the
      given optimization function.
  
      Parameters
      ----------
      pass_func : Optional[Callable[(Function, Module, PassContext) -> Function]]
          The transformation function or class.
  
      opt_level : int
          The optimization level of this module pass.
  
      name : Optional[str]
          The name of the function pass. The name could be empty. In this case, the
          name of the optimization function will be used as the pass name.
  
      required : Optional[List[str]]
          The list of passes that the module pass is dependent on.
  
      Returns
      -------
      create_function_pass : Union[Callable, FunctionPass]
  
          A decorator will be returned if pass_func is not provided,
          otherwise return the decorated result.
          The returned decorator has two behaviors depending on the input:
          A new FunctionPass will be returned when we decorate a pass function.
          A new FunctionPass class will be returned when we decorate a class type.
  
      Examples
      --------
      The following code block decorates a function pass class.
  
      .. code-block:: python
  
          @relay.transform.function_pass(opt_level=1)
          class TestReplaceFunc:
              def __init__(self, new_func):
                  self.new_func = new_func
  
              def transform_function(self, func, mod, ctx):
                  # just for demo purposes
                  # transform func to new_func
                  return self.new_func
  
          x = relay.var("x", shape=(10, 20))
          f1 = relay.Function([x], x)
          f2 = relay.Function([x], relay.log(x))
          # fpass is now a special pass that replaces every
          # function to f1
          fpass = TestReplaceFunc(f1)
          # now every function in input_mod is replaced by f1
          res_mod = fpass(input_mod)
  
  
      The following code creates a function pass by decorating
      a user defined transform function.
  
      .. code-block:: python
  
          @relay.transform.function_pass(opt_level=2)
          def transform(func, mod, ctx):
              # my transformations here.
              return func
  
          function_pass = transform
          assert isinstance(function_pass, transform.FunctionPass)
          assert function_pass.info.opt_level == 2
  
          # Given a module m, the optimization could be invoked as the follwoing:
          updated_mod = function_pass(m)
          # Now constant folding should have been applied to every function in
          # the provided module m. And the updated module will be returned.
      """
  
      if opt_level is None:
          raise ValueError("Please provide opt_level for the function pass.")
  
      required = required if required else []
      if not isinstance(required, (list, tuple)):
          raise TypeError("Required is expected to be the type of " + "list/tuple.")
  
      def create_function_pass(pass_arg):
          """Internal function that creates a function pass"""
          fname = name if name else pass_arg.__name__
          info = tvm.transform.PassInfo(opt_level, fname, required)
          if inspect.isclass(pass_arg):
              return _wrap_class_function_pass(pass_arg, info)
          if not isinstance(pass_arg, (types.FunctionType, types.LambdaType)):
              raise TypeError("pass_func must be a callable for Module pass")
          return _ffi_api.MakeFunctionPass(pass_arg, info)
  
      if pass_func:
          return create_function_pass(pass_func)
      return create_function_pass
  
  
  @function_pass(opt_level=1)
  class ChangeBatch:
      """
      Change the batch size.
  
      Parameters
      ----------
      data: Dict[relay.Var, int]
        A dictionary of all the params to change.
        The keys are all params, and the values are which dimension hold the batch.
  
      batch_size: int
        The batch size to change to.
  
      Returns
      -------
      pass: FunctionPass
        The pass.
      """
  
      def __init__(self, data, batch_size=16):
          self.data = data
          self.batch_size = batch_size
  
      def transform_function(self, func, mod, ctx):
          func = relay.Function(func.params, func.body, None, func.type_params, func.attrs)
          change_batch = self
  
          class ChangeBatchMutator(tvm.relay.ExprMutator):
              def visit_var(self, var):
                  if var in change_batch.data:
                      ty = var.type_annotation
                      new_shape = list(ty.shape)
                      new_shape[change_batch.data[var]] = change_batch.batch_size
                      return relay.Var(var.name_hint, relay.TensorType(new_shape, ty.dtype))
                  return var
  
          return ChangeBatchMutator().visit(func)
  
  
  def DenseToSparse(weight_name, weight_shape):
      """
      Rewrite qualified ```nn.dense operation``` to ```nn.sparse_dense```
      This pass is used in ```data_dep_optimization.bsr_dense```
      Parameters of this pass is generated by ```analysis.sparse_dense.process_params```
  
      Parameters
      ----------
      weight_name: Array[String]
        Names of weights which qualified sparse contrains
  
      weight_shape: Array[Array[IntImm]]
        Weights shape in BSR format.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered DenseToSparse pass.
      """
      return _ffi_api.DenseToSparse(weight_name, weight_shape)
  
  
  def Conv2dToSparse(weight_name, weight_shape, layout, kernel_size):
      """
      Rewrite qualified ```nn.conv2d operation``` to ```nn.sparse_conv2d```
  
      Parameters
      ----------
      weight_name: Array[String]
        Names of weights which qualified sparse contrains
  
      weight_shape: Array[Array[IntImm]]
        Weights shape in BSR format.
  
      layout : str
          layout of data
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered DenseToSparse pass.
      """
      return _ffi_api.Conv2dToSparse(weight_name, weight_shape, layout, kernel_size)
  
  
  def Conv2dToSparse2(layout, kernel_size, blocksize, sparsity_threshold):
      """
      Rewrite freezed ```nn.conv2d``` operation to ```nn.sparse_conv2d```
  
      Parameters
      ----------
      layout : str
          layout of data
  
      kernel_size : int
          kernel size of conv2d
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered DenseToSparse pass.
      """
      return _ffi_api.Conv2dToSparse2(layout, kernel_size, *blocksize, sparsity_threshold)
  
  
  def SimplifyFCTranspose(target_weight_name):
      """
      Rewrite ```y = nn.dense(x, transpose(w, [1, 0]))``` to ```y = nn.dense(x, wt)```
      This pass is used in ```data_dep_optimization.simplify_fc_transpose```
  
      Parameters
      ----------
      weight_name: Array[String]
        Names of weights which qualified ```y = nn.dense(x, transpose(w, [1, 0]))```
        This parameter is generated by ```analysis.search_fc_transpose``` function
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered SimplifyFCTranspose pass.
      """
      return _ffi_api.SimplifyFCTranspose(target_weight_name)
  
  
  def SimplifyExpr():
      """
      Simplify the Relay expression, including merging consecutive reshapes.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered SimplifyExpr pass.
      """
      return _ffi_api.SimplifyExpr()
  
  
  def PlanDevices(config):
      """
      Uses existing "on_device" and "device_copy" calls to infer the virtual device on which
      every Relay sub-expression should run and the result stored. Captures the result of that
      analysis using new "on_device" and "device_copy" calls. Sub-expressions which are
      not otherwise constrained are assigned to the default primitive virtual device describe by
      config. However data and computations which must be hosted on a CPU (such as shapes and
      shape functions) use the host virtual device of the config.
  
      Parameters
      ----------
      config : tvm.CompilationConfig
          The compilation configuration, specifying available targets and default devices.
  
      Returns
      -------
      ret : tvm.transforms.Pass
          The pass.
      """
      return _ffi_api.PlanDevices(config)
  
  
  def ManifestLifetimes():
      """
      Manifest the lifetimes of variables after allocations have been manifested, by inserting kill
      operations once variables become dead.
      """
      return _ffi_api.ManifestLifetimes()
  
  
  def FoldExplicitPadding():
      """
      FoldExplicitPadding finds explict padding before an op that can support
      implicit padding and fuses them.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered ImplicitPadding pass.
      """
      return _ffi_api.FoldExplicitPadding()
  
  
  def AnnotateSpans():
      """
      Annotate a program with span information by first generating its textual
      representation and then parsing it back into a Relay AST annotated with
      span information.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered AnnotateSpans pass.
      """
      return _ffi_api.AnnotateSpans()
  
  
  def FakeQuantizationToInteger(hard_fail=False, use_qat=False):
      # pylint: disable=anomalous-backslash-in-string
      """
      Find regions of the graph of the form
  
      .. code-block:: text
  
          x    w
          |    |
          dq   dq
           \\   /
            op1
             |
            op2
             |
             q
  
      where ``q == qnn.quantize`` and ``dq = qnn.dequantize``
      and rewrite them into integer versions of ``op1`` and ``op2``
  
      Rules for rewriting indivdual ops are in fake_quantization_to_integer.py
  
      Parameters
      ----------
      hard_fail : boolean
          How do deal with errors during graph rewriting.
          If true, raise an error.
          If false, skip rewriting the subgraph.
  
      use_qat : boolean
          To perform an additional QAT pass - convert enabled operations with dequantized inputs.
          Example: in the graph above op2 is not registered with the FakeQuantizationToInteger
          attribute, op1 operation can still be converted. Converted pattern below:
  
          .. code-block:: text
  
              x    w
              |    |
              \\   /
                op1
                |
                dq
                |
                op2
                |
                q
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered FakeQuantizationToInteger pass.
      """
      return _ffi_api.FakeQuantizationToInteger(hard_fail, use_qat)
  
  
  def FlattenAtrousConv():
      # pylint: disable=anomalous-backslash-in-string
      """
      The purpose of this pass is to find a sequence of space_to_batch_nd-conv2d-batch_to_space_nd
      operations:
  
      .. code-block:: text
  
        x     w
        |     |
        s2b   |
         \\   /
          conv2d
           |
           b2s
  
      and convert them into subgraphs with a convolution with the modified "dilation" and
      recalculated "padding" parameters.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered FlattenAtrousConv pass.
      """
      return _ffi_api.FlattenAtrousConv()
  
  
  def ToMixedPrecision(mixed_precision_type="float16", missing_op_mode=1):
      """
      Automatic mixed precision rewriter. Rewrite an FP32 relay graph into a version
      where as many operations as possible are in the target mixed_precision_type.
  
      Parameters
      ----------
      mixed_precision_type: str
        The target datatype to transform operations in the graph to use.
  
      missing_op_mode: int
        Determines how to handle ops not registered with FTVMMixedPrecisionConversionType
          0: Does not allow any missing ops. Will throw errors when encountering any.
          1: Allow missing ops but emit warnings.
          2: Allow missing ops and silently ignore them.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass.
      """
      if missing_op_mode < 0 or missing_op_mode > 2:
          raise ValueError("Missing op mode is either 0, 1, or 2")
      return _ffi_api.ToMixedPrecision(mixed_precision_type, missing_op_mode)
  
  
  def SplitArgs(max_function_args):
      """Split function with huge number of arguments to smaller pieces.
  
      Returns
      -------
      ret : tvm.transform.Pass
          The registered pass for constant folding.
      """
      return _ffi_api.SplitArgs(max_function_args)
  