code: def emit_ckpt_func(
      body,
      ckpt_func,
      node_list: List[Node],
      emit_node_func,
      delete_unused_value_func,
      level=0,
      in_ckpt=False,
  ):
      """Emit ckpt function in nested way
      Args:
          body: forward code, in recursive calls, this part will be checkpoint
          functions code
          ckpt_func: checkpoint functions code, in recursive calls, this part
          will be a buffer
          node_list (List[Node]): list of torch.fx.Node
          emit_node_func: function to emit a node
          delete_unused_value_func: function to delete unused value
          level (int, optional): checkpoint level. Defaults to 0.
          in_ckpt (bool, optional): indicates wether the func is in recursive
          call. Defaults to False.
      """
      inputs, outputs = _find_input_and_output_nodes(node_list)
  
      # if the current checkpoint function use int as label, using old generation method
      if isinstance(node_list[0].meta["activation_checkpoint"], int):
          label = node_list[0].meta["activation_checkpoint"]
          ckpt_fn_def = _gen_ckpt_fn_def(label, inputs)
          ckpt_func.append(f"{ckpt_fn_def}\n")
          for node in node_list:
              emit_node_func(node, ckpt_func)
              ckpt_func[-1] = "    " + ckpt_func[-1]
              delete_unused_value_func(node, ckpt_func)
  
          ckpt_func.append("    " + _gen_ckpt_output(outputs) + "\n\n")
          activation_offload = node_list[0].meta.get(
              "activation_offload", False
          )
          usage = _gen_ckpt_usage(
              label, activation_offload, inputs, outputs, False
          )
          usage += "\n"
          body.append(usage)
  
      # use nested ckpt function codegen
      else:
          # label given by each layer, e.g. if you are currently at level [0, 1, 1]
          # the label will be '0_1_1'
          label = "_".join(
              [
                  str(idx)
                  for idx in node_list[0].meta["activation_checkpoint"][
                      : level + 1
                  ]
              ]
          )
          ckpt_fn_def = _gen_ckpt_fn_def(label, inputs)
          ckpt_func.append(f"{ckpt_fn_def}\n")
  
          # if there is more level to fetch
          if level + 1 < len(
              node_list[0].meta["activation_checkpoint"]
          ):
              ckpt_regions = _find_nested_ckpt_regions(
                  node_list, level + 1
              )
              start_idx = [item[0] for item in ckpt_regions]
              end_idx = [item[1] for item in ckpt_regions]
  
              # use ckpt_func_buffer to store nested checkpoint functions
              ckpt_func_buffer = []
              node_idx = 0
              while 1:
                  if node_idx >= len(node_list):
                      break
  
                  if node_idx in start_idx:
                      ckpt_node_list = node_list[
                          node_idx : end_idx[start_idx.index(node_idx)]
                          + 1
                      ]
                      emit_ckpt_func(
                          ckpt_func,
                          ckpt_func_buffer,
                          ckpt_node_list,
                          emit_node_func,
                          delete_unused_value_func,
                          level + 1,
                          True,
                      )
                      node_idx += len(ckpt_node_list)
  
                  else:
                      node = node_list[node_idx]
                      emit_node_func(node, ckpt_func)
                      ckpt_func[-1] = "    " + ckpt_func[-1]
                      delete_unused_value_func(node, ckpt_func)
                      node_idx += 1
  
              ckpt_func.append(
                  "    " + _gen_ckpt_output(outputs) + "\n\n"
              )
              ckpt_func += ckpt_func_buffer
              activation_offload = node_list[0].meta.get(
                  "activation_offload", False
              )
              usage = (
                  _gen_ckpt_usage(
                      label, activation_offload, inputs, outputs, False
                  )
                  + "\n"
              )
              if in_ckpt:
                  usage = "    " + usage
              body.append(usage)
  
          # last level
          else:
              for node in node_list:
                  emit_node_func(node, ckpt_func)
                  ckpt_func[-1] = "    " + ckpt_func[-1]
                  delete_unused_value_func(node, ckpt_func)
  
              ckpt_func.append(
                  "    " + _gen_ckpt_output(outputs) + "\n\n"
              )
              activation_offload = node_list[0].meta.get(
                  "activation_offload", False
              )
              usage = (
                  _gen_ckpt_usage(
                      label, activation_offload, inputs, outputs, False
                  )
                  + "\n"
              )
              if in_ckpt:
                  usage = "    " + usage
              body.append(usage)
  
file_path: colossalai/colossalai/fx/codegen/activation_checkpoint_codegen.py
imports: from typing import Any, Dict, Iterable, List, Tuple
  import torch
  import colossalai
label_desc: Use image node_id 12 for calling a global function with example usage: emit_ckpt_func(body, ckpt_func, node_list, emit_node_func, delete_unused_value_func, level, in_ckpt) without return types