docstr: Search ckpt strategy with b = 0, then run the allocation algorithm again with b = √xy.
  Grid search over [√2/2 b, √2 b] for ``ckpt_opt`` over ``num_grids`` as in appendix A.
code: def grid_search(self) -> Set:
      """
      Search ckpt strategy with b = 0, then run the allocation algorithm again with b = √xy.
      Grid search over [√2/2 b, √2 b] for ``ckpt_opt`` over ``num_grids`` as in appendix A.
      """
      _, b_approx = self.run_chen_greedy(0)
      b_min, b_max = math.floor(b_approx / math.sqrt(2)), math.ceil(
          b_approx * math.sqrt(2)
      )
      b_opt = math.inf
      for b in range(b_min, b_max, (b_max - b_min) // self.num_grids):
          ckpt_intv, b_approx = self.run_chen_greedy(b)
          if b_approx < b_opt:
              b_opt = b_approx
              ckpt_opt = ckpt_intv
      return ckpt_opt
  
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 4 for calling the CheckpointSolverChen obj's underlying member method code with example usage: obj.grid_search() and returns: ckpt_opt