code: def _grid_sample_3d(
      data,
      grid,
      method="bilinear",
      layout="NCDHW",
      padding_mode="zeros",
      align_corners=True,
  ):
      """Applies bilinear/nearest sampling to input feature map.
  
      Given :math:`data` and :math:`grid` assuming NCDHW layout, then the output is computed by
  
      .. math::
  
          x_{src} = grid[batch, 0, z_{dst}, y_{dst}, x_{dst}] \\
          y_{src} = grid[batch, 1, z_{dst}, y_{dst}, x_{dst}] \\
          z_{src} = grid[batch, 2, z_{dst}, y_{dst}, x_{dst}] \\
          output[batch, channel, z_{src}, y_{dst}, x_{dst}]
          = G(data[batch, channel, z_{src}, y_{src}, x_{src})
  
      :math:`x_{dst}`, :math:`y_{dst}`, :math:`z_{dst}` enumerate all spatial locations
      in :math:`output`, and :math:`G()` denotes the interpolation method.
  
      The out-boundary points will be padded with zeros if padding_mode is "zeros", or
      border pixel value if padding_mode is "border", or
      inner pixel value if padding_mode is "reflection".
  
      The left-top corner (-1, -1, -1) and right-bottom corner (1, 1, 1) in grid will be map to
      (0, 0, 0) and (d - 1, h - 1, w - 1) of data if align_corners is "True", or
      (-0.5, -0.5, -0.5) and (d - 0.5, h - 0.5, w - 0.5) of data if align_corners is "False".
  
      The shape of the output will be
      (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3], grid.shape[4]).
  
      The operator assumes that :math:`grid` has been normalized to [-1, 1].
  
      grid_sample often cooperates with affine_grid which generates sampling grids for grid_sample.
  
      Parameters
      ----------
      data : tvm.Tensor
          5-D with shape [batch, in_channel, in_depth, in_height, in_width]
  
      grid : tvm.Tensor
          5-D with shape [batch, 3, out_depth, out_height, out_width]
  
      method : str
          The interpolation method "nearest", "bilinear"("trilinear") are supported.
  
      layout : str
          The layout of input data and the output.
  
      padding_mode : str
          The padding mode for outside grid values, "zeros", "border", "reflection" are supported.
  
      align_corners: bool
          Geometrically, we consider the pixels of the input as squares rather than points.
          If set to "True", the extrema ("-1" and "1") are considered as referring
          to the center points of the input corner pixels. If set to "False", they
          are instead considered as referring to the corner points of the input corner
          pixels, making the sampling more resolution agnostic.
  
      Returns
      -------
      Output : tvm.Tensor
          5-D with shape [batch, in_channel, out_depth, out_height, out_width]
      """
  
      assert method in (
          "bilinear",
          "nearest",
      ), f"{method} is not supported"
      assert padding_mode in (
          "zeros",
          "border",
          "reflection",
      ), f"{padding_mode} is not supported"
      assert layout == "NCDHW", f"{layout} is not supported"
  
      batch, in_channel, in_depth, in_height, in_width = data.shape
      out_depth, out_height, out_width = grid.shape[2:]
  
      def _get_pixel_value(n, c, d, h, w):
          return te.if_then_else(
              te.all(
                  d >= 0,
                  h >= 0,
                  w >= 0,
                  d < in_depth,
                  h < in_height,
                  w < in_width,
              ),
              data[n, c, d, h, w],
              tir.const(0.0, dtype=data.dtype),
          )
  
      def _compute_source_index(n, d, h, w):
          z = grid[n, 2, d, h, w]
          y = grid[n, 1, d, h, w]
          x = grid[n, 0, d, h, w]
  
          if align_corners:
              z = (z + 1) * (in_depth - 1) / 2
              y = (y + 1) * (in_height - 1) / 2
              x = (x + 1) * (in_width - 1) / 2
          else:
              z = -0.5 + (z + 1) * in_depth / 2
              y = -0.5 + (y + 1) * in_height / 2
              x = -0.5 + (x + 1) * in_width / 2
  
          if padding_mode == "reflection":
              z = _reflect_coordinates(z, in_depth)
              y = _reflect_coordinates(y, in_height)
              x = _reflect_coordinates(x, in_width)
              z = _clip_coordinates(z, in_depth)
              y = _clip_coordinates(y, in_height)
              x = _clip_coordinates(x, in_width)
          elif padding_mode == "border":
              z = _clip_coordinates(z, in_depth)
              y = _clip_coordinates(y, in_height)
              x = _clip_coordinates(x, in_width)
  
          return (z, y, x)
  
      def _clip_coordinates(x, size):
          return te.min(te.max(x, 0), size - 1)
  
      def _reflect_coordinates(x, size):
          def __refelection(x, size, corner_start):
              def __reflect(index, size, corner_start):
                  index_align_corner = te.abs(corner_start - index)
                  size_times = te.truncdiv(
                      index_align_corner.astype("int32"), size
                  ).astype("int32")
                  t = tir.Mod(size_times, 2)
                  extra = index_align_corner - size_times * size
                  return tir.if_then_else(
                      tir.EQ(t, 0),
                      extra + corner_start,
                      size - extra + corner_start,
                  )
  
              return tir.if_then_else(
                  tir.all(x >= corner_start, x <= size + corner_start),
                  x,
                  __reflect(x, size, corner_start),
              )
  
          if align_corners:
              return __refelection(x, size - 1, 0)
          return __refelection(x, size, -0.5)
  
      def _trilinear_sample(n, c, d, h, w):
          z, y, x = _compute_source_index(n, d, h, w)
          z0 = te.floor(z).astype("int32")
          y0 = te.floor(y).astype("int32")
          x0 = te.floor(x).astype("int32")
          z1 = z0 + tir.const(1, "int32")
          y1 = y0 + tir.const(1, "int32")
          x1 = x0 + tir.const(1, "int32")
  
          return (
              _get_pixel_value(n, c, z0, y0, x0)
              * (1 - (x - x0))
              * (1 - (y - y0))
              * (1 - (z - z0))
              + _get_pixel_value(n, c, z0, y0, x1)
              * (x - x0)
              * (1 - (y - y0))
              * (1 - (z - z0))
              + _get_pixel_value(n, c, z1, y1, x0)
              * (1 - (x - x0))
              * (y - y0)
              * (z - z0)
              + _get_pixel_value(n, c, z1, y1, x1)
              * (x - x0)
              * (y - y0)
              * (z - z0)
              + _get_pixel_value(n, c, z0, y1, x0)
              * (1 - (x - x0))
              * (y - y0)
              * (1 - (z - z0))
              + _get_pixel_value(n, c, z1, y0, x1)
              * (x - x0)
              * (1 - (y - y0))
              * (z - z0)
              + _get_pixel_value(n, c, z1, y0, x0)
              * (1 - (x - x0))
              * (1 - (y - y0))
              * (z - z0)
              + _get_pixel_value(n, c, z0, y1, x1)
              * (x - x0)
              * (y - y0)
              * (1 - (z - z0))
          )
  
      def _nearest_sample(n, c, d, h, w):
          z, y, x = _compute_source_index(n, d, h, w)
          z_new = te.nearbyint(z).astype("int32")
          y_new = te.nearbyint(y).astype("int32")
          x_new = te.nearbyint(x).astype("int32")
          return _get_pixel_value(n, c, z_new, y_new, x_new)
  
      if method == "bilinear":
          interpolation = _trilinear_sample
      else:  # method == "nearest"
          interpolation = _nearest_sample
  
      return te.compute(
          (batch, in_channel, out_depth, out_height, out_width),
          interpolation,
          tag="grid_sample",
      )
  
imports: from tvm import te, tir
label_desc: Use image node_id 3 for calling a global function with example usage: _grid_sample_3d(data, grid, method, layout, padding_mode, align_corners) and returns: te, te, z,  y,  x, te, __refelection, unknown, _get_pixel_value, tir, __refelection, tir