code: def match_computation(
      stats_baseline, key=["gnn", "dim_inner"], mode="sqrt"
  ):
      """Match computation budget by modifying :obj:`cfg.gnn.dim_inner`."""
      stats = get_stats()
      if stats != stats_baseline:
          # Phase 1: fast approximation
          while True:
              if mode == "sqrt":
                  scale = math.sqrt(stats_baseline / stats)
              elif mode == "linear":
                  scale = stats_baseline / stats
              step = (
                  int(round(cfg[key[0]][key[1]] * scale))
                  - cfg[key[0]][key[1]]
              )
              cfg[key[0]][key[1]] += step
              stats = get_stats()
              if abs(step) <= 1:
                  break
          # Phase 2: fine tune
          flag_init = 1 if stats < stats_baseline else -1
          step = 1
          while True:
              cfg[key[0]][key[1]] += flag_init * step
              stats = get_stats()
              flag = 1 if stats < stats_baseline else -1
              if stats == stats_baseline:
                  return stats
              if flag != flag_init:
                  if not cfg.model.match_upper:  # stats is SMALLER
                      if flag < 0:
                          cfg[key[0]][key[1]] -= flag_init * step
                      return get_stats()
                  else:
                      if flag > 0:
                          cfg[key[0]][key[1]] -= flag_init * step
                      return get_stats()
      return stats
  
imports: import math
  from torch_geometric.graphgym.config import cfg, set_cfg
  from torch_geometric.graphgym.model_builder import create_model
label_desc: Use image node_id 3 for calling a global function with example usage: match_computation(stats_baseline, key, mode) and returns: stats, stats, get_stats, get_stats