code: def depth_to_space(data, block_size, layout="NCHW", mode="DCR"):
      """Perform depth to space transformation on the data
  
      Parameters
      ----------
      data : tvm.te.Tensor
          4-D tensor in either NCHW or NHWC layout.
  
      block_size : int
          Size of blocks to compose from channel dimension.
  
      layout : string
          Either NCHW or NHWC, indicating data layout.
  
      mode : string
          Either DCR or CDR, indicates how channels should be accessed.
          In DCR, channels are interwoven in the Tensorflow style while
          in CDR channels are accessed sequentially as in Pytorch.
  
      Returns
      -------
      output : tvm.te.Tensor
          Output of shape [N, C / block_size**2, H * block_size, W * block_size]
      """
      if layout == "NCHW":
          in_n, in_c, in_h, in_w = data.shape
          channel_factor = tvm.tir.truncdiv(
              in_c, (block_size * block_size)
          )
          output_shape = [
              in_n,
              channel_factor,
              in_h * block_size,
              in_w * block_size,
          ]
      elif layout == "NHWC":
          in_n, in_h, in_w, in_c = data.shape
          channel_factor = tvm.tir.truncdiv(
              in_c, (block_size * block_size)
          )
          output_shape = [
              in_n,
              in_h * block_size,
              in_w * block_size,
              channel_factor,
          ]
      else:
          raise ValueError(
              "Only NCHW and NHWC layouts are currently supported."
          )
  
      def _get_indices(*indices):
          if layout == "NCHW":
              n, c, y, x = indices
          elif layout == "NHWC":
              n, y, x, c = indices
          return n, c, y, x
  
      def _get_pixel(n, c, y, x):
          block_x = tvm.tir.truncdiv(x, block_size)
          block_y = tvm.tir.truncdiv(y, block_size)
          idx_x = tvm.tir.truncmod(x, block_size)
          idx_y = tvm.tir.truncmod(y, block_size)
          if mode == "DCR":
              channel_idx = (
                  channel_factor * ((block_size * idx_y) + idx_x) + c
              )
          else:
              channel_idx = (c * block_size * block_size) + (
                  (block_size * idx_y) + idx_x
              )
  
          if layout == "NCHW":
              output = data(n, channel_idx, block_y, block_x)
          else:
              output = data(n, block_y, block_x, channel_idx)
          return output
  
      def _compute(*indices):
          n, c, y, x = _get_indices(*indices)
          return _get_pixel(n, c, y, x)
  
      return te.compute(
          output_shape,
          _compute,
          name="depth_to_space",
          tag=tag.INJECTIVE,
      )
  
imports: from __future__ import absolute_import
  import tvm
  from tvm import te
  from ..None import tag
label_desc: Use image node_id 1 for calling a global function with example usage: depth_to_space(data, block_size, layout, mode) and returns: te, n,  c,  y,  x, output, _get_pixel