class_docstr: A Stage represents schedule for one operation.
class_docstr_tok: ["A","Stage","represents","schedule","for","one","operation","."]
docstr: Defines the layout transformation for the current stage's tensor.
  
  The map from initial_indices to final_indices must be an
  invertible affine transformation.  This method may be called
  more than once for a given tensor, in which case each
  transformation is applied sequentially.
  
  If the stage is a ComputeOp, then the iteration order of the
  compute stage is rewritten to be a row-major traversal of the
  tensor, and the new loop iteration variables are returned.
  For all other stages, the loop iteration order is unmodified,
  and the return value is None.
  
  Parameters
  ----------
  mapping_function : Callable[..., List[tvm.tir.PrimExpr]]
  
      A callable that accepts N arguments of type tvm.tir.Var,
      and outputs a list of PrimExpr.  The input arguments
      represent the location of a value in the current stage's
      tensor, using the pre-transformation layout.  The return
      value of the function gives the location of that value in
      the current stage's tensor, using the post-transformation
      layout.
  
  Returns
  -------
  new_iter_vars : Optional[List[tvm.tir.IterVar]]
  
      If the stage is a ComputeOp, then the return will be the
      updated loop iteration variables over the data array, in
      the same order as the output values from the
      `mapping_function`.
  
      Otherwise, the return value is None.
  
  Examples
  --------
  .. code-block:: python
  
      # ``A`` is a tensor whose compute definition is in NHWC
      # format, and should be transformed into NCHWc format.
  
      s[A].transform_layout(
          lambda n,h,w,c: [n, c//4, h, w, c%4]
      )
  
  
  .. code-block:: python
  
      # ``A`` is a tensor whose compute definition is in an
      # arbitrary format, and should be transformed such that
      # the last index is split, with the slower-changing index
      # of the split placed at the slowest changing dimension.
  
      s[A].transform_layout(
          lambda *indices, i: [i//4, *indices, i%4]
      )
  
  .. code-block:: python
  
      # ``B`` is a tensor defined by te.compute to be a copy of
      # ``A`, and should be transformed such that ``B``'s layout
      # is a transpose of ``A``'s layout.  The loop iteration
      # that computes ``B`` will correspond to ``B``'s memory
      # layout.
  
      A = te.placeholder([n,m])
      B = te.compute(A.shape, lambda i,j: A[i,j])
      s = te.create_schedule(B.op)
  
      s[B].transform_layout(lambda i,j: [j,i])
code: def transform_layout(
      self, mapping_function: Callable[..., List[tvm.tir.PrimExpr]]
  ):
      """Defines the layout transformation for the current stage's tensor.
  
      The map from initial_indices to final_indices must be an
      invertible affine transformation.  This method may be called
      more than once for a given tensor, in which case each
      transformation is applied sequentially.
  
      If the stage is a ComputeOp, then the iteration order of the
      compute stage is rewritten to be a row-major traversal of the
      tensor, and the new loop iteration variables are returned.
      For all other stages, the loop iteration order is unmodified,
      and the return value is None.
  
      Parameters
      ----------
      mapping_function : Callable[..., List[tvm.tir.PrimExpr]]
  
          A callable that accepts N arguments of type tvm.tir.Var,
          and outputs a list of PrimExpr.  The input arguments
          represent the location of a value in the current stage's
          tensor, using the pre-transformation layout.  The return
          value of the function gives the location of that value in
          the current stage's tensor, using the post-transformation
          layout.
  
      Returns
      -------
      new_iter_vars : Optional[List[tvm.tir.IterVar]]
  
          If the stage is a ComputeOp, then the return will be the
          updated loop iteration variables over the data array, in
          the same order as the output values from the
          `mapping_function`.
  
          Otherwise, the return value is None.
  
      Examples
      --------
      .. code-block:: python
  
          # ``A`` is a tensor whose compute definition is in NHWC
          # format, and should be transformed into NCHWc format.
  
          s[A].transform_layout(
              lambda n,h,w,c: [n, c//4, h, w, c%4]
          )
  
  
      .. code-block:: python
  
          # ``A`` is a tensor whose compute definition is in an
          # arbitrary format, and should be transformed such that
          # the last index is split, with the slower-changing index
          # of the split placed at the slowest changing dimension.
  
          s[A].transform_layout(
              lambda *indices, i: [i//4, *indices, i%4]
          )
  
      .. code-block:: python
  
          # ``B`` is a tensor defined by te.compute to be a copy of
          # ``A`, and should be transformed such that ``B``'s layout
          # is a transpose of ``A``'s layout.  The loop iteration
          # that computes ``B`` will correspond to ``B``'s memory
          # layout.
  
          A = te.placeholder([n,m])
          B = te.compute(A.shape, lambda i,j: A[i,j])
          s = te.create_schedule(B.op)
  
          s[B].transform_layout(lambda i,j: [j,i])
  
      """
  
      ndim = len(self.op.output(0).shape)
      index_map, axis_separators = IndexMap.from_func_with_separators(
          mapping_function, ndim=ndim, index_dtype="int32"
      )
  
      new_iter_vars = _ffi_api.StageTransformLayout(
          self, index_map.initial_indices, index_map.final_indices
      )
      _ffi_api.StageSetAxisSeparators(self, axis_separators)
  
      return new_iter_vars or None
  
imports: import collections
  import inspect
  from typing import Callable, List
  import tvm._ffi
  from tvm._ffi.base import string_types
  from tvm.ir import container
  from tvm.runtime import Object, convert
  from tvm.tir import Buffer, IndexMap, IterVar, Var
  from .None import _ffi_api
  from .None import tensor
label_desc: Use image node_id 21 for calling the Stage obj's underlying member method code with example usage: obj.transform_layout(mapping_function) and returns: unknown