code: def _add_node_slice(
      chunk_nodes: List[Node],
      region_idx: int,
      chunk_nodes_dim: Dict,
      node_idx: int,
      body: List[str],
      node: Node,
  ) -> List[str]:
      """
      add chunk slice for input nodes
      """
      for chunk_node_idx, chunk_node in enumerate(
          chunk_nodes[region_idx]
      ):
          # inputs node
          if isinstance(
              chunk_nodes_dim[region_idx][chunk_node_idx], dict
          ):
              for idx, dim in chunk_nodes_dim[region_idx][
                  chunk_node_idx
              ].items():
                  if idx == node_idx:
                      chunk_slice = _gen_chunk_slice_dim(
                          dim[0],
                          "chunk_idx",
                          get_node_shape(chunk_node),
                      )
                      body[-1] = _replace_name(
                          body[-1],
                          chunk_node.name,
                          chunk_node.name + chunk_slice,
                      )
          # outputs node
          else:
              if chunk_node.name == node.name or (
                  chunk_node.name
                  in [i.name for i in node.all_input_nodes]
              ):
                  chunk_slice = _gen_chunk_slice_dim(
                      chunk_nodes_dim[region_idx][chunk_node_idx],
                      "chunk_idx",
                      get_node_shape(chunk_node),
                  )
                  if get_node_name(chunk_node) in ["split", "unbind"]:
                      split_chunk_slice = ""
                      for i in range(
                          len(chunk_node.meta["tensor_meta"])
                      ):
                          split_chunk_slice += "%s[%d]%s, " % (
                              chunk_node.name,
                              i,
                              chunk_slice,
                          )
                      split_chunk_slice = split_chunk_slice[:-2]
                      body[-1] = _replace_name(
                          body[-1], chunk_node.name, split_chunk_slice
                      )
                  else:
                      body[-1] = _replace_name(
                          body[-1],
                          chunk_node.name,
                          chunk_node.name + chunk_slice,
                      )
      return body
  
file_path: ColossalAI/colossalai/autochunk/autochunk_codegen.py
imports: from typing import Any, Callable, Dict, Iterable, List, Tuple
  import torch
  import colossalai
  from colossalai.fx._compatibility import is_compatible_with_meta
  from colossalai.fx.codegen.activation_checkpoint_codegen import CODEGEN_AVAILABLE
  from torch.fx.node import Argument, Node, _get_qualified_name, _type_repr, map_arg
  from .search_chunk import SearchChunk
  from .utils import delete_free_var_from_last_use, get_logger, get_node_name, get_node_shape
label_desc: Use image node_id 8 for calling a global function with example usage: _add_node_slice(chunk_nodes, region_idx, chunk_nodes_dim, node_idx, body, node) and returns: body