code: def emit_code_with_nested_activation_checkpoint(
      body, ckpt_func, nodes, emit_node_func, delete_unused_value_func
  ):
      """Emit code with nested activation checkpoint
      When we detect some of the node.activation_checkpoint is a List, we will use
      this function to emit the activation checkpoint codes.
      Args:
          body: forward code
          ckpt_func: checkpoint functions code
          nodes: graph.nodes
          emit_node_func: function to emit node
          delete_unused_value_func: function to remove the unused value
      """
      ckpt_regions = _find_nested_ckpt_regions(nodes, 0)
      start_idx = [item[0] for item in ckpt_regions]
      end_idx = [item[1] for item in ckpt_regions]
  
      # find the offload regions
      offload_regions, offload_labels = _find_offload_regions(nodes)
      offload_starts = [item[0] for item in offload_regions]
      offload_ends = [item[1] for item in offload_regions]
      offload_inputs = []
      offload_outputs = []
      within_offload_region = False
  
      node_list = list(nodes)
  
      # find the input and output var names for each offload region
      for idx, (start, end) in enumerate(offload_regions):
          offload_node_list = node_list[start : end + 1]
          inputs, outputs = _find_input_and_output_nodes(
              offload_node_list
          )
          offload_inputs.append(inputs)
          offload_outputs.append(outputs)
  
      # this flag is to prevent repeated insert of save tensors
      # hooks definition in ckpt_func
      is_hook_inserted = False
      node_idx = 0
      while 1:
          # break if we finish the processing all the nodes
          if node_idx >= len(node_list):
              break
  
          # process ckpt_regions
          if node_idx in start_idx:
              ckpt_node_list = node_list[
                  node_idx : end_idx[start_idx.index(node_idx)] + 1
              ]
              emit_ckpt_func(
                  body,
                  ckpt_func,
                  ckpt_node_list,
                  emit_node_func,
                  delete_unused_value_func,
              )
              node_idx += len(ckpt_node_list)
  
          # process node in forward function
          else:
              node = node_list[node_idx]
  
              if node_idx in offload_starts:
                  offload_label = offload_labels[
                      offload_starts.index(node_idx)
                  ]
                  _, offload_input, offload_bar = offload_label
                  within_offload_region = True
  
                  # insert hook functions if needed
                  if not is_hook_inserted:
                      (
                          pack_hook,
                          unpack_hook,
                      ) = _gen_saved_tensors_hooks()
                      ckpt_func.insert(
                          0, "\n".join([pack_hook, unpack_hook]) + "\n"
                      )
                      is_hook_inserted = True
  
                  if offload_input and offload_bar:
                      body.append(_gen_save_on_cpu_context())
  
                  elif offload_input:
                      for par in offload_inputs[offload_label[0]]:
                          body.append(
                              f"setattr({par}, 'offload', True)\n"
                          )
                      body.append(
                          _gen_save_tensors_hooks_context(
                              offload_input=True
                          )
                      )
  
                  else:
                      for par in offload_inputs[offload_label[0]]:
                          body.append(
                              f"setattr({par}, 'offload', False)\n"
                          )
                      body.append(
                          _gen_save_tensors_hooks_context(
                              offload_input=False
                          )
                      )
  
              if within_offload_region:
                  emit_node_func(node, body)
                  body[-1] = "    " + body[-1]
                  delete_unused_value_func(node, body)
  
              else:
                  emit_node_func(node, body)
                  delete_unused_value_func(node, body)
  
              if node_idx in offload_ends:
                  within_offload_region = False
  
              node_idx += 1
  
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 13 for calling a global function with example usage: emit_code_with_nested_activation_checkpoint(body, ckpt_func, nodes, emit_node_func, delete_unused_value_func) without return types