code: def helper_test(
      nm,
      gen,
      train,
      max_memory_allowed,
      max_kernels_allowed,
      all_jitted=False,
  ):
      tms = []
      for _ in range(4):
          early_gen = [
              x.realize() if isinstance(x, Tensor) else x for x in gen()
          ]
          GlobalCounters.reset()
          Device[Device.DEFAULT].synchronize()
          st = time.perf_counter_ns()
          train(*early_gen)
          Device[Device.DEFAULT].synchronize()
          tms.append(time.perf_counter_ns() - st)
      mem_used = GlobalCounters.mem_used - global_mem_used
  
      # TODO: jit should expose this correctly with graph
      kernels_used = (
          len(train.jit_cache) if hasattr(train, "jit_cache") else None
      )
      print(
          f"{nm}: used {mem_used/1e9:.2f} GB and {kernels_used} kernels in {min(tms)/1e6:.2f} ms"
      )
      assert (
          mem_used / 1e9 < max_memory_allowed
      ), f"{nm} used more than {max_memory_allowed:.2f} GB"
      assert (
          not kernels_used or kernels_used <= max_kernels_allowed
      ), f"{nm} used more than {max_kernels_allowed} kernels"
      if all_jitted:
          assert (
              kernels_used > 0
              and kernels_used == GlobalCounters.kernel_count
              or (
                  kernels_used == 1
                  and getattr(Device[Device.DEFAULT], "graph", None)
              )
          ), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted"
  
imports: import unittest, time, gc
  import numpy
  from tinygrad.tensor import Tensor
  from tinygrad.nn import optim
  from tinygrad.nn.state import get_parameters
  from tinygrad.jit import TinyJit
  from tinygrad import Device, GlobalCounters, dtypes
  from tinygrad.helpers import CI
  from tinygrad.shape.symbolic import Variable
  from extra.lr_scheduler import OneCycleLR
  from test.helpers import derandomize_model
  from examples.gpt2 import Transformer, MODEL_PARAMS
  from examples.hlb_cifar10 import SpeedyResNet, hyp
  from examples.llama import Transformer, MODEL_PARAMS
  from examples.stable_diffusion import UNetModel, ResBlock
label_desc: Use image node_id 1 for calling a global function with example usage: helper_test(nm, gen, train, max_memory_allowed, max_kernels_allowed, all_jitted) without return types