code: def emit_code_with_chunk(
      body: List[str],
      nodes: Iterable[Node],
      emit_node_func: Callable,
      delete_unused_value_func: Callable,
      search_chunk: SearchChunk,
      chunk_infos: List,
      eval_mem: bool = False,
  ):
      """
      Emit code with chunk according to chunk_infos.
  
      It will generate a for loop in chunk regions, and
      replace inputs and outputs of regions with chunked variables.
  
      Args:
          body: forward code
          nodes: graph.nodes
          emit_node_func: function to emit node
          delete_unused_value_func: function to remove the unused value
          search_chunk: the class to search all chunks
          chunk_infos: store all information about all chunks.
      """
      node_list = list(nodes)
  
      # chunk region
      chunk_starts = [i["region"][0] for i in chunk_infos]
      chunk_ends = [i["region"][1] for i in chunk_infos]
  
      # chunk inputs
      chunk_inputs = [
          i["inputs"] for i in chunk_infos
      ]  # input with chunk
      chunk_inputs_non_chunk = [
          i["inputs_non_chunk"] for i in chunk_infos
      ]  # input without chunk
      chunk_inputs_dim = [
          i["inputs_dim"] for i in chunk_infos
      ]  # input chunk dim
      chunk_inputs_names = [j.name for i in chunk_inputs for j in i] + [
          j.name for i in chunk_inputs_non_chunk for j in i
      ]
  
      # chunk outputs
      chunk_outputs = [i["outputs"] for i in chunk_infos]
      chunk_outputs_non_tensor = [
          i["outputs_non_tensor"] for i in chunk_infos
      ]
      chunk_outputs_dim = [i["outputs_dim"] for i in chunk_infos]
  
      node_list = search_chunk.reorder_graph.reorder_node_list(
          node_list
      )
      node_idx = 0
      region_idx = 0
      within_chunk_region = False
  
      if eval_mem:
          body.append(
              "init_memory = torch.cuda.memory_allocated() / 1024**2\n"
          )
  
      while node_idx < len(node_list):
          node = node_list[node_idx]
  
          # if is chunk start, generate for loop start
          if node_idx in chunk_starts:
              within_chunk_region = True
              region_idx = chunk_starts.index(node_idx)
              body.append(
                  _gen_loop_start(
                      chunk_inputs[region_idx],
                      chunk_outputs[region_idx],
                      chunk_outputs_dim[region_idx],
                      chunk_infos[region_idx]["chunk_size"],
                  )
              )
  
          if within_chunk_region:
              emit_node_func(node, body)
              # replace input var with chunk var
              body = _add_node_slice(
                  chunk_inputs,
                  region_idx,
                  chunk_inputs_dim,
                  node_idx,
                  body,
                  node,
              )
              # replace output var with chunk var
              body = _add_node_slice(
                  chunk_outputs,
                  region_idx,
                  chunk_outputs_dim,
                  node_idx,
                  body,
                  node,
              )
              # new tensor like
              body = _replace_new_tensor_like_shape(
                  search_chunk,
                  chunk_infos,
                  region_idx,
                  node_idx,
                  node,
                  body,
              )
              # new tensor
              body = _replace_new_tensor_shape(
                  search_chunk,
                  chunk_infos,
                  region_idx,
                  node_idx,
                  node,
                  body,
              )
              # reassign reshape size
              body[-1] = _replace_reshape_size(
                  body[-1],
                  node.name,
                  chunk_infos[region_idx]["reshape_size"],
              )
              body[-1] = "    " + body[-1]
              delete_unused_value_func(node, body, chunk_inputs_names)
              if eval_mem:
                  body.append(
                      "    if chunk_idx == 0:\n        print('%s', torch.cuda.max_memory_allocated() / 1024**2 - init_memory);  torch.cuda.reset_peak_memory_stats()\n"
                      % (node.name)
                  )
          else:
              emit_node_func(node, body)
              if node_idx not in chunk_inputs:
                  delete_unused_value_func(
                      node, body, chunk_inputs_names
                  )
              if eval_mem:
                  body.append(
                      "print('%s', torch.cuda.max_memory_allocated() / 1024**2 - init_memory);  torch.cuda.reset_peak_memory_stats()\n"
                      % (node.name)
                  )
  
          # generate chunk region end
          if node_idx in chunk_ends:
              body.append(
                  _gen_loop_end(
                      chunk_inputs[region_idx],
                      chunk_inputs_non_chunk[region_idx],
                      node_list,
                      chunk_ends[region_idx],
                      chunk_outputs_non_tensor[region_idx],
                      search_chunk,
                  )
              )
              within_chunk_region = False
  
          node_idx += 1
  
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 9 for calling a global function with example usage: emit_code_with_chunk(body, nodes, emit_node_func, delete_unused_value_func, search_chunk, chunk_infos, eval_mem) without return types