code: def space_to_depth(data, block_size, layout="NCHW"):
      """Perform space to depth 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 decompose into channel dimension.
  
      layout : string
          Either NCHW or NHWC, indicating data layout.
  
      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
          output_shape = [
              in_n,
              in_c * block_size * block_size,
              tvm.tir.truncdiv(in_h, block_size),
              tvm.tir.truncdiv(in_w, block_size),
          ]
      elif layout == "NHWC":
          in_n, in_h, in_w, in_c = data.shape
          output_shape = [
              in_n,
              tvm.tir.truncdiv(in_h, block_size),
              tvm.tir.truncdiv(in_w, block_size),
              in_c * block_size * block_size,
          ]
      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_offset = tvm.tir.truncdiv(c, in_c)
          channel_idx = tvm.tir.truncmod(c, in_c)
          x_idx = tvm.tir.truncmod(block_offset, block_size)
          y_idx = tvm.tir.truncdiv(block_offset, block_size)
  
          if layout == "NCHW":
              output = data(
                  n,
                  channel_idx,
                  y_idx + (y * block_size),
                  x_idx + (x * block_size),
              )
          else:
              output = data(
                  n,
                  y_idx + (y * block_size),
                  x_idx + (x * block_size),
                  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="space_to_depth",
          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: space_to_depth(data, block_size, layout) and returns: te, n,  c,  y,  x, output, _get_pixel