code: def beam_search(
      lin: Linearizer, rawbufs, amt: int, allow_test_size=True
  ) -> Linearizer:
      key = {
          "ast": str(lin.ast),
          "amt": amt,
          "allow_test_size": allow_test_size,
          "device": lin.opts.device,
      }
      if (
          (val := diskcache_get("beam_search", key)) is not None
          and not getenv("IGNORE_BEAM_CACHE")
          and CACHELEVEL >= 1
      ):
          ret = lin.copy()
          for o in val[len(lin.applied_opts) :]:
              ret.apply_opt(o)
          return ret
  
      beam: List[Tuple[Linearizer, float]] = []
      seen_libs = set()
  
      default_parallel = 1 if lin.opts.device in {"CUDA", "HIP"} else 0
      pool = (
          multiprocessing.Pool(
              multiprocessing.cpu_count(), _init_worker
          )
          if getenv("PARALLEL", default_parallel)
          else None
      )
  
      try:
          var_vals = {k: (k.max + k.min) // 2 for k in lin.ast.vars()}
          exiting, st = False, time.perf_counter()
          dev = Device[lin.opts.device]
          assert isinstance(dev, Compiled)
          while not exiting:
              acted_lins = (
                  flatten(
                      [
                          get_linearizer_actions(
                              lin, include_0=False
                          ).values()
                          for lin, _ in beam
                      ]
                  )
                  if len(beam)
                  else [lin]
              )
              timed_lins: List[Tuple[Linearizer, float]] = []
              _compile_fn = functools.partial(
                  _try_compile_linearized_w_idx, device=lin.opts.device
              )
              for i, proc in (
                  pool.imap_unordered(
                      _compile_fn, enumerate(acted_lins)
                  )
                  if pool is not None
                  else map(_compile_fn, enumerate(acted_lins))
              ):
                  if proc is None:
                      continue
                  lib, global_size, local_size = proc
                  if lib in seen_libs:
                      continue
                  seen_libs.add(lib)
                  tms = _time_program(
                      lin.ast,
                      dev,
                      lib,
                      global_size,
                      local_size,
                      var_vals,
                      rawbufs,
                      early_stop=beam[0][1] * 3 if len(beam) else 1.0,
                  )
                  timed_lins.append((acted_lins[i], min(tms)))
                  if DEBUG >= 2:
                      print(
                          f"\r{time.perf_counter() - st:7.2f}s: {timed_lins[-1][1]*1e6:12.2f} us       {len(timed_lins):4d}/{len(acted_lins):4d}         {timed_lins[-1][0].colored_shape()}\033[K",
                          end="",
                      )  # noqa: E501
  
              # done
              opts = sorted(timed_lins, key=lambda x: x[1])
              exiting = len(opts) == 0 or (
                  len(beam) > 0 and beam[0][1] <= opts[0][1]
              )
              if not exiting:
                  beam = opts[:amt]
              assert len(beam) > 0, "no BEAM items succeeded?!?"
              if DEBUG >= 2:
                  print(
                      f"\r{time.perf_counter() - st:7.2f}s:",
                      colored(
                          f"{beam[0][1]*1e6:12.2f} us",
                          "green" if exiting else None,
                      ),
                      f"from {len(acted_lins):3d} -> {len(opts):3d} actions\033[K",
                      beam[0][0].colored_shape(),
                  )  # noqa: E501
          if pool is not None:
              pool.close()  # the pool is closed
      except KeyboardInterrupt as e:
          if pool is not None:
              pool.terminate()
          raise e
  
      if CACHELEVEL >= 1:
          diskcache_put("beam_search", key, beam[0][0].applied_opts)
      if DEBUG >= 3:
          print(beam[0][0].applied_opts)
      return beam[0][0]
  
imports: from typing import Dict, List, cast, DefaultDict, Optional, Tuple, Callable
  import itertools, functools, random, math, time, multiprocessing, traceback, signal
  from tinygrad.device import Device, Compiled, Buffer, CompiledASTRunner
  from tinygrad.ops import MemBuffer, LazyOp
  from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, to_function_name
  from tinygrad.dtype import ImageDType
  from tinygrad.codegen.linearizer import Linearizer
  from collections import defaultdict
  from tinygrad.tensor import Tensor
  from tinygrad.shape.symbolic import sym_infer
  from tinygrad.codegen.kernel import Opt, OptOps
label_desc: Use image node_id 8 for calling a global function with example usage: beam_search(lin, rawbufs, amt, allow_test_size) and returns: beam, ret