docstr: This is the simple implementation of Algorithm 3 in https://arxiv.org/abs/1604.06174.
  Note that this algorithm targets at memory optimization only, using techniques in appendix A.
  
  Usage:
      Assume that we have a ``GraphModule``, and we have already done the extractions
      to the graph to retrieve all information needed, then we could use the following
      code to find a solution using ``CheckpointSolverChen``:
      >>> solver = CheckpointSolverChen(gm.graph)
      >>> chen_graph = solver.solve()
      >>> gm.graph = chen_graph    # set the graph to a new graph
  
  Args:
      graph (Graph): The computing graph to be optimized.
      cnode (List[str], optional): Common node List, should be the subset of input. Defaults to None.
      num_grids (int, optional): Number of grids to search for b. Defaults to 6.
code: def __init__(
      self, graph: Graph, cnode: List[str] = None, num_grids: int = 6
  ):
      """
      This is the simple implementation of Algorithm 3 in https://arxiv.org/abs/1604.06174.
      Note that this algorithm targets at memory optimization only, using techniques in appendix A.
  
      Usage:
          Assume that we have a ``GraphModule``, and we have already done the extractions
          to the graph to retrieve all information needed, then we could use the following
          code to find a solution using ``CheckpointSolverChen``:
          >>> solver = CheckpointSolverChen(gm.graph)
          >>> chen_graph = solver.solve()
          >>> gm.graph = chen_graph    # set the graph to a new graph
  
      Args:
          graph (Graph): The computing graph to be optimized.
          cnode (List[str], optional): Common node List, should be the subset of input. Defaults to None.
          num_grids (int, optional): Number of grids to search for b. Defaults to 6.
      """
      super().__init__(graph, 0, 0, True, cnode)
      self.num_grids = num_grids
  
file_path: colossalai/colossalai/auto_parallel/checkpoint/ckpt_solver_chen.py
imports: import math
  from copy import deepcopy
  from typing import List, Set, Tuple
  from torch.fx import Graph, Node
  from colossalai.fx.profiler import calculate_fwd_in, calculate_fwd_tmp
  from .ckpt_solver_base import CheckpointSolverBase
label_desc: Use image node_id 1 to create a new CheckpointSolverChen object from inherited base classes: CheckpointSolverBase with example: obj = CheckpointSolverChen(graph, cnode, num_grids)