code: def partition_for_vitis_ai(mod, params=None, dpu=None, **opts):
      """Partition the Relay expression for offloading operators to Vitis AI DPU
  
      Parameters
      ----------
      mod : Module
          The module to run passes on.
      params : Optional[Dict[str, NDArray]]
          Constant input parameters.
      dpu : str
          The DPU identifier (e.g. DPUCZDX8G-zcu104, DPUCADF8H)
  
      Returns
      -------
      ret : Module
      """
  
      if dpu is None:
          raise ValueError(
              "Please pass Vitis AI DPU identifier to the partitioning function"
          )
  
      if params:
          mod["main"] = bind_params_by_name(mod["main"], params)
  
      desired_layouts_in_partition = {
          "nn.conv2d": ["NHWC", "default"],
          "nn.upsampling": ["NHWC"],
          "image.resize2d": ["NHWC"],
      }
      desired_layouts_in_main = {
          "nn.conv2d": ["NCHW", "default"],
          "nn.upsampling": ["NCHW"],
          "image.resize2d": ["NCHW"],
      }
      seq = tvm.transform.Sequential(
          [
              transform.RemoveUnusedFunctions(),
              transform.ConvertLayout(desired_layouts_in_partition),
              transform.FoldConstant(),
              transform.InferType(),
              VitisAIAnnotationPass("vitis_ai", dpu, params),
              transform.MergeCompilerRegions(),
              transform.PartitionGraph(),
              transform.RemoveUnusedFunctions(),
              transform.ConvertLayout(desired_layouts_in_main),
              transform.FoldConstant(),
          ]
      )
  
      with tvm.transform.PassContext(opt_level=3):
          return seq(mod)
  
file_path: tvm/python/tvm/relay/op/contrib/vitis_ai.py
imports: import warnings
  import numpy
  from tvm import relay
  import tvm._ffi
  from tvm.relay import transform
  from tvm.relay.expr import Tuple, TupleGetItem
  from tvm.relay.build_module import bind_params_by_name
  from tvm.relay.op.annotation import compiler_begin, compiler_end
label_desc: Use image node_id 3 for calling a global function with example usage: partition_for_vitis_ai(mod, params, dpu) and returns: seq