code: def _gen_loop_start(
      chunk_input: List[Node],
      chunk_output: List[Node],
      chunk_output_dim: int,
      chunk_size=2,
  ) -> str:
      """
      Generate chunk loop start
  
      eg. chunk_result = torch.empty([100, 100], dtype=input_node.dtype, device=input_node.device)
          chunk_size = 32
          for chunk_idx in range(0, 100, 32):
              ......
  
      Args:
          chunk_input (List[Node]): chunk input node
          chunk_output (Node): chunk output node
          chunk_output_dim (int): chunk output node chunk dim
          chunk_size (int): chunk size. Defaults to 2.
  
      Returns:
          context (str): generated str
      """
      input_node = chunk_input[0]
  
      context = ""
      for i in range(len(chunk_output)):
          shape_str = str(list(get_node_shape(chunk_output[i])))
          if get_node_name(chunk_output[i]) in ["split", "unbind"]:
              tensor_str = (
                  "torch.empty(%s, dtype=%s.dtype, device=%s.device), "
                  % (
                      shape_str,
                      input_node.name,
                      input_node.name,
                  )
              )
              tensor_str = tensor_str * len(
                  chunk_output[i].meta["tensor_meta"]
              )
              tensor_str = "[" + tensor_str[:-2] + "]"
              context += "%s = %s;  " % (
                  chunk_output[i].name,
                  tensor_str,
              )
          else:
              context += (
                  "%s = torch.empty(%s, dtype=%s.dtype, device=%s.device);  "
                  % (
                      chunk_output[i].name,
                      shape_str,
                      input_node.name,
                      input_node.name,
                  )
              )
  
      out_shape = get_node_shape(chunk_output[0])
      chunk_shape = out_shape[chunk_output_dim[0]]
      context += (
          "chunk_size = %d\nfor chunk_idx in range(0, %d, chunk_size):\n"
          % (chunk_size, chunk_shape)
      )
      return context
  
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 2 for calling a global function with example usage: _gen_loop_start(chunk_input, chunk_output, chunk_output_dim, chunk_size) and returns: context