IQ.Pilot Release Commit @ 0798119

tinygrad_repo/test/external/external_benchmark_am.py

@@ -0,0 +1,9 @@
+from tinygrad.helpers import Profiling
+from tinygrad import Device
+
+if __name__ == "__main__":
+  am = Device["AMD"]
+
+  # kfd is 0.55ms!
+  with Profiling("allocation 127.7mb"):
+    am.allocator.alloc(int(127.7*1024*1024))

tinygrad_repo/test/external/external_benchmark_bert_matmuls.py

@@ -0,0 +1,18 @@
+from tinygrad import Tensor, dtypes
+dtypes.default_float = dtypes.float16
+from tinygrad.dtype import to_dtype
+from tinygrad.helpers import getenv
+
+if __name__ == "__main__":
+  # matmuls in bert layers
+  BS = getenv("BS", 96//6)
+  acc_dtype = to_dtype(getenv("ACC_DTYPE", "half"))
+  tensors = [
+    (Tensor.empty(BS, 512, 1024), Tensor.empty(1024, 1024).T),                                          # linear to get qkv
+    (Tensor.empty(BS, 512, 16, 64).permute(0,2,1,3), Tensor.empty(BS, 512, 16, 64).permute(0,2,3,1)),   # q@k
+    (Tensor.empty(BS, 16, 512, 512), Tensor.empty(BS, 512, 16, 64).permute(0,2,1,3)),                   # qk@v
+  ]
+  for t0, t1 in tensors:
+    print(f"{t0.shape=}, {t0.uop.st.is_expanded()=}, {t1.shape=}, {t1.uop.st.is_expanded()=}")
+    for _ in range(5):
+      t0.dot(t1, dtype=acc_dtype).realize()

tinygrad_repo/test/external/external_benchmark_bert_softmax.py

@@ -0,0 +1,17 @@
+from tinygrad import Tensor, dtypes, GlobalCounters
+dtypes.default_float = dtypes.float16
+from tinygrad.dtype import to_dtype
+from tinygrad.helpers import getenv
+from test.backend.test_softmax_fusion import single_kernel_softmax
+
+if __name__ == "__main__":
+  # softmax in bert layers
+  BS = getenv("BS", 96//6)
+  acc_dtype = to_dtype(getenv("ACC_DTYPE", "half"))
+  t = Tensor.empty(BS, 16, 512, 512)
+  t.softmax(-1, dtype="half").realize()
+
+  # test single kernel softmax
+  GlobalCounters.reset()
+  single_kernel_softmax(t, -1, acc_dtype).realize()
+

tinygrad_repo/test/external/external_benchmark_disk_raw.py

@@ -0,0 +1,8 @@
+import pathlib
+from tinygrad import Tensor, Device, Context
+from tinygrad.helpers import getenv
+
+if __name__ == "__main__":
+  with Context(DEBUG=2):
+    disk_llama = Tensor(pathlib.Path(getenv("TESTFILE", "/raid/weights/LLaMA-3/8B/consolidated.00.pth")))
+    device_llama = disk_llama.to(Device.DEFAULT).realize()

tinygrad_repo/test/external/external_benchmark_hip_compile.py

@@ -0,0 +1,31 @@
+import random, os
+from tinygrad.helpers import Timing
+from tinygrad.runtime.ops_hip import compile_hip, HIPDevice
+from tinygrad.runtime.ops_cl import compile_cl, CLDevice
+
+# OMP_NUM_THREADS=1 strace -tt -f -e trace=file python3 test/external/external_benchmark_hip_compile.py
+# AMD_COMGR_REDIRECT_LOGS=stdout AMD_COMGR_EMIT_VERBOSE_LOGS=1 python3 test/external/external_benchmark_hip_compile.py
+
+# issue is in https://github.com/ROCm-Developer-Tools/clr/
+
+if __name__ == "__main__":
+  HIPDevice()
+  CLDevice()
+
+  # warmup
+  name = "none"+str(random.randint(0, 1000000))
+  compile_cl.__wrapped__(f"void {name}() {{}}")
+  print("compile cl warmed up")
+  compile_hip.__wrapped__(f"void {name}() {{}}")
+  print("compile hip warmed up")
+
+  print("**** benchmark ****")
+  name = "none"+str(random.randint(0, 1000000))
+  # this uses AMD_COMGR_ACTION_COMPILE_SOURCE_TO_BC, then it links the lib on the next step
+  with Timing("compile cl:  "): compile_cl.__wrapped__(f"void {name}() {{}}")
+  # this uses AMD_COMGR_ACTION_COMPILE_SOURCE_WITH_DEVICE_LIBS_TO_BC, much slower
+  with Timing("compile hip: "): compile_hip.__wrapped__(f"void {name}() {{}}")
+  os._exit(0)
+
+
+

tinygrad_repo/test/external/external_benchmark_keccak.py

@@ -0,0 +1,20 @@
+from tinygrad import Tensor, dtypes
+from tinygrad.engine.jit import TinyJit
+from tinygrad.helpers import Timing, getenv
+
+if __name__ == "__main__":
+  BS = getenv("BS", 2**14)
+  BLOCKSIZE = getenv("BLOCKSIZE", 4096)
+  HASHFN = getenv("HASHFN", "shake_128")
+  NRUNS = getenv("NRUNS", 5)
+
+  @TinyJit
+  def hasher(data: Tensor): return data.keccak(HASHFN)
+
+  t = Tensor.randn(BS, BLOCKSIZE, dtype=dtypes.uint8).realize()
+  ds_mib = t.nbytes() / 1024**2
+
+  print(f"--- benchmarking (hash: {HASHFN}, data size: {ds_mib} MiB, block size: {BLOCKSIZE} B, batch size: {BS})")
+  for i in range(NRUNS):
+    with Timing(f"run: {i+1}, elapsed time: ", (lambda et: f", throughput: {ds_mib / (et*1e-9):.2f} MiB/s")):
+      hasher(t).realize()

tinygrad_repo/test/external/external_benchmark_kernel_launch.py

@@ -0,0 +1,38 @@
+import time
+from tinygrad import Tensor, TinyJit, Device, Context
+from tinygrad.helpers import Profiling, Timing, GlobalCounters
+
+# python3 test/speed/external_test_speed_v_torch.py TestSpeed.test_add_a
+
+@TinyJit
+def plus(a:Tensor, b:Tensor): return a+b
+
+if __name__ == "__main__":
+  a = Tensor([1]).realize()
+  b = Tensor([1]).realize()
+  for i in range(5):
+    with Timing(prefix=f"{i}:"):
+      c = plus(a,b)
+      Device[c.device].synchronize()
+  assert c.item() == 2
+  for i in range(5):
+    st = time.perf_counter()
+    c = plus(a,b)
+    et = time.perf_counter() - st
+    Device[c.device].synchronize()
+    print(f"nosync  {i}: {et*1e6:.2f} us")
+  for i in range(5):
+    st = time.perf_counter()
+    c = plus(a,b)
+    Device[c.device].synchronize()
+    et = time.perf_counter() - st
+    print(f"precise {i}: {et*1e6:.2f} us")
+  assert GlobalCounters.time_sum_s == 0
+  with Context(DEBUG=2):
+    st = time.perf_counter()
+    c = plus(a,b)
+    Device[c.device].synchronize()
+    et = time.perf_counter() - st
+  print(f"kernel {GlobalCounters.time_sum_s*1e3:.2f} ms / full {et*1e3:.2f} ms -- {et/(GlobalCounters.time_sum_s+1e-12):.2f} x")
+  with Profiling():
+    c = plus(a,b)

tinygrad_repo/test/external/external_benchmark_llama_schedule.py

@@ -0,0 +1,35 @@
+from tinygrad import nn, Tensor, dtypes
+from tinygrad.helpers import DEV, Timing
+
+from extra.models.llama import Transformer
+from examples.llama3 import MODEL_PARAMS
+
+if __name__ == "__main__":
+  DEV.value = "NULL"
+  Tensor.training = True
+  #model_size = "8B"
+  model_size = "405B"
+
+  with Timing("total "):
+    with Timing("***** create model in    "):
+      model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=nn.Linear, embedding=nn.Embedding,
+                          max_context=1024, jit=True, disable_kv_cache=True)
+
+    with Timing("***** fake state in      "):
+      Tensor.realize(*[p.assign(Tensor.empty(*p.shape, device=p.device, dtype=p.dtype)) for p in nn.state.get_parameters(model)])
+
+    with Timing("***** create optim in    "):
+      opt = nn.optim.AdamW(nn.state.get_parameters(model))
+
+    with Timing("***** run model in       "):
+      toks = Tensor.empty(1, 1024, dtype=dtypes.int)
+      out = model(toks, 0, temperature=float('nan'))
+
+    with Timing("***** backward in        "):
+      out.mean().backward()
+
+    with Timing("***** realize in        "):
+      out.realize()
+
+    with Timing("***** step in           "):
+      opt.step()

tinygrad_repo/test/external/external_benchmark_multitensor_allreduce.py

@@ -0,0 +1,49 @@
+from tinygrad import Tensor, Device, GlobalCounters, TinyJit, dtypes
+from tinygrad.helpers import getenv, Context, DEBUG
+
+def test(devs: list[str], N: int, iters:int = 10, name:str = "allreduce"):
+  @TinyJit
+  def f(t: Tensor) -> Tensor: t.sum(0).realize()
+
+  secs, gflops, gbs = 0, 0, 0
+  for i in range(-3, iters):
+    t = Tensor.empty((len(devs), N))
+    t = t.shard(devs, 0).realize()
+    GlobalCounters.reset()
+    f(t)
+    for d in devs: Device[d].synchronize()
+
+    if i < 0: continue # warm up jit
+    i_secs = GlobalCounters.time_sum_s
+    i_gflops = GlobalCounters.global_ops/i_secs/10**9
+    i_gbs = (N*4)/i_secs/10**9
+    print(f"{name} iter {i+1}/{iters}: {i_secs:.6f} sec {i_gflops:.2f} GFLOP/s {i_gbs:.2f} GB/s")
+    secs += i_secs
+    gflops += i_gflops
+    gbs += i_gbs
+
+  return (gflops/iters, gbs/iters, secs/iters)
+
+def run(sz, n_gpus=6, iters=10, ring=0, all2all=0):
+  devs = tuple([f"{Device.DEFAULT}:{x}" for x in range(n_gpus)])
+  N = sz // dtypes.float32.itemsize
+  name = "all2all" if all2all else ("ring" if ring else "naive")
+  with Context(RING=(2 if ring else 0), ALL2ALL=(2 if all2all else 0), JIT_BATCH_SIZE=0, DEBUG=max(DEBUG.value, 2)):
+    return test(devs, N, iters=iters, name=name)
+
+def main():
+  n_gpus = getenv("GPUS", 6)
+  iters = getenv("ITERS", 10)
+  sz = getenv("SZ", 1000) * 10**6 # size of data on each gpu
+  print(f"Using {sz/10**9:.2f} GB of numbers on each of {n_gpus} GPUs, {n_gpus*sz/10**9:.2f} GB total.")
+
+  results = {}
+  for name, kwargs in [("naive", {}), ("ring", {"ring": 2}), ("all2all", {"all2all": 2})]:
+    results[name] = run(sz, n_gpus=n_gpus, iters=iters, **kwargs)
+
+  print("\n=== RESULTS ===")
+  for name, (gflops, gbs, secs) in results.items():
+    print(f"{name.upper()}:\n  {secs:.6f} seconds/iter\n  {gflops:.2f} GFLOP/s\n  {gbs:.2f} GB/s")
+
+if __name__ == "__main__":
+  main()

tinygrad_repo/test/external/external_benchmark_op_conv.py

@@ -0,0 +1,102 @@
+# ruff: noqa: E501 E712 F401
+from dataclasses import replace
+from tinygrad import dtypes, Device
+from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
+from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import
+from tinygrad.engine.realize import get_runtime
+from tinygrad.codegen import to_program
+from tinygrad.helpers import dedup, getenv
+from tinygrad.device import Buffer
+from tinygrad.dtype import ImageDType, Invalid
+
+# PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+
+def vision_conv_143():
+  c0 = UOp.param(0, dtypes.imageh((16, 1024, 4)))
+  c2 = UOp.range(32, 3, AxisType.LOOP)
+  c5 = UOp.range(128, 4, AxisType.LOOP)
+  c8 = UOp.range(16, 2, AxisType.LOOP)
+  c16 = UOp.range(7, 0, AxisType.REDUCE)
+  c17 = c8*2+c16
+  c24 = ((c17<3)!=True)&(c17<35)
+  c26 = UOp.range(7, 1, AxisType.REDUCE)
+  c27 = c2*2+c26
+  c32 = ((c27<3)!=True)&(c27<67)
+  c34 = UOp.param(1, dtypes.imageh((32, 1024, 4)))
+  c38 = c5//2
+  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.weakint, Invalid))
+  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
+  c49 = UOp.param(2, dtypes.imageh((64, 49, 4)))
+  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
+  c63 = UOp.param(3, dtypes.float.ptr(128))
+  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
+  c67 = c0.index((c2*128+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
+
+  opts = None
+  # JITBEAM=2
+  # (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.SWAP, axis=1, arg=2))
+  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
+
+def vision_conv_153():
+  c0 = UOp.param(0, dtypes.imageh((8, 1024, 4)))
+  c2 = UOp.range(16, 3, AxisType.LOOP)
+  c5 = UOp.range(256, 4, AxisType.LOOP)
+  c8 = UOp.range(8, 2, AxisType.LOOP)
+  c16 = UOp.range(7, 0, AxisType.REDUCE)
+  c17 = c8*2+c16
+  c24 = ((c17<3)!=True)&(c17<19)
+  c26 = UOp.range(7, 1, AxisType.REDUCE)
+  c27 = c2*2+c26
+  c32 = ((c27<3)!=True)&(c27<35)
+  c34 = UOp.param(1, dtypes.imageh((16, 1024, 4)))
+  c38 = c5//2
+  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.weakint, Invalid))
+  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
+  c49 = UOp.param(2, dtypes.imageh((128, 49, 4)))
+  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
+  c63 = UOp.param(3, dtypes.float.ptr(256))
+  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
+  c67 = c0.index((c2*256+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
+
+  opts = None
+  # JITBEAM=2
+  # (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.SWAP, axis=1, arg=2))
+  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
+
+def dm_conv_172():
+  c0 = UOp.param(0, dtypes.imageh((1, 240, 4)))
+  c2 = UOp.range(960, 4, AxisType.LOOP)
+  c5 = UOp.param(1, dtypes.imageh((8, 384, 4)))
+  c7 = UOp.range(32, 0, AxisType.REDUCE)
+  c10 = UOp.range(4, 1, AxisType.REDUCE)
+  c13 = UOp.range(12, 3, AxisType.REDUCE)
+  c18 = UOp.range(8, 2, AxisType.REDUCE)
+  c23 = UOp.param(2, dtypes.imageh((240, 128, 4)))
+  c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
+  c37 = UOp.param(3, dtypes.float.ptr(960))
+  c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
+  c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
+  c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
+  c55 = c0.index(c2, ptr=True).store(c53).end(c2)
+
+  opts = None
+  # JITBEAM=2
+  # (Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.GROUPTOP, axis=1, arg=32), Opt(op=OptOps.UNROLL, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.GROUP, axis=1, arg=0))
+  return c55.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
+
+ast = {143: vision_conv_143, 153: vision_conv_153, 172: dm_conv_172}[getenv("NUM", 143)]()
+
+renderer = Device.default.renderer
+allocator = Device.default.allocator
+
+ps = to_program(ast, renderer)
+rt = get_runtime(Device.DEFAULT, ps)
+
+gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.PARAM]), key=lambda u: u.arg)
+# print(len(gs))
+# print([g.dtype for g in gs])
+bufs = [Buffer(ps.arg.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]
+
+gsize, lsize = ps.arg.launch_dims({})
+t = rt(*[b._buf for b in bufs], global_size=gsize, local_size=lsize, vals=ps.arg.vals({}), wait=True)
+print(f"{t*1e6:.2f} us")

tinygrad_repo/test/external/external_benchmark_pyrender.py

@@ -0,0 +1,34 @@
+# benchmark speed of pyrender for all created UOps saved with TRACK_MATCH_STATS=2
+import functools, pickle
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.helpers import tqdm, temp, time_to_str, cpu_profile
+
+BENCHMARK_OPS = {Ops.INDEX, Ops.STAGE}
+
+@functools.cache
+def create_uop(a:int) -> UOp:
+  op, dtype, src, arg, *rest = trace.uop_fields[a]
+  return UOp(op, dtype, tuple(create_uop(s) for s in src), arg, *rest)
+
+if __name__ == "__main__":
+  # load rewrite trace
+  with open(temp("rewrites.pkl", append_user=True), "rb") as f:
+    trace = pickle.load(f)
+
+  # benchmark
+  result:list[tuple[str, int]] = []
+  try:
+    for steps in tqdm(trace.rewrites):
+      for r in steps:
+        for _,yn,_,__ in r.matches:
+          y = create_uop(yn)
+          if y.op in BENCHMARK_OPS:
+            with cpu_profile("pyrender") as e:
+              try: ren = y.render()
+              except Exception: ren = "PYRENDER_ERR"
+            result.append((ren, float(e.en-e.st)/1e6))
+  finally:
+    N = 10
+    print(f"Slowst {N} renders from {len(result)} samples:")
+    for ren,tm in sorted(result, key=lambda x:x[1], reverse=True)[:N]:
+      print(f"{time_to_str(tm).strip():<10s} {ren}")

tinygrad_repo/test/external/external_benchmark_resnet.py

@@ -0,0 +1,127 @@
+import functools
+import time
+import unittest
+
+from tinygrad import Tensor, TinyJit, GlobalCounters, Device
+from tinygrad.helpers import getenv, Context
+from tinygrad.nn.optim import SGD
+from tinygrad.nn.state import get_parameters
+from tinygrad.engine.realize import run_linear
+
+from extra.models import resnet
+from examples.mlperf.initializers import Conv2dHeNormal, Linear
+from examples.hlb_cifar10 import UnsyncedBatchNorm
+
+# benchmark memory or kernel count: DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py
+# benchmark speed:                  BEAM=2 JITCNT=10 DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py
+# benchmark only one layer:         BEAM=2 DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py BenchmarkResnetTrain.test_layer1_2
+# inspect:                          DEBUG=2 BEAM=2 DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py
+# inspect 1x1 convs:                DEBUG=2 BEAM=2 CONV=2 DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py
+# inspect 3x3 convs:                DEBUG=2 BEAM=2 CONV=2 DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py
+# inspect 3x3 convs with batchnorm: DEBUG=2 BEAM=2 CONV=2 BN=1 DEFAULT_FLOAT=HALF python test/external/external_benchmark_resnet.py
+# etc
+
+# use ASSIGN=0 to disable batchnorm/optimizer assigns
+
+# memory will be slightly high with JITCNT > 1
+
+bs = getenv("BS", 64)
+
+class BenchmarkResnetTrain(unittest.TestCase):
+  def _get_layer(self, layer_i, slice_i):
+    # isolate to conv, with or without BN
+    conv = getenv("CONV", 0)
+    bn = getenv("BN", 0)
+
+    if not hasattr(self, 'model'):
+      resnet.Conv2d = Conv2dHeNormal
+      resnet.Linear = Linear
+      if not getenv("SYNCBN"): resnet.BatchNorm = functools.partial(UnsyncedBatchNorm, num_devices=1)
+      self.model = resnet.ResNet50()
+      self.layers = [self.model.layer1, self.model.layer2, self.model.layer3, self.model.layer4]
+
+    layer = self.layers[layer_i][slice_i]
+    xy = 112 >> layer_i
+    xy >>= (1 if slice_i > 0 or layer_i == 0 else 0)  # layer 1 is preceded by maxpool2d
+    name = f"layer{layer_i+1} slice{slice_i+1}"
+
+    # get specific conv
+    if conv:
+      convs = [layer.conv1, layer.conv2, layer.conv3] + ([layer.downsample[0]] if layer.downsample else [])
+      bns = [layer.bn1, layer.bn2, layer.bn3] + ([layer.downsample[1]] if layer.downsample else [])
+      f = [convs[conv-1]]
+      if bn: f.append(bns[conv-1])
+      f.append(Tensor.relu)
+      cin = f[0].in_channels
+      if conv == 3: xy //= convs[1].stride
+      return f"{name} conv{conv} x{str((bs, cin, xy, xy)):20s} k{str(f[0].weight.shape):20s}" + (" bn" if bn else ""), f, cin, xy
+
+    cin = layer.conv1.in_channels
+    return f"{name} x{(bs, cin, xy, xy)}", [layer], cin, xy
+  def _test_layer(self, name, layer, cin, xy):
+    optim = SGD(get_parameters(layer), bs / 128 * 1.0)  # need sgd for some params but not consequential for benchmarking
+    with Context(TRACK_MATCH_STATS=0): Tensor.realize(*[t.assign(t.detach().contiguous()) for t in get_parameters(optim)])
+
+    JITCNT = getenv("JITCNT", 1)
+    Tensor.training = True
+    @TinyJit
+    def step(x):
+      optim.zero_grad()
+      x.grad = None
+
+      y = x.sequential(layer).contiguous().contiguous_backward()
+      y.sum().backward()
+      if getenv("ASSIGN", 1): linear, var_vals = Tensor.linear_with_vars(y, x.grad, *optim.schedule_step())
+      else: linear, var_vals = Tensor.linear_with_vars(y, x.grad, *[t.grad for t in optim.params])
+
+      for _ in range(JITCNT):
+        run_linear(linear, var_vals)
+
+    CNT = getenv("CNT", 5)
+    best_tm = None
+    flops, mem_used, mem, kernels = None, None, None, None
+    for i in range(CNT):
+      with Context(TRACK_MATCH_STATS=0): x = Tensor.randn(bs, cin, xy, xy).realize()
+      GlobalCounters.reset()
+
+      st = time.perf_counter()
+      step(x)
+      Device[Device.DEFAULT].synchronize()
+      et = time.perf_counter()
+
+      flops = GlobalCounters.global_ops / JITCNT
+      mem_used = GlobalCounters.mem_used  # a little high with JITCNT > 1 fsr
+      mem = GlobalCounters.global_mem / JITCNT
+      if kernels is None: kernels = GlobalCounters.kernel_count // JITCNT
+      tm = (et-st) / JITCNT
+      if best_tm is None or tm < best_tm: best_tm = tm
+    print(f"\r{name:38s}: {best_tm * 1000:>9.2f} ms, {flops / 10**12 / best_tm:>6.2f} tflops, {mem / 10**9 / best_tm:>5.0f} GB/s, "
+          f"{mem_used / 10**9: 6.2f} GB used, {kernels:>5d} kernels")
+    return best_tm, flops, mem, kernels
+
+  def test_layer1_1(self): self._est(*self._test_layer(*self._get_layer(0, 0)), 1)
+  def test_layer1_2(self): self._est(*self._test_layer(*self._get_layer(0, 1)), 2)
+  def test_layer2_1(self): self._est(*self._test_layer(*self._get_layer(1, 0)), 1)
+  def test_layer2_2(self): self._est(*self._test_layer(*self._get_layer(1, 1)), 3)
+  def test_layer3_1(self): self._est(*self._test_layer(*self._get_layer(2, 0)), 1)
+  def test_layer3_2(self): self._est(*self._test_layer(*self._get_layer(2, 1)), 5)
+  def test_layer4_1(self): self._est(*self._test_layer(*self._get_layer(3, 0)), 1)
+  def test_layer4_2(self): self._est(*self._test_layer(*self._get_layer(3, 1)), 2)
+
+  est_tm, est_flops, est_mem, est_kernels = 0, 0, 0, 0
+
+  @classmethod
+  def _est(cls, tm, flops, mem, kernels, mult):
+    cls.est_tm += tm * mult
+    cls.est_flops += flops * mult
+    cls.est_mem += mem * mult
+    cls.est_kernels += kernels * mult
+
+  @classmethod
+  def tearDownClass(cls):
+    print(f"\restimated step tm: {cls.est_tm * 1000.0:.2f} ms, {cls.est_flops / 10 ** 12 / cls.est_tm:.3f} tflops, "
+          f"{cls.est_mem / 10 ** 9 / cls.est_tm:.2f} GB/s, {cls.est_kernels} kernels")
+
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_benchmark_schedule.py

@@ -0,0 +1,45 @@
+from extra.models.resnet import ResNet50
+from tinygrad import Tensor, nn, Device
+from tinygrad.helpers import Profiling, Timing, getenv
+from tinygrad.uop.ops import Ops
+from tinygrad.codegen import full_rewrite_to_sink
+from tinygrad.codegen.late.linearizer import linearize
+from tinygrad.uop.spec import type_verify, spec_program
+
+if __name__ == "__main__":
+  mdl = ResNet50()
+  for p in nn.state.get_parameters(mdl): p.replace(Tensor.empty(p.shape))
+  img = Tensor.empty(64, 3, 224, 224)
+
+  PROFILE = getenv("PYPROFILE", 0)
+  FORWARD_ONLY = getenv("FORWARD_ONLY", 0)
+  SCHEDULE_ONLY = getenv("SCHEDULE_ONLY", 0)
+  LINEARIZE = bool(getenv("LINEARIZE", 1))
+
+  with Timing("all "):
+    with Timing("***** model tensor in    "):
+      out = mdl(img)
+
+    if not FORWARD_ONLY:
+      with Timing("***** model schedule in  "):
+        with Profiling(PROFILE >= 3):
+          linear = out.schedule_linear()
+
+      if not SCHEDULE_ONLY:
+        asts = list({call.src[0].key:call.src[0] for call in linear.src if call.src[0].op is Ops.SINK}.values())
+        if (restrict_kernel := getenv("RESTRICT_KERNEL", -1)) != -1: asts = asts[restrict_kernel:restrict_kernel+1]
+
+        with Profiling(PROFILE, fn="/tmp/rewrite.prof"):
+          with Timing("***** model rewrite in   "):
+            rewritten_uops = []
+            for u in asts:
+              rewritten_uops.append(full_rewrite_to_sink(u, ren=Device.default.renderer))
+
+        if LINEARIZE:
+          with Timing("***** model linearize in "):
+            uops_line = []
+            for u in rewritten_uops:
+              uops_line.append(linearize(u))
+          with Timing("***** model verify in    "):
+            for u in uops_line: type_verify(u, spec_program)
+          print(sum(len(u) for u in uops_line))

tinygrad_repo/test/external/external_cl_half_max.py

@@ -0,0 +1,13 @@
+from tinygrad.runtime.ops_cl import CLDevice, CLProgram, compile_cl
+
+if __name__ == "__main__":
+  dev = CLDevice()
+  lib = compile_cl("""
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__kernel void test(__global half *out, __global half *a, __global half *b) {
+  int gid = get_global_id(0);
+  out[gid] = max(a[gid], b[gid]);
+}
+""")
+  prg = CLProgram(dev, "test", lib)
+

tinygrad_repo/test/external/external_fuzz_ampt.py

@@ -0,0 +1,109 @@
+import random
+from typing import Optional
+from tinygrad.helpers import round_up
+from tinygrad.runtime.support.am.amdev import AMPageTableTraverseContext
+from test.external.external_test_am import helper_read_entry_components, FakeAM
+
+class AMPTFuzzer:
+  def __init__(self, total_size):
+    self.total_size = total_size
+    self.alloc_payload = 0
+    self.d = FakeAM()
+
+    self.allocations: dict[int, tuple[int, int]] = {} # ptr -> (size, pattern)
+
+    self.min_alloc_size = 0x1000
+    self.max_alloc_size = int(total_size * 0.1)
+    self.alloc_probability = 0.7
+
+  def generate_pattern(self, ptr: int, size: int) -> int: return random.randint(0, 0xff)
+
+  def fill_memory(self, va, size: int, pattern: int):
+    ctx = AMPageTableTraverseContext(self.d, self.d.mm.root_page_table, va.va_addr)
+    pages = list(ctx.next(size))
+
+    for _offset, _pt, _pte_idx, _n_ptes, _pte_covers in pages:
+      _vaddr = va.va_addr + _offset
+
+      for i in range(_n_ptes):
+        pte = helper_read_entry_components(_pt.entries[_pte_idx + i])
+        self.d.vram[pte['paddr']] = pattern # Mark this page
+        assert pte['valid'] == 1
+
+        # If page has contiguous fragment, all range should be this valid memory
+        frags_cnt = pte['fragment']
+        contig_range = (1 << (frags_cnt + 12))
+        start_vaddr = _vaddr & ~(contig_range - 1)
+        start_paddr = pte['paddr'] - (_vaddr - start_vaddr)
+        contig_ptes = contig_range // _pte_covers
+        assert contig_ptes > 0
+
+        ctx = AMPageTableTraverseContext(self.d, self.d.mm.root_page_table, start_vaddr)
+        frags_l = list(ctx.next(contig_range))
+        for f_offset, f_pt, f_pte_idx, f_n_ptes, f_pte_covers in frags_l:
+          for j in range(f_n_ptes):
+            f_pte = helper_read_entry_components(f_pt.entries[f_pte_idx + j])
+            assert f_pte['valid'] == 1
+            assert f_pte['paddr'] == start_paddr+f_offset+j*f_pte_covers, f"paddr {f_pte['paddr']:#x} not {start_paddr+f_offset+j*f_pte_covers:#x}"
+
+        _vaddr += _pte_covers
+        _offset += _pte_covers
+
+    return pages
+
+  def verify_memory(self, pages, pattern: int) -> bool:
+    for _offset, _pt, _pte_idx, _n_ptes, _pte_covers in pages:
+      for i in range(_n_ptes):
+        pte = helper_read_entry_components(_pt.entries[_pte_idx + i])
+        if self.d.vram[pte['paddr']] != pattern: return False
+        if pte['valid'] == 0: return False
+
+    return True
+
+  def random_alloc(self) -> Optional[int]:
+    if self.total_size - self.alloc_payload < self.min_alloc_size: return None
+
+    size = random.randint(self.min_alloc_size, min(self.max_alloc_size, self.total_size - self.alloc_payload))
+    size = round_up(size, (2 << 20) if size > (4 << 20) else (4 << 10))
+
+    try: ptr = self.d.mm.valloc(size)
+    except MemoryError:
+      print(f"Failed to allocate {size} bytes. Payload size is {self.alloc_payload}, so fragmenation is {(size / self.total_size)*100.0:.2f}%")
+      return None
+
+    pattern = self.generate_pattern(ptr, size)
+    pages = self.fill_memory(ptr, size, pattern)
+    self.allocations[ptr.va_addr] = (size, pattern, pages, ptr)
+    self.alloc_payload += size
+    print(f"Allocated {size} bytes at {ptr.va_addr:x}, pattern: {pattern:02x}")
+    return ptr
+
+  def random_free(self) -> bool:
+    if not self.allocations: return False
+
+    ptr = random.choice(list(self.allocations.keys()))
+    size, pattern, pages, vm = self.allocations[ptr]
+
+    # Verify pattern before freeing
+    if not self.verify_memory(pages, pattern):
+      raise RuntimeError(f"Memory corruption detected at {vm.va_addr:x}!")
+
+    print(f"Freeing {size} bytes at {vm.va_addr:x}, pattern verified: {pattern:02x}")
+    self.alloc_payload -= size
+    self.d.mm.vfree(vm)
+    del self.allocations[ptr]
+    return True
+
+  def run(self):
+    for i in range(10000000):
+      if (random.random() < self.alloc_probability or not self.allocations): self.random_alloc()
+      else: self.random_free()
+
+    print("\nCleaning up remaining allocations...")
+    while self.allocations: self.random_free()
+
+    print("Fuzzing completed successfully!")
+
+if __name__ == "__main__":
+  fuzzer = AMPTFuzzer(1 << 30)
+  fuzzer.run()

tinygrad_repo/test/external/external_fuzz_beam_timeout_recovery.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+"""
+Stress test for beam timeout + device recovery on AM devices.
+
+Usage:
+  DEV=AMD python test/external/external_fuzz_beam_timeout_recovery.py
+"""
+from tinygrad import Tensor, Device
+from tinygrad.helpers import Context
+from tinygrad.runtime.ops_amd import AMDDevice
+
+if __name__ == "__main__":
+  dev = Device["AMD"]
+  assert isinstance(dev, AMDDevice) and dev.is_am(), "not am"
+
+  N = 10000
+  for i in range(N):
+    with Context(DEBUG=0, BEAM=0):
+      a = Tensor.rand(4096, 4096, device="AMD").contiguous().realize()
+      b = Tensor.rand(4096, 4096, device="AMD").contiguous().realize()
+      c = a.matmul(b)
+      c.realize()
+    try: dev.synchronize(timeout=1)
+    except RuntimeError as e: print(e)
+    with Context(DEBUG=0, BEAM=0):
+      a = Tensor.ones(512, 512, device="AMD").contiguous().realize()
+      b = Tensor.ones(512, 512, device="AMD").contiguous().realize()
+      result = a.matmul(b).realize()[0, 0].item()
+    assert result == 512.0, f"iter {i}: got {result}"
+    print(f"  iter {i+1}/{N}: ok")
+  print(f"=== All {N} iterations passed ===")

tinygrad_repo/test/external/external_fuzz_hcq_mp.py

@@ -0,0 +1,44 @@
+import subprocess
+import random
+import time
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from tinygrad.helpers import getenv
+
+# checks that HCQ drivers can be killed during operation without causing issues
+
+def run_test(i, full_run=False, force_ok=False):
+  print(f"\rRunning iteration {i}...", end=" ", flush=True)
+
+  p = subprocess.Popen(["python3", "test/test_tiny.py", "TestTiny.test_plus"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+  if not full_run:
+    time.sleep(random.uniform(0, 1200) / 1000.0)
+    p.kill()
+    _, stderr = p.communicate()
+  else:
+    _, stderr = p.communicate()
+    stderr_text = stderr.decode()
+    assert ("Ran 1 test in" in stderr_text and "OK" in stderr_text) or (not force_ok and "Failed to take lock file" in stderr_text), stderr_text
+
+if __name__ == "__main__":
+  max_workers = getenv("MAX_WORKERS", 4)
+  with ProcessPoolExecutor(max_workers=max_workers) as executor:
+    futures = []
+    for i in range(1000000):
+      if i % 100 == 0:
+        # wait for everything we launched so far
+        for f in as_completed(futures):
+          try:
+            f.result()
+          except Exception as e:
+            print(f"\nError in iteration: {e}")
+        futures = []
+
+        # do a full run in the main proc
+        run_test(i, True, force_ok=True)
+      else:
+        futures.append(executor.submit(run_test, i, bool(getenv("FULL_RUN", 0))))
+
+      # keep list small
+      if len(futures) > max_workers * 2:
+        futures = [f for f in futures if not f.done()]

tinygrad_repo/test/external/external_fuzz_hcq_signals.py

@@ -0,0 +1,31 @@
+import random
+from tinygrad import Device
+from tinygrad.helpers import getenv, DEBUG
+
+def main():
+  seed = getenv("SEED", 1337)
+  n_gpus = getenv("GPUS", 3)
+  iters = getenv("ITERS", 10000000)
+  only_compute = bool(getenv("ONLY_COMPUTE", 0))
+
+  print(f"{n_gpus} GPUs for {iters} iterations, {only_compute=}, seed {seed}")
+  devs = tuple([Device[f"{Device.DEFAULT}:{x}"] for x in range(n_gpus)])
+
+  for i in range(iters):
+    dev = random.choice(devs)
+    q_t = random.choice([dev.hw_copy_queue_t, dev.hw_compute_queue_t] if not only_compute else [dev.hw_compute_queue_t])
+
+    deps_sigs = random.randint(0, len(devs))
+    wait_devs = random.sample(devs, deps_sigs)
+
+    q = q_t()
+    for d in wait_devs: q.wait(d.timeline_signal, d.timeline_value - 1)
+    q.wait(dev.timeline_signal, dev.timeline_value - 1).signal(dev.timeline_signal, dev.timeline_value).submit(dev)
+    dev.timeline_value += 1
+
+    if sync:=random.randint(0, 10) < 3: dev.synchronize()
+    if DEBUG >= 2: print(f"{i}: {q_t} {dev.device_id} timeline {dev.timeline_value}, wait for {[d.device_id for d in wait_devs]}, {sync=}")
+    elif i % 100 == 0: print(f"\rCompleted {i} iterations", end='')
+
+if __name__ == "__main__":
+  main()

tinygrad_repo/test/external/external_fuzz_qcom_cpu_cache.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+from tinygrad.tensor import Tensor
+import numpy as np
+
+while True:
+  arr = np.ones(1000000, dtype=np.uint8)
+  print(f"numpy: {(arr + 1)[:10]}")
+
+  ptr = arr.ctypes.data
+  tensor = Tensor.from_blob(ptr, arr.shape, dtype='uint8', device='QCOM').realize() + 1
+  print(f"from_blob: {tensor.numpy()[:10]}")

tinygrad_repo/test/external/external_fuzz_sdma_warm_start.py

@@ -0,0 +1,55 @@
+import subprocess, sys, os, random
+
+CHILD_SCRIPT = """
+import os, random
+import numpy as np
+from tinygrad import Tensor, Device
+from tinygrad.runtime.ops_amd import AMDDevice
+
+dev = Device["AMD"]
+for i in range({N}):
+  sz = random.randint(1, {MAX_SZ})
+  data = np.random.randint(0, 256, sz, dtype=np.uint8)
+  t = Tensor(data, device="AMD").contiguous().realize()
+  dev.synchronize()
+  result = t.numpy()
+  assert (result == data).all(), f"Data mismatch at iter {{i}}"
+""".strip()
+
+def run_child(n_ops, max_sz, timeout):
+  env = os.environ.copy()
+  env.setdefault("SDMA_RING_SIZE", "4096")
+
+  script = CHILD_SCRIPT.format(N=n_ops, MAX_SZ=max_sz)
+  p = subprocess.Popen([sys.executable, "-c", script], stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)
+
+  try:
+    _, stderr = p.communicate(timeout=timeout)
+    return ("ok" if p.returncode == 0 else "fail"), stderr.decode(errors='replace')
+  except subprocess.TimeoutExpired:
+    p.kill()
+    p.communicate()
+    return "timeout", "TIMEOUT: SDMA ring likely stuck"
+
+if __name__ == "__main__":
+  n_iters = int(os.environ.get("FUZZ_ITERS", "10000"))
+  timeout = int(os.environ.get("FUZZ_TIMEOUT", "10"))
+  max_sz = int(os.environ.get("FUZZ_MAX_SZ", "65536"))
+
+  timeouts = 0
+  failures = 0
+
+  for i in range(n_iters):
+    # Run child with many ops to stress the small sdma ring buffer across warm starts
+    n_ops = random.randint(20, 100)
+    status, stderr = run_child(n_ops=n_ops, max_sz=max_sz, timeout=timeout)
+    if status == "timeout":
+      timeouts += 1
+      print(f"\tstderr: {stderr[:500]}")
+    elif status == "fail":
+      failures += 1
+      print(f"\tstderr: {stderr[:500]}")
+    else:
+      print(f"iter {i}: ok (n_ops={n_ops})")
+
+  print(f"\n=== Results: {n_iters} iterations, {timeouts} timeouts, {failures} failures ===")

tinygrad_repo/test/external/external_fuzz_tlsf.py

@@ -0,0 +1,81 @@
+import random
+from typing import Dict, Optional
+from tinygrad.helpers import getenv
+from tinygrad.runtime.support.memory import TLSFAllocator
+
+class AllocatorFuzzer:
+  def __init__(self, total_size):
+    self.total_size = total_size
+    self.alloc_payload = 0
+    self.mv = memoryview(bytearray(total_size))
+    self.alloctor = TLSFAllocator(total_size, block_size=16)
+
+    self.allocations: Dict[int, tuple[int, int]] = {} # ptr -> (size, pattern)
+
+    self.min_alloc_size = 16
+    self.max_alloc_size = int(total_size * 0.3)
+    self.alloc_probability = 0.7
+
+  def generate_pattern(self, ptr: int, size: int) -> int: return (ptr * 31 + size * 17) & 0xFF
+
+  def fill_memory(self, ptr: int, size: int, pattern: int):
+    for i in range(min(size, 32)):
+      self.mv[ptr + i] = pattern
+      self.mv[ptr + (size - 1 - i)] = pattern
+
+  def verify_memory(self, ptr: int, size: int, pattern: int) -> bool:
+    for i in range(min(size, 32)):
+      assert self.mv[ptr + i] == pattern
+      assert self.mv[ptr + (size - 1 - i)] == pattern
+    return True
+
+  def random_alloc(self) -> Optional[int]:
+    if self.total_size - self.alloc_payload < self.min_alloc_size: return None
+
+    size = random.randint(self.min_alloc_size, min(self.max_alloc_size, self.total_size - self.alloc_payload))
+
+    try:
+      ptr = self.alloctor.alloc(size)
+    except MemoryError:
+      print(f"Failed to allocate {size} bytes. Payload size is {self.alloc_payload}, so fragmenation is {(size / self.total_size)*100.0:.2f}%")
+      return None
+
+    pattern = self.generate_pattern(ptr, size)
+    self.fill_memory(ptr, size, pattern)
+    self.allocations[ptr] = (size, pattern)
+    self.alloc_payload += size
+    print(f"Allocated {size} bytes at {ptr:x}, pattern: {pattern:02x}")
+    return ptr
+
+  def random_free(self) -> bool:
+    if not self.allocations: return False
+
+    ptr = random.choice(list(self.allocations.keys()))
+    size, pattern = self.allocations[ptr]
+
+    # Verify pattern before freeing
+    if not self.verify_memory(ptr, size, pattern):
+      raise RuntimeError(f"Memory corruption detected at {ptr:x}!")
+
+    print(f"Freeing {size} bytes at {ptr:x}, pattern verified: {pattern:02x}")
+    self.alloc_payload -= size
+    self.alloctor.free(ptr)
+    del self.allocations[ptr]
+    return True
+
+  def run(self):
+    for i in range(getenv("ITERS", 100000)):
+      if (random.random() < self.alloc_probability or not self.allocations): self.random_alloc()
+      else: self.random_free()
+
+    print("\nCleaning up remaining allocations...")
+    while self.allocations: self.random_free()
+
+    print("Fuzzing completed successfully!")
+
+if __name__ == "__main__":
+  SEED = getenv("SEED", 42)
+  random.seed(SEED)
+
+  fuzzer = AllocatorFuzzer(1 << 30)
+  fuzzer.run()

tinygrad_repo/test/external/external_gpu_fail_osx.py

@@ -0,0 +1,16 @@
+# ugh, OS X OpenCL doesn't support half
+from tinygrad.runtime.ops_cl import CLDevice, CLProgram, CLCompiler
+
+src = """#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__kernel void max_half(__global half* data0, const __global half* data1) {
+  int gidx0 = get_group_id(0);
+  data0[gidx0] = max(data1[gidx0], (half)0.0);
+}"""
+
+if __name__ == "__main__":
+  dev = CLDevice()
+  print("created device")
+  lib = CLCompiler(dev, "test").compile(src)
+  print("created lib", len(lib))
+  prg = CLProgram(dev, "max_half", lib)
+  print("created prg")

tinygrad_repo/test/external/external_jit_failure.py

@@ -0,0 +1,17 @@
+from tinygrad import Tensor, TinyJit, Device
+import numpy as np
+
+GPUS = 4
+N = 128
+ds = tuple([Device.canonicalize(f"{Device.DEFAULT}:{i}") for i in range(GPUS)])
+t = Tensor.rand(N, N, N).shard(ds, 0)
+n = t.numpy()
+
+@TinyJit
+def allreduce(t:Tensor) -> Tensor:
+  return t.sum(0) #.realize()
+
+for i in range(10):
+  print(i)
+  tn = allreduce(t).numpy()
+  np.testing.assert_allclose(tn, n.sum(0), atol=1e-4, rtol=1e-4)

tinygrad_repo/test/external/external_llm_eval.py

@@ -0,0 +1,50 @@
+# eval for OpenAI API server
+# uses Meta's exact ARC-Challenge prompt template from lm-evaluation-harness llama3 tasks
+import argparse, re, pyarrow.parquet as pq
+from openai import OpenAI
+from tinygrad.helpers import fetch, colored
+
+LABEL = ["A", "B", "C", "D"]
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--port", "-p", type=int, default=8000)
+  parser.add_argument("--limit", "-L", type=int, default=None)
+  parser.add_argument("--max_tokens", "-T", type=int, default=4096)
+  parser.add_argument("--offset", "-O", type=int, default=0)
+  parser.add_argument("--temperature", "-t", type=float, default=0.0)
+  parser.add_argument("--no_think", action="store_true", help="disable thinking (prefills empty think block via assistant message)")
+  parser.add_argument("--debug", action="store_true")
+  args = parser.parse_args()
+
+  client = OpenAI(base_url=f"http://127.0.0.1:{args.port}/v1", api_key="tinygrad")
+  dat = fetch("https://huggingface.co/datasets/allenai/ai2_arc/resolve/main/ARC-Challenge/test-00000-of-00001.parquet")
+  table = pq.read_table(dat)
+
+  num_correct, num_answered = 0, 0
+  # filter to 4-choice questions and normalize labels to A/B/C/D (matches Meta's eval)
+  rows = [(q, c, a) for q, c, a in zip(table["question"], table["choices"], table["answerKey"]) if len(c["label"]) == 4]
+  total_questions = min(len(rows), args.offset + args.limit) if args.limit else len(rows)
+  for question, choices, answer in rows[args.offset:total_questions]:
+    phrasing = "Given the following question and four candidate answers (A, B, C and D), choose the best answer.\n" +\
+               f"Question: {question}\n" + '\n'.join([f"{l}. {t}" for l, t in zip(LABEL, choices['text'])]) +\
+               '\nYour response should end with "The best answer is [the_answer_letter]"' +\
+               " where the [the_answer_letter] is one of A, B, C or D."
+    messages = [{"role": "user", "content": phrasing}]
+    if args.no_think: messages.append({"role": "assistant", "content": "<think>\n\n</think>\n\n"})
+    resp = client.chat.completions.create(model="test", messages=messages,
+                                          max_tokens=args.max_tokens, temperature=args.temperature)
+    # normalize answer key (some use 1/2/3/4 instead of A/B/C/D)
+    correct = answer.as_py().strip()
+    if correct not in LABEL: correct = LABEL[int(correct) - 1]
+    # extract answer: take last single capital letter A-D from response (prompt asks model to end with the answer)
+    text = resp.choices[0].message.content.strip()
+    if args.debug: print(f"\n--- PROMPT ---\n{phrasing}\n--- RESPONSE ---\n{text}\n---")
+    m = re.findall(r'\b([A-D])\b', text)
+    given = m[-1] if m else text[:1]
+    num_correct += correct == given
+    num_answered += 1
+    print(f"{num_answered:4d}/{total_questions:4d}  "+\
+          f"Correct Answer: {correct}  "+\
+          f"Given Answer: {colored(given, 'green' if correct==given else 'red')}  "+\
+          f"Percent: {num_correct*100.0/num_answered:.2f}%")

tinygrad_repo/test/external/external_metal_compile_fail.py

@@ -0,0 +1,230 @@
+# AssertionError: Error Domain=AGXMetalG15X_B0 Code=3 "Compiler encountered an internal error"
+
+src = """
+#include <metal_stdlib>
+using namespace metal;
+kernel void r_64_32_8_16_4_6_6_4(device float* data0, const device float* data1,
+                                 uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {
+  int gidx0 = gid.x; /* 64 */
+  int lidx2 = lid.x; /* 8 */
+  int gidx1 = gid.y; /* 32 */
+  int lidx3 = lid.y; /* 16 */
+  int alu0 = ((gidx0*4096)+(gidx1*16)+(lidx2*512)+lidx3);
+  int alu1 = ((gidx0*147456)+(gidx1*576)+(lidx2*18432)+(lidx3*36));
+  float acc0 = 0.0f;
+  float acc1 = 0.0f;
+  float acc2 = 0.0f;
+  float acc3 = 0.0f;
+  float acc4 = 0.0f;
+  float acc5 = 0.0f;
+  float acc6 = 0.0f;
+  float acc7 = 0.0f;
+  float acc8 = 0.0f;
+  float acc9 = 0.0f;
+  float acc10 = 0.0f;
+  float acc11 = 0.0f;
+  float acc12 = 0.0f;
+  float acc13 = 0.0f;
+  float acc14 = 0.0f;
+  float acc15 = 0.0f;
+  float acc16 = 0.0f;
+  float acc17 = 0.0f;
+  float acc18 = 0.0f;
+  float acc19 = 0.0f;
+  float acc20 = 0.0f;
+  float acc21 = 0.0f;
+  float acc22 = 0.0f;
+  float acc23 = 0.0f;
+  float acc24 = 0.0f;
+  float acc25 = 0.0f;
+  float acc26 = 0.0f;
+  float acc27 = 0.0f;
+  float acc28 = 0.0f;
+  float acc29 = 0.0f;
+  float acc30 = 0.0f;
+  float acc31 = 0.0f;
+  float acc32 = 0.0f;
+  float acc33 = 0.0f;
+  float acc34 = 0.0f;
+  float acc35 = 0.0f;
+  for (int ridx0 = 0; ridx0 < 4; ridx0++) {
+    int alu2 = (ridx0*6);
+    int alu3 = (alu2+1);
+    int alu4 = (alu2+2);
+    int alu5 = (alu2+3);
+    int alu6 = (alu2+4);
+    int alu7 = (alu2+5);
+    int alu8 = (alu2%7);
+    int alu9 = ((alu8+1)%7);
+    int alu10 = ((alu8+2)%7);
+    int alu11 = ((alu8+3)%7);
+    int alu12 = ((alu8+4)%7);
+    int alu13 = ((alu8+5)%7);
+    int alu14 = ((((alu0+(alu3/21))%262144)*144)+(((alu3/7)%3)*3)+(alu9*36));
+    int alu15 = ((((alu0+(alu4/21))%262144)*144)+(((alu4/7)%3)*3)+(alu10*36));
+    int alu16 = ((((alu0+(alu5/21))%262144)*144)+(((alu5/7)%3)*3)+(alu11*36));
+    int alu17 = ((((alu0+(alu6/21))%262144)*144)+(((alu6/7)%3)*3)+(alu12*36));
+    int alu18 = ((((alu0+(alu7/21))%262144)*144)+(((alu7/7)%3)*3)+(alu13*36));
+    int alu19 = (alu8%7);
+    int alu20 = ((((alu0+(alu2/21))%262144)*144)+(((alu2/7)%3)*3)+(alu19*36));
+    bool alu21 = ((alu2<16)&(alu13<4));
+    bool alu22 = ((alu2<17)&(alu12<4));
+    bool alu23 = ((alu2<18)&(alu11<4));
+    bool alu24 = ((alu2<19)&(alu10<4));
+    bool alu25 = ((alu2<20)&(alu9<4));
+    bool alu26 = ((alu2<21)&(alu19<4));
+    float val0 = (alu25?*(data1+alu14+1):0.0f);
+    float val1 = (alu25?*(data1+alu14+2):0.0f);
+    float val2 = (alu25?*(data1+alu14+9):0.0f);
+    float val3 = (alu25?*(data1+alu14+10):0.0f);
+    float val4 = (alu25?*(data1+alu14+11):0.0f);
+    float val5 = (alu25?*(data1+alu14+18):0.0f);
+    float val6 = (alu25?*(data1+alu14+19):0.0f);
+    float val7 = (alu25?*(data1+alu14+20):0.0f);
+    float val8 = (alu25?*(data1+alu14+27):0.0f);
+    float val9 = (alu25?*(data1+alu14+28):0.0f);
+    float val10 = (alu25?*(data1+alu14+29):0.0f);
+    float val11 = (alu24?*(data1+alu15+1):0.0f);
+    float val12 = (alu24?*(data1+alu15+2):0.0f);
+    float val13 = (alu24?*(data1+alu15+9):0.0f);
+    float val14 = (alu24?*(data1+alu15+10):0.0f);
+    float val15 = (alu24?*(data1+alu15+11):0.0f);
+    float val16 = (alu24?*(data1+alu15+18):0.0f);
+    float val17 = (alu24?*(data1+alu15+19):0.0f);
+    float val18 = (alu24?*(data1+alu15+20):0.0f);
+    float val19 = (alu24?*(data1+alu15+27):0.0f);
+    float val20 = (alu24?*(data1+alu15+28):0.0f);
+    float val21 = (alu24?*(data1+alu15+29):0.0f);
+    float val22 = (alu23?*(data1+alu16+1):0.0f);
+    float val23 = (alu23?*(data1+alu16+2):0.0f);
+    float val24 = (alu23?*(data1+alu16+9):0.0f);
+    float val25 = (alu23?*(data1+alu16+10):0.0f);
+    float val26 = (alu23?*(data1+alu16+11):0.0f);
+    float val27 = (alu23?*(data1+alu16+18):0.0f);
+    float val28 = (alu23?*(data1+alu16+19):0.0f);
+    float val29 = (alu23?*(data1+alu16+20):0.0f);
+    float val30 = (alu23?*(data1+alu16+27):0.0f);
+    float val31 = (alu23?*(data1+alu16+28):0.0f);
+    float val32 = (alu23?*(data1+alu16+29):0.0f);
+    float val33 = (alu22?*(data1+alu17+1):0.0f);
+    float val34 = (alu22?*(data1+alu17+2):0.0f);
+    float val35 = (alu22?*(data1+alu17+9):0.0f);
+    float val36 = (alu22?*(data1+alu17+10):0.0f);
+    float val37 = (alu22?*(data1+alu17+11):0.0f);
+    float val38 = (alu22?*(data1+alu17+18):0.0f);
+    float val39 = (alu22?*(data1+alu17+19):0.0f);
+    float val40 = (alu22?*(data1+alu17+20):0.0f);
+    float val41 = (alu22?*(data1+alu17+27):0.0f);
+    float val42 = (alu22?*(data1+alu17+28):0.0f);
+    float val43 = (alu22?*(data1+alu17+29):0.0f);
+    float val44 = (alu21?*(data1+alu18+1):0.0f);
+    float val45 = (alu21?*(data1+alu18+2):0.0f);
+    float val46 = (alu21?*(data1+alu18+9):0.0f);
+    float val47 = (alu21?*(data1+alu18+10):0.0f);
+    float val48 = (alu21?*(data1+alu18+11):0.0f);
+    float val49 = (alu21?*(data1+alu18+18):0.0f);
+    float val50 = (alu21?*(data1+alu18+19):0.0f);
+    float val51 = (alu21?*(data1+alu18+20):0.0f);
+    float val52 = (alu21?*(data1+alu18+27):0.0f);
+    float val53 = (alu21?*(data1+alu18+28):0.0f);
+    float val54 = (alu21?*(data1+alu18+29):0.0f);
+    float val55 = (alu26?*(data1+alu20+1):0.0f);
+    float val56 = (alu26?*(data1+alu20+2):0.0f);
+    float val57 = (alu26?*(data1+alu20+9):0.0f);
+    float val58 = (alu26?*(data1+alu20+10):0.0f);
+    float val59 = (alu26?*(data1+alu20+11):0.0f);
+    float val60 = (alu26?*(data1+alu20+18):0.0f);
+    float val61 = (alu26?*(data1+alu20+19):0.0f);
+    float val62 = (alu26?*(data1+alu20+20):0.0f);
+    float val63 = (alu26?*(data1+alu20+27):0.0f);
+    float val64 = (alu26?*(data1+alu20+28):0.0f);
+    float val65 = (alu26?*(data1+alu20+29):0.0f);
+    float val66 = (alu25?*(data1+alu14):0.0f);
+    float val67 = (alu24?*(data1+alu15):0.0f);
+    float val68 = (alu23?*(data1+alu16):0.0f);
+    float val69 = (alu22?*(data1+alu17):0.0f);
+    float val70 = (alu21?*(data1+alu18):0.0f);
+    float val71 = (alu26?*(data1+alu20):0.0f);
+    acc0 = (acc0+val71);
+    acc1 = (acc1+val66);
+    acc2 = (acc2+val67);
+    acc3 = (acc3+val68);
+    acc4 = (acc4+val69);
+    acc5 = (acc5+val70);
+    acc6 = (acc6+val57+val55);
+    acc7 = (acc7+val2+val0);
+    acc8 = (acc8+val13+val11);
+    acc9 = (acc9+val24+val22);
+    acc10 = (acc10+val35+val33);
+    acc11 = (acc11+val46+val44);
+    acc12 = (acc12+val60+val58+val56);
+    acc13 = (acc13+val5+val3+val1);
+    acc14 = (acc14+val16+val14+val12);
+    acc15 = (acc15+val27+val25+val23);
+    acc16 = (acc16+val38+val36+val34);
+    acc17 = (acc17+val49+val47+val45);
+    acc18 = (acc18+val63+val61+val59);
+    acc19 = (acc19+val8+val6+val4);
+    acc20 = (acc20+val19+val17+val15);
+    acc21 = (acc21+val30+val28+val26);
+    acc22 = (acc22+val41+val39+val37);
+    acc23 = (acc23+val52+val50+val48);
+    acc24 = (acc24+val64+val62);
+    acc25 = (acc25+val9+val7);
+    acc26 = (acc26+val20+val18);
+    acc27 = (acc27+val31+val29);
+    acc28 = (acc28+val42+val40);
+    acc29 = (acc29+val53+val51);
+    acc30 = (acc30+val65);
+    acc31 = (acc31+val10);
+    acc32 = (acc32+val21);
+    acc33 = (acc33+val32);
+    acc34 = (acc34+val43);
+    acc35 = (acc35+val54);
+  }
+  *(data0+alu1+1) = acc6;
+  *(data0+alu1+2) = acc12;
+  *(data0+alu1+3) = acc18;
+  *(data0+alu1+4) = acc24;
+  *(data0+alu1+5) = acc30;
+  *(data0+alu1+6) = acc1;
+  *(data0+alu1+7) = acc7;
+  *(data0+alu1+8) = acc13;
+  *(data0+alu1+9) = acc19;
+  *(data0+alu1+10) = acc25;
+  *(data0+alu1+11) = acc31;
+  *(data0+alu1+12) = acc2;
+  *(data0+alu1+13) = acc8;
+  *(data0+alu1+14) = acc14;
+  *(data0+alu1+15) = acc20;
+  *(data0+alu1+16) = acc26;
+  *(data0+alu1+17) = acc32;
+  *(data0+alu1+18) = acc3;
+  *(data0+alu1+19) = acc9;
+  *(data0+alu1+20) = acc15;
+  *(data0+alu1+21) = acc21;
+  *(data0+alu1+22) = acc27;
+  *(data0+alu1+23) = acc33;
+  *(data0+alu1+24) = acc4;
+  *(data0+alu1+25) = acc10;
+  *(data0+alu1+26) = acc16;
+  *(data0+alu1+27) = acc22;
+  *(data0+alu1+28) = acc28;
+  *(data0+alu1+29) = acc34;
+  *(data0+alu1+30) = acc5;
+  *(data0+alu1+31) = acc11;
+  *(data0+alu1+32) = acc17;
+  *(data0+alu1+33) = acc23;
+  *(data0+alu1+34) = acc29;
+  *(data0+alu1+35) = acc35;
+  *(data0+alu1) = acc0;
+}
+"""
+
+from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
+
+if __name__ == "__main__":
+  dev = MetalDevice("METAL")
+  lib = MetalCompiler().compile(src)
+  prg = MetalProgram(dev, "r_64_32_8_16_4_6_6_4", lib)
+

tinygrad_repo/test/external/external_model_benchmark.py

@@ -0,0 +1,139 @@
+import csv, pathlib, time
+import numpy as np
+import torch
+torch.set_num_threads(1)
+import onnxruntime as ort
+from onnx2torch import convert
+from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.helpers import OSX, DEBUG, fetch, getenv
+from tinygrad.dtype import _to_np_dtype
+from tinygrad import Tensor, Device, Context, dtypes
+
+MODELS = {
+  "resnet50": "https://github.com/onnx/models/raw/main/validated/vision/classification/resnet/model/resnet50-caffe2-v1-9.onnx",
+  "openpilot": "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx",
+  "efficientnet": "https://github.com/onnx/models/raw/main/validated/vision/classification/efficientnet-lite4/model/efficientnet-lite4-11.onnx",
+  "shufflenet": "https://github.com/onnx/models/raw/main/validated/vision/classification/shufflenet/model/shufflenet-9.onnx",
+  # TODO: precision issue
+  # "commavq": "https://huggingface.co/commaai/commavq-gpt2m/resolve/main/gpt2m.onnx",
+  "dm": "https://github.com/commaai/openpilot/raw/ba7f840a06dbc8ae3c45b3b4976c88a21895aed0/selfdrive/modeld/models/dmonitoring_model.onnx",
+
+  # broken in torch MPS
+  # "zfnet": "https://github.com/onnx/models/raw/main/archive/vision/classification/zfnet-512/model/zfnet512-9.onnx",
+  # TypeError: BatchNormalization() got an unexpected keyword argument 'is_test'
+  # "densenet": "https://github.com/onnx/models/raw/main/archive/vision/classification/densenet-121/model/densenet-3.onnx",
+  # AssertionError: only onnx version >= 10 supported for slice
+  # "bert": "https://github.com/onnx/models/raw/main/archive/text/machine_comprehension/bert-squad/model/bertsquad-8.onnx",
+  # really slow
+  # "resnet18": "https://github.com/onnx/models/raw/main/archive/vision/classification/resnet/model/resnet18-v2-7.onnx",
+}
+half_models = ["openpilot", "commavq"]
+
+CSV = {}
+open_csv = None
+
+def benchmark(mnm, nm, fxn):
+  tms = []
+  for _ in range(3):
+    st = time.perf_counter_ns()
+    ret = fxn()
+    tms.append(time.perf_counter_ns() - st)
+  print(f"{mnm:15s} {nm:25s} {min(tms)*1e-6:7.2f} ms")
+  CSV[nm] = min(tms)*1e-6
+  return min(tms), ret
+
+#BASE = pathlib.Path(__file__).parents[2] / "weights" / "onnx"
+BASE = pathlib.Path("/tmp/onnx")
+def benchmark_model(m, devices, validate_outs=False):
+  torch.manual_seed(1)
+  global open_csv, CSV
+  CSV = {"model": m}
+
+  fn = fetch(MODELS[m])
+  runner = OnnxRunner(fn)
+  output_names = runner.graph_outputs
+  input_shapes = {name: tuple(s if isinstance(s, int) and s != 0 else 1 for s in spec.shape) for name, spec in runner.graph_inputs.items()}
+  input_types = {name: spec.dtype for name, spec in runner.graph_inputs.items()}
+  np_inputs = {k:torch.randn(shp).numpy().astype(_to_np_dtype(input_types[k])) for k,shp in input_shapes.items()}
+  assert len(input_shapes) < 30, f"too many input shapes {len(input_shapes)}"
+
+  # print input names
+  if DEBUG >= 2: print(list(runner.graph_inputs))
+  for device in devices:
+    with Context(DEV=device):
+      inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
+      tinygrad_model = runner.to(device)
+      benchmark(m, f"tinygrad_{device.lower()}_jitless", lambda: {k:v.numpy() for k,v in tinygrad_model(inputs).items()})
+
+      from tinygrad.engine.jit import TinyJit
+      tinygrad_jitted_model = TinyJit(lambda **kwargs: {k:v.realize() for k,v in tinygrad_model(kwargs).items()})
+      for _ in range(3): {k:v.numpy() for k,v in tinygrad_jitted_model(**inputs).items()}
+      benchmark(m, f"tinygrad_{device.lower()}_jit", lambda: {k:v.numpy() for k,v in tinygrad_jitted_model(**inputs).items()}) # noqa: F821
+      del inputs, tinygrad_model, tinygrad_jitted_model
+
+  # convert model to torch
+  try:
+    torch_model = convert(fn)
+  except Exception as e:
+    # model conversion failed
+    print(f"{m:16s}onnx2torch {type(e).__name__:>25}")
+  else:
+    torch_inputs = [torch.tensor(x) for x in np_inputs.values()]
+    try: benchmark(m, "torch_cpu", lambda: torch_model(*torch_inputs))
+    except Exception as e: print(f"{m:16s}torch_cpu {type(e).__name__:>25}")
+
+    torch_device = "mps" if OSX else "cuda"
+    torch_mps_model = torch_model.to(torch_device)
+    torch_mps_inputs = [x.to(torch_device) for x in torch_inputs]
+    try: benchmark(m, f"torch_{torch_device}", lambda: torch_mps_model(*torch_mps_inputs))
+    except Exception as e: print(f"{m:16s}torch_{torch_device} {type(e).__name__:>25}")
+
+  # bench onnxruntime
+  ort_options = ort.SessionOptions()
+  ort_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+  ort_options.log_severity_level = 3  # no warnings
+  for backend in ["CPU", "CUDA" if not OSX else "CoreML"]:  # https://onnxruntime.ai/docs/execution-providers/
+    provider = backend+"ExecutionProvider"
+    if provider not in ort.get_available_providers(): continue
+    ort_sess = ort.InferenceSession(str(fn), ort_options, [provider])
+    try:
+      benchmark(m, f"onnxruntime_{backend.lower()}", lambda: ort_sess.run(output_names, np_inputs))
+    except Exception as e: print(f"{m:16s}onnxruntime_{backend.lower()} {type(e).__name__:>25}")
+    del ort_sess
+
+  if validate_outs:
+    for device in devices:
+      rtol, atol = 2e-3, 2e-3  # tolerance for fp16 models
+      with Context(DEV=device):
+        # force half inputs to float for numerical stability when validating
+        # this will rely on automatic dtype promotion for converting half weights inside the graph
+        if m in half_models:
+          inputs = {k:Tensor(inp, dtype=dtypes.float32) if inp.dtype == np.float16 else Tensor(inp) for k,inp in np_inputs.items()}
+        else:
+          inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
+        tinygrad_model = runner.to(device)
+        tinygrad_out = tinygrad_model(inputs)
+
+        ort_sess = ort.InferenceSession(str(fn), ort_options, ["CPUExecutionProvider"])
+        onnx_out = ort_sess.run(output_names, np_inputs)
+        onnx_out = dict([*list(zip(output_names, onnx_out))])
+
+        assert_allclose(tinygrad_out, onnx_out, rtol=rtol, atol=atol)
+        print(f"{m:16s}outputs validated on {device=} with rtol={rtol:.1e}, atol={atol:.1e}")
+
+  if open_csv is None:
+    open_csv = csv.DictWriter(open('onnx_inference_speed.csv', 'w', newline=''), fieldnames=list(CSV.keys()))
+    open_csv.writeheader()
+  open_csv.writerow(CSV)
+
+def assert_allclose(tiny_out:dict, onnx_out:dict, rtol, atol):
+  assert tiny_out.keys() == onnx_out.keys()
+  for k in tiny_out.keys():
+    tiny_v, onnx_v = tiny_out[k], onnx_out[k]
+    np.testing.assert_allclose(tiny_v.numpy(), onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tiny_out.keys()}")
+
+if __name__ == "__main__":
+  devices = [Device.DEFAULT] if getenv("NOCLANG") else [Device.DEFAULT, "CPU"]
+  if (model:=getenv("MODEL", "")) != "": benchmark_model(model, devices, validate_outs=True)
+  else:
+    for m in MODELS: benchmark_model(m, devices, validate_outs=True)

tinygrad_repo/test/external/external_multi_gpu.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+# cd extra/disassemblers/ && git clone --recursive github.com:geohot/cuda_ioctl_sniffer.git
+# LD_PRELOAD=$PWD/extra/disassemblers/cuda_ioctl_sniffer/out/sniff.so DEV=CL python3 test/external/external_multi_gpu.py
+import numpy as np
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import colored, Timing, getenv
+from tinygrad.device import Device
+
+d0, d1 = f'{Device.DEFAULT}:0', f'{Device.DEFAULT}:1'
+
+def sync():
+  Device[d0].synchronize()
+  Device[d1].synchronize()
+
+if __name__ == "__main__":
+  print("GPU devices", d0, d1)
+  sz = getenv("N", 1024*1024*256)  # 1 GB
+
+  with Timing("GPU initial sync: "): sync()
+
+  with Timing("CPU creation: ", on_exit=lambda x: f", {(sz*4*2)/x:.2f} GB/sec"):
+    c0 = (Tensor.ones(sz, device="CPU")/2).realize()
+    c1 = (Tensor.ones(sz, device="CPU")/4).realize()
+    print(c0.uop.base.realized)
+    print(c1.uop.base.realized)
+
+  with Timing("CPU -> 0: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    a0 = c0.to(d0).realize()
+    sync()
+  with Timing("CPU -> 1: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    b1 = c1.to(d1).realize()
+    sync()
+
+  # cross copy. this is (sometimes) going through the CPU
+  with Timing("0 -> 1: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    a1 = a0.to(d1).realize()
+    sync()
+  with Timing("1 -> 0: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    b0 = b1.to(d0).realize()
+    sync()
+
+  # sum
+  with Timing("0+0 -> 0 (sum): ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    ab0 = (a0 + b0).realize()
+    sync()
+  with Timing("1+1 -> 1 (sum): ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    ab1 = (a1 + b1).realize()
+    sync()
+
+  # cross device sum (does this work?)
+  with Timing(colored("0+1 -> 0 (sum): ", "red"), on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    abx0 = (a0 + b1.to(d0)).realize()
+    sync()
+
+  with Timing(colored("1+0 -> 1 (sum): ", "red"), on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    abx1 = (b1 + a0.to(d1)).realize()
+    sync()
+
+  # copy back
+  # NOTE: half of this slowness is caused by allocating memory on the CPU
+  with Timing("0 -> CPU: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    cc0 = ab0.numpy()
+  with Timing("1 -> CPU: ", on_exit=lambda x: f", {(sz*4)/x:.2f} GB/sec"):
+    cc1 = ab1.numpy()
+
+  # same
+  print("testing")
+  np.testing.assert_allclose(cc0, cc1)
+
+  # same (cross)
+  print("testing (cross)")
+  np.testing.assert_allclose(cc0, abx0.numpy())
+  np.testing.assert_allclose(cc0, abx1.numpy())
+
+  # devices
+  print(ab0)
+  print(ab1)
+  print(abx0)
+  print(abx1)

tinygrad_repo/test/external/external_nhwc_conv.py

@@ -0,0 +1,38 @@
+from tinygrad import Tensor, nn, Context, GlobalCounters
+
+if __name__ == "__main__":
+  conv = nn.Conv2d(64, 128, 3)
+  img = Tensor.randn((1,64,128,128))
+  with Context(DEBUG=0, BEAM=0):
+    Tensor.realize(img, conv.weight, conv.bias)
+
+  tst = conv(img).permute(0,2,3,1).realize()
+  print(tst.shape)
+
+  print("NEW")
+  img_perm = img.permute(0,2,3,1).contiguous()
+  print(img_perm.shape)
+  pp = img_perm.permute(0,3,1,2)._pool((3,3)).permute(0,2,3,4,5,1)
+
+  def hwio(pp, conv):
+    pp = pp.unsqueeze(-1)
+    weight = conv.weight.permute(2,3,1,0).contiguous()
+    print(pp.shape, weight.shape, (pp*weight).shape)
+    return (pp * weight).sum([-4,-3, -2])
+
+  def ohwi(pp, conv):
+    pp = pp.unsqueeze(-4)
+    weight = conv.weight.permute(0,2,3,1).contiguous()
+    print(pp.shape, weight.shape, (pp*weight).shape)
+    return (pp * weight).sum([-3,-2,-1])
+
+  for f in [hwio, ohwi]:
+    GlobalCounters.reset()
+    print("\n**************", f.__name__, "**************")
+    out = f(pp, conv)
+    out.realize()
+    print(out.shape)
+
+    with Context(DEBUG=0, BEAM=0):
+      err = (tst-out).square()
+      print(err.mean().item(), err.max().item())

tinygrad_repo/test/external/external_openpilot_image_warp.py

@@ -0,0 +1,55 @@
+import time
+from tinygrad.tensor import Tensor, Device
+
+MODEL_WIDTH = 512
+MODEL_HEIGHT = 256
+MODEL_FRAME_SIZE = MODEL_WIDTH * MODEL_HEIGHT * 3 // 2
+IMG_INPUT_SHAPE = (1, 12, 128, 256)
+
+def tensor_arange(end): return Tensor([float(i) for i in range(end)])
+def tensor_round(tensor:Tensor): return (tensor + 0.5).floor()
+
+h_src, w_src = 1208, 1928
+h_dst, w_dst = MODEL_HEIGHT, MODEL_WIDTH
+x = tensor_arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst)
+y = tensor_arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst)
+ones = Tensor.ones_like(x)
+dst_coords = x.reshape((1,-1)).cat(y.reshape((1,-1))).cat(ones.reshape((1,-1)))
+
+def warp_perspective_tinygrad(src:Tensor, M_inv:Tensor) -> Tensor:
+  src_coords = M_inv @ dst_coords
+  src_coords = src_coords / src_coords[2:3, :]
+
+  x_src = src_coords[0].reshape(h_dst, w_dst)
+  y_src = src_coords[1].reshape(h_dst, w_dst)
+
+  x_nearest = tensor_round(x_src).clip(0, w_src - 1).cast('int')
+  y_nearest = tensor_round(y_src).clip(0, h_src - 1).cast('int')
+
+  # TODO: make 2d indexing fast
+  idx = y_nearest*src.shape[1] + x_nearest
+  dst = src.flatten()[idx]
+  return dst.reshape(h_dst, w_dst)
+
+if __name__ == "__main__":
+  from tinygrad.engine.jit import TinyJit
+  update_img_jit = TinyJit(warp_perspective_tinygrad, prune=True)
+
+  step_times = []
+  for _ in range(10):
+    # regenerate inputs
+    inputs = [Tensor.randn(1928,1208), Tensor.randn(3,3)]
+    Tensor.realize(*inputs)
+    Device.default.synchronize()
+
+    # do the warp
+    st = time.perf_counter()
+    out = update_img_jit(*inputs)
+    mt = time.perf_counter()
+    val = out.contiguous().realize()
+    Device.default.synchronize()
+    et = time.perf_counter()
+
+    # measure the time
+    step_times.append((et-st)*1e3)
+    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")

tinygrad_repo/test/external/external_osx_profiling.py

@@ -0,0 +1,41 @@
+from tinygrad.runtime.ops_cl import CLProgram, CL, CLBuffer
+from tinygrad import dtypes
+import time
+
+N = 1000000
+a = CLBuffer(N, dtypes.float32)
+b = CLBuffer(N, dtypes.float32)
+c = CLBuffer(N, dtypes.float32)
+
+prg = CLProgram("test", """__kernel void test(__global float *a, __global float *b, __global float *c) {
+  int idx = get_global_id(0);
+  a[idx] = b[idx] + c[idx];
+}""")
+prg.clprgs[0](CL.cl_queue[0], [N,], None, a._buf, b._buf, c._buf)
+t1 = time.monotonic_ns()
+e1 = prg.clprgs[0](CL.cl_queue[0], [N,], None, a._buf, b._buf, c._buf)
+CL.synchronize()
+t2 = time.monotonic_ns()
+time.sleep(3)
+t3 = time.monotonic_ns()
+e2 = prg.clprgs[0](CL.cl_queue[0], [N,], None, a._buf, b._buf, c._buf)
+CL.synchronize()
+t4 = time.monotonic_ns()
+
+print(e1.profile.queued)
+print(e1.profile.submit)
+print(e1.profile.start)
+print(e1.profile.end)
+
+print(e1, e2)
+print(t2-t1, e1.profile.end - e1.profile.start)
+print(t4-t3, e2.profile.end - e2.profile.start)
+print(t3-t2, e2.profile.queued-e1.profile.end)
+print((t3-t2) / (e2.profile.start-e1.profile.end), "ratio")
+
+print("ratio since boot", t1/e1.profile.start)
+
+print(e1.profile.start)
+print(e1.profile.end)
+print(e2.profile.start)
+print(e2.profile.end)

tinygrad_repo/test/external/external_test_am.py

@@ -0,0 +1,224 @@
+import unittest
+from tinygrad.runtime.support.am.amdev import AMMemoryManager, AMPageTableEntry
+from tinygrad.runtime.support.am.ip import AM_GMC
+from tinygrad.runtime.support.hcq import MMIOInterface
+from tinygrad.runtime.support.memory import PageTableTraverseContext, AddrSpace
+from tinygrad.runtime.autogen.am import am
+from tinygrad.helpers import mv_address
+
+class FakeGMC(AM_GMC):
+  def __init__(self, adev):
+    self.adev = adev
+    self.vm_base = 0x0
+    self.address_space_mask = (1 << 44) - 1
+  def init_hw(self): pass
+  def flush_tlb(self, *args, **kwargs): pass
+
+class FakePCIDev:
+  def __init__(self): self.regions = [(0xc12300000000, 0xc12400000000, 0x0)]
+
+class FakeAM:
+  def __init__(self):
+    self.is_booting, self.smi_dev = True, False
+    self.pcidev = FakePCIDev()
+    self.vram_size = (512 << 20)
+    self.vram_mv = memoryview(bytearray(self.vram_size))
+    self.vram = MMIOInterface(mv_address(self.vram_mv), self.vram_mv.nbytes)
+    self.gmc = FakeGMC(self)
+    self.mm = AMMemoryManager(self, self.vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39], va_bits=48,
+      first_lv=am.AMDGPU_VM_PDB2, va_base=AMMemoryManager.va_allocator.base,
+      palloc_ranges=[(1 << (i + 12), (2 << 20) if i >= 9 else 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
+    self.is_booting = False
+    self.ip_ver = {am.GC_HWIP: (11, 0, 0)}
+  def paddr2cpu(self, paddr:int) -> int: return paddr + mv_address(self.vram)
+  def paddr2mc(self, paddr:int) -> int: return paddr
+  def paddr2xgmi(self, paddr:int) -> int: return paddr
+  def xgmi2paddr(self, xgmi_paddr:int) -> int: return xgmi_paddr
+
+#  * PTE format:
+#  * 63:59 reserved
+#  * 58:57 reserved
+#  * 56 F
+#  * 55 L
+#  * 54 reserved
+#  * 53:52 SW
+#  * 51 T
+#  * 50:48 mtype
+#  * 47:12 4k physical page base address
+#  * 11:7 fragment
+#  * 6 write
+#  * 5 read
+#  * 4 exe
+#  * 3 Z
+#  * 2 snooped
+#  * 1 system
+#  * 0 valid
+def helper_read_entry_components(entry_val):
+  return {"paddr": entry_val & 0x0000FFFFFFFFF000, "fragment":(entry_val >> 7) & 0x1f, "valid": entry_val & 0x1,
+          "read": (entry_val >> 5) & 0x1, "write": (entry_val >> 6) & 0x1, "exec": (entry_val >> 4) & 0x1,
+          "mtype": (entry_val >> 48) & 0x7, "T": (entry_val >> 51) & 0x1, "L": (entry_val >> 55) & 0x1, "F": (entry_val >> 56) & 0x1}
+
+def helper_va(va:int): return va + AMMemoryManager.va_allocator.base
+
+class TestAMPageTable(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.d = [FakeAM() for _ in range(2)]
+
+  def test_page_table_walkers(self):
+    mm = self.d[0].mm
+
+    for va,sz in [(0x10000, 0x3000), (0x11000, 0x300000), (0x10000, 0x2000), (0x11000, 0x5000),
+                  (0x2000000, 0x2000), (0x4000000, 0x4000000), (0x38000, 0x303000), (0x8000, 0x1000)]:
+      mm.map_range(vaddr=helper_va(va), size=sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+
+      ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, helper_va(va))
+      results = list(ctx.next(sz))
+
+      total_covered = 0
+      for tup in results:
+        _offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
+        total_covered += _n_ptes * _pte_covers
+
+      assert total_covered == sz, f"Expected total coverage {total_covered} to be {sz}"
+
+      for tup in results:
+        _offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
+        for i in range(_n_ptes):
+          pte = helper_read_entry_components(_pt.entries[_pte_idx + i])
+          assert pte['paddr'] == va + _offset + i * _pte_covers, f"Expected paddr {pte['paddr']:#x} to be {va + _offset + i * _pte_covers:#x}"
+          assert pte['valid'] == 1
+
+      mm.unmap_range(helper_va(va), sz)
+
+      for tup in results:
+        _offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
+        for i in range(_n_ptes):
+          pte = helper_read_entry_components(_pt.entries[_pte_idx + i])
+          assert pte['paddr'] == 0
+          assert pte['valid'] == 0
+
+  def test_map_notaligned(self):
+    mm0 = self.d[0].mm
+
+    for (va1,sz1),(va2,sz2) in [((0x10000, (0x1000)), (0x11000, (2 << 20)))]:
+      mm0.map_range(vaddr=helper_va(va1), size=sz1, paddrs=[(va1, sz1)], aspace=AddrSpace.PHYS)
+      mm0.map_range(vaddr=helper_va(va2), size=sz2, paddrs=[(va2, sz2)], aspace=AddrSpace.PHYS)
+      mm0.unmap_range(helper_va(va2), sz2)
+      mm0.unmap_range(helper_va(va1), sz1)
+
+  def test_double_map(self):
+    mm0 = self.d[0].mm
+
+    for va,sz in [(0x10000, 0x3000), (0x1000000, 0x1000000), (0x12000, 0x4000)]:
+      exteranl_va = helper_va(va)
+      mm0.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+
+      with self.assertRaises(AssertionError):
+        mm0.map_range(vaddr=exteranl_va, size=0x1000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+
+      with self.assertRaises(AssertionError):
+        mm0.map_range(vaddr=exteranl_va, size=0x100000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+
+      with self.assertRaises(AssertionError):
+        mm0.map_range(vaddr=exteranl_va + 0x1000, size=0x1000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+
+      with self.assertRaises(AssertionError):
+        mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+
+      mm0.unmap_range(vaddr=exteranl_va, size=sz)
+
+      # Finally can map and check paddrs
+      mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(0xdead0000, 0x1000), (0xdead1000, 0xff000)], aspace=AddrSpace.PHYS)
+
+      ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, exteranl_va + 0x2000)
+      for tup in ctx.next(0x100000):
+        _offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
+        for i in range(_n_ptes):
+          pte = helper_read_entry_components(_pt.entries[_pte_idx + i])
+          assert pte['paddr'] == 0xdead0000 + _offset + i * _pte_covers, f"paddr {pte['paddr']:#x} not {0xdead0000 + _offset + i * _pte_covers:#x}"
+          assert pte['valid'] == 1
+
+      mm0.unmap_range(vaddr=exteranl_va + 0x2000, size=0x100000)
+
+  def test_try_bad_unmap(self):
+    mm0 = self.d[0].mm
+
+    with self.assertRaises(AssertionError):
+      mm0.unmap_range(helper_va(0x10000), 0x3000)
+
+    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)], aspace=AddrSpace.PHYS)
+    mm0.unmap_range(helper_va(0x10000), 0x3000)
+
+    with self.assertRaises(AssertionError):
+      mm0.unmap_range(helper_va(0x10000), 0x3000)
+
+    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)], aspace=AddrSpace.PHYS)
+    mm0.unmap_range(helper_va(0x10000), 0x3000)
+
+    with self.assertRaises(AssertionError):
+      mm0.unmap_range(helper_va(0x10000), 0x3000)
+
+  def test_free_pt(self):
+    mm0 = self.d[0].mm
+
+    # offset from start
+    for off in [0, 0x3000, 0x10000]:
+      mm0.map_range(helper_va(0x1000000) + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000), aspace=AddrSpace.PHYS)
+      mm0.unmap_range(helper_va(0x1000000) + off, (2 << 20) - off)
+      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
+      mm0.unmap_range(helper_va(0x1000000), 2 << 20)
+
+    # offset from end
+    for off in [0x1000, 0x20000]:
+      mm0.map_range(helper_va(0x1000000), (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000), aspace=AddrSpace.PHYS)
+      mm0.unmap_range(helper_va(0x1000000), (2 << 20) - off)
+      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
+      mm0.unmap_range(helper_va(0x1000000), 2 << 20)
+
+  def test_inspect_mode(self):
+    mm0 = self.d[0].mm
+
+    # Map a few disjoint ranges inside a larger region.
+    mappings = [(0x10000, 0x3000), (0x20000, 0x2000), (0x1000000, 2 << 20)]
+    for va, sz in mappings:
+      mm0.map_range(helper_va(va), sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+
+    # Inspect over the whole region: should visit all mapped pages.
+    ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, helper_va(0x0), inspect=True)
+    visited = set()
+    for _off, pt, pte_idx, n_ptes, pte_covers in ctx.next(0x4000000):
+      for i in range(n_ptes):
+        pte = helper_read_entry_components(pt.entries[pte_idx + i])
+        if pte['valid']:
+          for p in range(0, pte_covers, 0x1000): visited.add(pte['paddr'] + p)
+
+    expected_pages = {va + off for va, sz in mappings for off in range(0, sz, 0x1000)}
+    assert visited == expected_pages
+
+    for va, sz in mappings:
+      mm0.unmap_range(helper_va(va), sz)
+
+    # Inspect after unmap: should find no valid entries.
+    ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, helper_va(0x0), inspect=True)
+    for _off, pt, pte_idx, n_ptes, pte_covers in ctx.next(0x4000000):
+      for i in range(n_ptes): assert not pt.valid(pte_idx + i)
+
+  def test_frag_size(self):
+    mm0 = self.d[0].mm
+
+    def must_cover_checker(va, sz):
+      ans = (1 << (mm0._frag_size(va=va, sz=sz, must_cover=True) + 12))
+      assert va % ans == 0 and sz % ans == 0 and (va % (2 * ans) != 0 or sz % (2 * ans) != 0), f"va {va:#x} sz {sz:#x} ans {ans:#x}"
+
+    def not_cover_checker(va, sz):
+      ans = (1 << (mm0._frag_size(va=va, sz=sz, must_cover=False) + 12))
+      assert va % ans == 0 and ans <= sz and (va % (2 * ans) != 0 or (2 * ans) > sz), f"va {va:#x} sz {sz:#x} ans {ans:#x}"
+
+    for va, sz in [(0x0, 0x1000), (0x1000, 0x2000), (0x1000, 0x3000), (0x2000, 0x2000), (0x4000, 0x8000), (0x8000, 0x4000), (0x10000, 0x4000),
+                   (0x0, 0x4000), (0x10000, 0x4000), (0x10000, 0x40000), (0x10001000, 0x40000), (0x100001000, 0x3000)]:
+      must_cover_checker(va, sz)
+      not_cover_checker(va, sz)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/external/external_test_am_fault_recovery.py

@@ -0,0 +1,139 @@
+# ruff: noqa: F405
+import unittest, subprocess, os
+from tinygrad.helpers import DEV
+from tinygrad.runtime.autogen.amd.rdna3.ins import *  # noqa: F403
+from tinygrad.renderer.amd.dsl import s, v, Inst, NULL
+
+def assemble_kernel(insts:list[Inst], name:str="test") -> str:
+  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "wavefront_size32": 1, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
+  from test.amd.disasm import disasm as _disasm
+  disasm = "\n".join(_disasm(inst) for inst in insts)
+  hsasrc = f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{disasm}\n"
+  return hsasrc + f".rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + "\n".join(f".amdhsa_{k} {v}" for k, v in kd.items()) + "\n.end_amdhsa_kernel"
+
+def _run(code:str, timeout:float=15.0) -> subprocess.CompletedProcess:
+  # TODO: AM_RESET is required for now, so subprocesses
+  return subprocess.run(["python", "-c", code], env={**os.environ, "AMD": "1"}, capture_output=True, text=True, timeout=timeout)
+
+def _run_asm(asm_src:str) -> subprocess.CompletedProcess:
+  return _run('from tinygrad.device import Device; from tinygrad.runtime.ops_amd import AMDProgram; '
+              'from tinygrad.runtime.support.compiler_amd import HIPCompiler; dev = Device["AMD"]; '
+              f'AMDProgram(dev, "test", HIPCompiler(dev.arch).compile("""{asm_src}"""))('
+              'dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)')
+
+def _verify_recovery() -> subprocess.CompletedProcess:
+  return _run('from tinygrad import Tensor; t = Tensor([1.0, 2.0], device="AMD").realize(); assert (t + 1).numpy().tolist() == [2.0, 3.0]')
+
+_ILLEGAL_INST_ASM = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n.byte 0xff,0xff,0xff,0xff\ns_endpgm\n" \
+  ".rodata\n.p2align 6\n.amdhsa_kernel test\n.amdhsa_next_free_vgpr 8\n.amdhsa_next_free_sgpr 8\n" \
+  ".amdhsa_wavefront_size32 1\n.amdhsa_user_sgpr_kernarg_segment_ptr 1\n.amdhsa_kernarg_size 8\n.end_amdhsa_kernel"
+
+@unittest.skipIf(DEV.device != "AMD" or not DEV.interface.startswith("MOCK"), "AMD with AM driver required")
+class TestAMFaultRecovery(unittest.TestCase):
+  def _run_kernel(self, insts: list[Inst]) -> subprocess.CompletedProcess: return _run_asm(assemble_kernel(insts))
+
+  def _assert_fault_and_recovery(self, result:subprocess.CompletedProcess):
+    if result.stdout.strip(): print(f"\nstdout: {result.stdout.strip()}")
+    if result.stderr.strip(): print(f"\nstderr: {result.stderr.strip()}")
+    self.assertNotEqual(result.returncode, 0, f"Expected fault but succeeded: {result.stdout}")
+    self.assertEqual(_verify_recovery().returncode, 0)
+
+
+class TestGlobalMemoryFaults(TestAMFaultRecovery):
+  def test_global_load_unmapped(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_global_store_unmapped(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
+             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_global_null_ptr(self):
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_global_misaligned_b64(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0001), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b64(v[2:3], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_global_misaligned_b128(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0004), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b128(v[2:5], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+
+class TestSMEMFaults(TestAMFaultRecovery):
+  def test_smem_null_base(self):
+    insts = [s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
+             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_smem_unmapped_address(self):
+    insts = [s_mov_b32(s[2], 0xBEEF0000), s_mov_b32(s[3], 0xDEAD),
+             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_smem_misaligned_b64(self):
+    insts = [s_mov_b32(s[2], 0xBEEF0004), s_mov_b32(s[3], 0xDEAD),
+             s_load_b64(s[4:5], s[2:3], 0, soffset=NULL), s_waitcnt(lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_smem_misaligned_b128(self):
+    insts = [s_mov_b32(s[2], 0xBEEF0004), s_mov_b32(s[3], 0xDEAD),
+             s_load_b128(s[4:7], s[2:3], 0, soffset=NULL), s_waitcnt(lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+
+class TestIllegalInstruction(TestAMFaultRecovery):
+  def test_malformed_encoding(self):
+    self._assert_fault_and_recovery(_run_asm(_ILLEGAL_INST_ASM))
+
+
+class TestFlatFaults(TestAMFaultRecovery):
+  def test_flat_load_unmapped(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD),
+             flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0, lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_flat_store_unmapped(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
+             flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(vmcnt=0, lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+
+class TestAtomicFaults(TestAMFaultRecovery):
+  def test_global_atomic_unmapped(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
+             global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+  def test_flat_atomic_unmapped(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
+             flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(vmcnt=0, lgkmcnt=0), s_endpgm()]
+    self._assert_fault_and_recovery(self._run_kernel(insts))
+
+
+class TestRecovery(TestAMFaultRecovery):
+  def test_recovery_after_memviol(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    self.assertNotEqual(self._run_kernel(insts).returncode, 0)
+    self.assertEqual(_verify_recovery().returncode, 0)
+
+  def test_recovery_after_illegal_inst(self):
+    self.assertNotEqual(_run_asm(_ILLEGAL_INST_ASM).returncode, 0)
+    self.assertEqual(_verify_recovery().returncode, 0)
+
+  def test_multiple_faults_recovery(self):
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(vmcnt=0), s_endpgm()]
+    for _ in range(3):
+      self.assertNotEqual(self._run_kernel(insts).returncode, 0)
+      self.assertEqual(_verify_recovery().returncode, 0)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/external/external_test_amd.py

@@ -0,0 +1,28 @@
+import unittest
+from tinygrad import Device, Tensor
+from tinygrad.schedule import create_schedule
+from tinygrad.runtime.ops_amd import AMDDevice
+
+class TestAMD(unittest.TestCase):
+  @classmethod
+  def setUpClass(self):
+    TestAMD.d0: AMDDevice = Device["AMD"]
+    TestAMD.a = Tensor([0.,1.], device="AMD").realize()
+    TestAMD.b = self.a + 1
+    si = create_schedule([self.b.uop])[-1]
+    TestAMD.d0_runner = TestAMD.d0.get_runner(*si.ast)
+    TestAMD.b.uop.buffer.allocate()
+
+  def test_amd_ring_64bit_doorbell(self):
+    TestAMD.d0.pm4_write_pointer[0] = TestAMD.d0.pm4_write_pointer[0] + (2 << 32) - TestAMD.d0.pm4_ring.size // 4
+    for _ in range(2000):
+      TestAMD.d0_runner.clprg(TestAMD.b.uop.buffer._buf, TestAMD.a.uop.buffer._buf,
+                              global_size=TestAMD.d0_runner.global_size, local_size=TestAMD.d0_runner.local_size)
+      TestAMD.d0_runner.clprg(TestAMD.a.uop.buffer._buf, TestAMD.b.uop.buffer._buf,
+                              global_size=TestAMD.d0_runner.global_size, local_size=TestAMD.d0_runner.local_size)
+    val = TestAMD.a.uop.buffer.as_memoryview().cast("f")[0]
+    assert val == 4000.0, f"got val {val}"
+
+if __name__ == "__main__":
+  unittest.main()
+

tinygrad_repo/test/external/external_test_datasets.py

@@ -0,0 +1,172 @@
+from extra.datasets.kits19 import iterate, preprocess
+from examples.mlperf.dataloader import batch_load_unet3d, batch_load_retinanet
+from test.external.mlperf_retinanet.coco_utils import get_openimages
+from test.external.mlperf_retinanet.openimages import postprocess_targets
+from test.external.mlperf_retinanet.presets import DetectionPresetTrain, DetectionPresetEval
+from test.external.mlperf_retinanet.model.transform import GeneralizedRCNNTransform
+from test.external.mlperf_unet3d.kits19 import PytTrain, PytVal
+from tinygrad.helpers import temp
+from pathlib import Path
+from pycocotools.coco import COCO
+
+import json
+import nibabel as nib
+import numpy as np
+import os
+import PIL
+import random
+import tempfile
+import torch
+import unittest
+
+class ExternalTestDatasets(unittest.TestCase):
+  def _set_seed(self):
+    np.random.seed(42)
+    random.seed(42)
+    torch.manual_seed(42)
+
+class TestKiTS19Dataset(ExternalTestDatasets):
+  def _create_samples(self, val, num_samples=2):
+    self._set_seed()
+
+    img, lbl = np.random.rand(190, 392, 392).astype(np.float32), np.random.randint(0, 100, size=(190, 392, 392)).astype(np.uint8)
+    img, lbl = nib.Nifti1Image(img, np.eye(4)), nib.Nifti1Image(lbl, np.eye(4))
+    dataset = "val" if val else "train"
+    preproc_pth = Path(tempfile.gettempdir() + f"/{dataset}")
+
+    for i in range(num_samples):
+      os.makedirs(tempfile.gettempdir() + f"/case_000{i}", exist_ok=True)
+      nib.save(img, temp(f"case_000{i}/imaging.nii.gz"))
+      nib.save(lbl, temp(f"case_000{i}/segmentation.nii.gz"))
+
+      preproc_img, preproc_lbl = preprocess(Path(tempfile.gettempdir()) / f"case_000{i}")
+      preproc_img_pth, preproc_lbl_pth = temp(f"{dataset}/case_000{i}_x.npy"), temp(f"{dataset}/case_000{i}_y.npy")
+
+      os.makedirs(preproc_pth, exist_ok=True)
+      np.save(preproc_img_pth, preproc_img, allow_pickle=False)
+      np.save(preproc_lbl_pth, preproc_lbl, allow_pickle=False)
+
+    return preproc_pth, list(preproc_pth.glob("*_x.npy")), list(preproc_pth.glob("*_y.npy"))
+
+  def _create_ref_dataloader(self, preproc_img_pths, preproc_lbl_pths, val):
+    if val:
+      dataset = PytVal(preproc_img_pths, preproc_lbl_pths)
+    else:
+      dataset = PytTrain(preproc_img_pths, preproc_lbl_pths, patch_size=(128, 128, 128), oversampling=0.4)
+
+    return iter(dataset)
+
+  def _create_tinygrad_dataloader(self, preproc_pth, val, batch_size=1, shuffle=False, seed=42, use_old_dataloader=False):
+    if use_old_dataloader:
+      dataset = iterate(list(Path(tempfile.gettempdir()).glob("case_*")), preprocessed_dir=preproc_pth, val=val, shuffle=shuffle, bs=batch_size)
+    else:
+      dataset = batch_load_unet3d(preproc_pth, batch_size=batch_size, val=val, shuffle=shuffle, seed=seed)
+
+    return iter(dataset)
+
+  @unittest.skip("flaky")
+  def test_training_set(self):
+    preproc_pth, preproc_img_pths, preproc_lbl_pths = self._create_samples(False)
+    ref_dataset = self._create_ref_dataloader(preproc_img_pths, preproc_lbl_pths, False)
+    tinygrad_dataset = self._create_tinygrad_dataloader(preproc_pth, False)
+
+    for ref_sample, tinygrad_sample in zip(ref_dataset, tinygrad_dataset):
+      self._set_seed()
+
+      np.testing.assert_equal(tinygrad_sample[0][:, 0].numpy(), ref_sample[0])
+      np.testing.assert_equal(tinygrad_sample[1][:, 0].numpy(), ref_sample[1])
+
+  def test_validation_set(self):
+    preproc_pth, preproc_img_pths, preproc_lbl_pths = self._create_samples(True)
+    ref_dataset = self._create_ref_dataloader(preproc_img_pths, preproc_lbl_pths, True)
+    tinygrad_dataset = self._create_tinygrad_dataloader(preproc_pth, True, use_old_dataloader=True)
+
+    for ref_sample, tinygrad_sample in zip(ref_dataset, tinygrad_dataset):
+      np.testing.assert_equal(tinygrad_sample[0][:, 0], ref_sample[0])
+      np.testing.assert_equal(tinygrad_sample[1], ref_sample[1])
+
+class TestOpenImagesDataset(ExternalTestDatasets):
+  @classmethod
+  def setUpClass(cls):
+    cls.img_mean = [0.485, 0.456, 0.406]
+    cls.img_std = [0.229, 0.224, 0.225]
+    cls.img_size = (800, 800)
+
+  def _create_samples(self, subset):
+    os.makedirs((base_dir := Path(tempfile.gettempdir() + "/openimages")) / f"{subset}/data", exist_ok=True)
+    os.makedirs(base_dir / Path(f"{subset}/labels"), exist_ok=True)
+
+    lbls, img_size = ["cls_1", "cls_2"], (447, 1024)
+    cats = [{"id": i, "name": c, "supercategory": None} for i, c in enumerate(lbls)]
+    imgs = [
+      {
+        "id": i, "file_name": f"img_{i}.jpg",
+        "height": img_size[0], "width": img_size[1],
+        "subset": subset, "license": None, "coco_url": None
+      }
+      for i in range(len(lbls))
+    ]
+    annots = [
+      {
+        "id": i, "image_id": i,
+        "category_id": 0, "bbox": [23.217183744, 31.75409775, 964.1241282560001, 326.09017434000003],
+        "area": 314391.4050683996, "IsOccluded": 0,
+        "IsInside": 0, "IsDepiction": 0,
+        "IsTruncated": 0, "IsGroupOf": 0,
+        "iscrowd": 0
+      }
+      for i in range(len(lbls))
+    ]
+    info = {"dataset": "openimages_mlperf", "version": "v6"}
+    coco_annotations = {"info": info, "licenses": [], "categories": cats, "images": imgs, "annotations": annots}
+
+    with open(ann_file:=base_dir / Path(f"{subset}/labels/openimages-mlperf.json"), "w") as fp:
+      json.dump(coco_annotations, fp)
+
+    for i in range(len(lbls)):
+      img = PIL.Image.new("RGB", img_size[::-1])
+      img.save(base_dir / Path(f"{subset}/data/img_{i}.jpg"))
+
+    return base_dir, ann_file
+
+  def _create_ref_dataloader(self, base_dir, ann_file, subset):
+    self._set_seed()
+    transforms = DetectionPresetTrain("hflip") if subset == "train" else DetectionPresetEval()
+    return iter(get_openimages(ann_file.stem, base_dir, subset, transforms))
+
+  def _create_tinygrad_dataloader(self, base_dir, ann_file, subset, batch_size=1, seed=42):
+    return iter(batch_load_retinanet(COCO(ann_file), subset == "validation", base_dir, batch_size=batch_size, shuffle=False, seed=seed))
+
+  def _normalize_img(self, img):
+    return ((img / 255.0) - np.array(self.img_mean)) / np.array(self.img_std)
+
+  def test_training_set(self):
+    base_dir, ann_file = self._create_samples((subset:="train"))
+    transform = GeneralizedRCNNTransform(self.img_size, self.img_mean, self.img_std)
+    anchors = torch.ones((120087, 4))
+
+    tinygrad_dataloader = self._create_tinygrad_dataloader(base_dir, ann_file, subset)
+    ref_dataloader = self._create_ref_dataloader(base_dir, ann_file, subset)
+
+    for ((tinygrad_img, tinygrad_boxes, tinygrad_labels, _, _, _), (ref_img, ref_tgt)) in zip(tinygrad_dataloader, ref_dataloader):
+      ref_img, ref_tgt = transform(ref_img.unsqueeze(0), [ref_tgt])
+      ref_tgt = postprocess_targets(ref_tgt, anchors.unsqueeze(0))
+      ref_boxes, ref_labels = ref_tgt[0]["boxes"], ref_tgt[0]["labels"]
+
+      np.testing.assert_allclose(self._normalize_img(tinygrad_img.numpy()), ref_img.tensors.transpose(1, 3).numpy(), rtol=1e-6)
+      np.testing.assert_equal(tinygrad_boxes[0].numpy(), ref_boxes.numpy())
+      np.testing.assert_equal(tinygrad_labels[0].numpy(), ref_labels.numpy())
+
+  def test_validation_set(self):
+    base_dir, ann_file = self._create_samples((subset:="validation"))
+    transform = GeneralizedRCNNTransform(self.img_size, self.img_mean, self.img_std)
+
+    tinygrad_dataloader = self._create_tinygrad_dataloader(base_dir, ann_file, subset)
+    ref_dataloader = self._create_ref_dataloader(base_dir, ann_file, "val")
+
+    for ((tinygrad_img, _, _, _), (ref_img, _)) in zip(tinygrad_dataloader, ref_dataloader):
+      ref_img, _ = transform(ref_img.unsqueeze(0))
+      np.testing.assert_allclose(self._normalize_img(tinygrad_img.numpy()), ref_img.tensors.transpose(1, 3).numpy(), rtol=1e-6)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_dev_var.py

@@ -0,0 +1,26 @@
+import subprocess, unittest, os, sys
+from tinygrad.device import Device
+
+class TestTinygradSlow(unittest.TestCase):
+  def test_env_overwrite_default_device(self):
+    subprocess.run([f'DEV={Device.DEFAULT} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                    shell=True, check=True)
+
+    if Device.DEFAULT != "CPU":
+      # setting device via DEV
+      subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+      subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+      subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+
+class TestRunAsModule(unittest.TestCase):
+  def test_module_runs(self):
+    p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+      env={**os.environ, "DEBUG": "1"}, timeout=40,)
+    out = (p.stdout + p.stderr).decode()
+    self.assertEqual(p.returncode, 0, msg=out)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_example.py

@@ -0,0 +1,74 @@
+import unittest, sys
+from tinygrad import Device
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import getenv, OSX
+
+def multidevice_test(fxn):
+  exclude_devices = getenv("EXCLUDE_DEVICES", "").split(",")
+  def ret(self):
+    for device in Device._devices:
+      # broken on OSX USB AMD, why?
+      if device in ["DISK", "NPY", "FAKE", "DSP", "NULL"] or (OSX and device in ["AMD"]): continue
+      if sys.stdout.isatty(): print(device)
+      if device in exclude_devices:
+        if sys.stdout.isatty(): print(f"WARNING: {device} test is excluded")
+        continue
+      with self.subTest(device=device):
+        try:
+          Device[device]
+        except Exception:
+          if sys.stdout.isatty(): print(f"WARNING: {device} test isn't running")
+          continue
+        fxn(self, device)
+  return ret
+
+class TestExample(unittest.TestCase):
+  @multidevice_test
+  def test_convert_to_cpu(self, device):
+    a = Tensor([[1,2],[3,4]], device=device)
+    assert a.numpy().shape == (2,2)
+    b = a.to("CPU")
+    assert b.numpy().shape == (2,2)
+
+  @multidevice_test
+  def test_2_plus_3(self, device):
+    a = Tensor([2], device=device)
+    b = Tensor([3], device=device)
+    result = a + b
+    print(f"{a.numpy()} + {b.numpy()} = {result.numpy()}")
+    assert result.numpy()[0] == 5.
+
+  @multidevice_test
+  def test_example_readme(self, device):
+    x = Tensor.eye(3).clone().to(device)
+    y = Tensor([[2.0,0,-2.0]], device=device)
+    z = y.matmul(x).sum()
+    z.backward()
+
+    x.grad.numpy()  # dz/dx
+    y.grad.numpy()  # dz/dy
+
+    assert x.grad.device == device
+    assert y.grad.device == device
+
+  @multidevice_test
+  def test_example_matmul(self, device):
+    try:
+      Device[device]
+    except Exception:
+      print(f"WARNING: {device} test isn't running")
+      return
+
+    x = Tensor.eye(8).clone().to(device)
+    y = Tensor.eye(8).clone().to(device)
+    z = y.matmul(x).sum()
+    z.backward()
+
+    x.grad.numpy()  # dz/dx
+    y.grad.numpy()  # dz/dy
+
+    assert x.grad.device == device
+    assert y.grad.device == device
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_gpu_crash.py

@@ -0,0 +1,138 @@
+# ruff: noqa: F405
+"""Tests for GPU crash scenarios using AMD assembly to trigger invalid operations.
+
+These tests intentionally cause GPU faults to verify error handling.
+Run with: DEV=AMD python -m pytest test/external/external_test_gpu_crash.py -v
+"""
+import unittest, re, importlib
+from tinygrad.device import Device
+from tinygrad.renderer.amd.dsl import s, v, Inst, NULL
+
+RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "s_waitcnt": "s_waitcnt", "s_endpgm": "s_endpgm",
+                   "global_load_b32": "global_load_dword", "global_store_b32": "global_store_dword",
+                   "global_atomic_add_u32": "global_atomic_add", "flat_load_b32": "flat_load_dword",
+                   "flat_store_b32": "flat_store_dword", "flat_atomic_add_u32": "flat_atomic_add", "s_load_b32": "s_load_dword"}
+
+def assemble(code:str, name:str="test", is_cdna:bool=False) -> str:
+  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
+  if is_cdna: kd["accum_offset"] = 8
+  else: kd["wavefront_size32"] = 1
+  return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \
+         "\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel"
+
+@unittest.skipIf(Device.DEFAULT != "AMD", "AMD required")
+class TestGPUCrash(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+    cls.dev = Device["AMD"]
+    cls.compiler = HIPCompiler(cls.dev.arch)
+    cls.is_cdna = cls.dev.target[0] < 10
+    ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins')
+    for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items():
+      setattr(cls, rdna3_name, getattr(ins, cdna3_name if cls.is_cdna else rdna3_name))
+
+  def setUp(self):
+    # Verify device works before each test
+    from tinygrad import Tensor
+    try:
+      t = Tensor([1.0, 2.0], device="AMD").realize()
+      assert (t + 1).numpy().tolist() == [2.0, 3.0]
+    except Exception:
+      self.fail("Device not working before test")
+
+  def _run(self, code: str):
+    from tinygrad.runtime.ops_amd import AMDProgram
+    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code, is_cdna=self.is_cdna)))
+    prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)
+
+  def _run_insts(self, insts: list[Inst]):
+    from test.amd.disasm import disasm
+    self._run("\n".join(disasm(i) for i in insts))
+
+  def _assert_gpu_fault(self, func):
+    """Assert that func raises a RuntimeError indicating a GPU fault (not a setup error)."""
+    with self.assertRaises(RuntimeError) as cm:
+      func()
+    err_msg = str(cm.exception).lower()
+    # Verify it's a GPU fault, not a setup/device initialization error
+    self.assertTrue(
+      re.search(r'fault|hang|timeout|illegal|memviol', err_msg),
+      f"Expected GPU fault error, got: {cm.exception}"
+    )
+
+
+class TestOutOfBoundsMemoryAccess(TestGPUCrash):
+  """Tests for out-of-bounds memory accesses."""
+
+  def test_global_load_null_ptr(self):
+    """Global load from NULL pointer."""
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0),
+             self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_global_store_null_ptr(self):
+    """Global store to NULL pointer."""
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF),
+             self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_global_load_unmapped_high_address(self):
+    """Global load from high unmapped address (0xDEAD00000000)."""
+    insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD),
+             self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_global_store_unmapped_high_address(self):
+    """Global store to high unmapped address."""
+    insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 0x12345678),
+             self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_global_atomic_unmapped(self):
+    """Atomic operation on unmapped memory."""
+    insts = [self.v_mov_b32_e32(v[0], 0xBEEF0000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 1),
+             self.global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+
+class TestSMEMFaults(TestGPUCrash):
+  """Tests for scalar memory (SMEM) faults."""
+
+  def test_smem_load_null(self):
+    """SMEM load from NULL base."""
+    insts = [self.s_mov_b32(s[2], 0), self.s_mov_b32(s[3], 0),
+             self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_smem_load_unmapped(self):
+    """SMEM load from unmapped address."""
+    insts = [self.s_mov_b32(s[2], 0xBEEF0000), self.s_mov_b32(s[3], 0xDEAD),
+             self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+
+class TestFlatMemoryFaults(TestGPUCrash):
+  """Tests for FLAT memory instruction faults."""
+
+  def test_flat_load_null(self):
+    """FLAT load from NULL address."""
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0),
+             self.flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_flat_store_null(self):
+    """FLAT store to NULL address."""
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF),
+             self.flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+  def test_flat_atomic_null(self):
+    """FLAT atomic on NULL address."""
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 1),
+             self.flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    self._assert_gpu_fault(lambda: self._run_insts(insts))
+
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/external/external_test_hcq.py

@@ -0,0 +1,328 @@
+import unittest, ctypes, struct, time, array
+from tinygrad import Device, Tensor, dtypes
+from tinygrad.helpers import to_mv, DEV
+from tinygrad.device import Buffer, BufferSpec
+from tinygrad.engine.realize import get_runtime
+from tinygrad.codegen import to_program
+
+def _time_queue(q, d):
+  st = time.perf_counter()
+  q.signal(d.timeline_signal, d.timeline_value)
+  q.submit(d)
+  d._wait_signal(d.timeline_signal, d.timeline_value)
+  d.timeline_value += 1
+  return time.perf_counter() - st
+
+@unittest.skipUnless(Device.DEFAULT in ["NV", "AMD"], "Runs only on NV or AMD")
+class TestHCQ(unittest.TestCase):
+  @classmethod
+  def setUpClass(self):
+    TestHCQ.d0 = Device[Device.DEFAULT]
+    #TestHCQ.d1: AMDDevice = Device["AMD:1"]
+    TestHCQ.a = Tensor([0.,1.], device=Device.DEFAULT).realize()
+    TestHCQ.b = self.a + 1
+    linear = self.b.schedule_linear()
+    TestHCQ.prg = to_program(linear.src[-1].src[0], TestHCQ.d0.renderer)
+    TestHCQ.runtime = get_runtime(TestHCQ.d0.device, TestHCQ.prg)
+    TestHCQ.b.uop.buffer.allocate()
+    # wow that's a lot of abstraction layers
+    TestHCQ.addr = struct.pack("QQ", TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf)
+    TestHCQ.addr2 = struct.pack("QQ", TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf)
+    TestHCQ.kernargs_off = TestHCQ.runtime.kernargs_offset
+    TestHCQ.kernargs_size = TestHCQ.runtime.kernargs_alloc_size
+    ctypes.memmove(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_off, TestHCQ.addr, len(TestHCQ.addr))
+    ctypes.memmove(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size+TestHCQ.kernargs_off, TestHCQ.addr2, len(TestHCQ.addr2))
+
+    if Device.DEFAULT == "AMD":
+      from tinygrad.runtime.ops_amd import HWQueue, HWPM4Queue
+      TestHCQ.compute_queue = HWPM4Queue
+      TestHCQ.copy_queue = HWQueue
+    elif Device.DEFAULT == "NV":
+      from tinygrad.runtime.ops_nv import HWQueue, HWQueue
+      # nv need to copy constbuffer there as well
+      to_mv(TestHCQ.d0.kernargs_ptr, 0x160).cast('I')[:] = array.array('I', TestHCQ.runtime.constbuffer_0)
+      to_mv(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, 0x160).cast('I')[:] = array.array('I', TestHCQ.runtime.constbuffer_0)
+      TestHCQ.compute_queue = HWQueue
+      TestHCQ.copy_queue = HWQueue
+
+  def setUp(self):
+    TestHCQ.d0.synchronize()
+    TestHCQ.a.uop.buffer.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
+    TestHCQ.b.uop.buffer.copyin(memoryview(bytearray(struct.pack("ff", 0, 0))))
+    TestHCQ.d0.synchronize() # wait for copyins to complete
+
+  def test_run_1000_times_one_submit(self):
+    temp_signal, temp_value = TestHCQ.d0._alloc_signal(value=0), 0
+    q = TestHCQ.compute_queue()
+    for _ in range(1000):
+      q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
+      q.signal(temp_signal, temp_value + 1).wait(temp_signal, temp_value + 1)
+      temp_value += 1
+
+      q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
+      q.signal(temp_signal, temp_value + 1).wait(temp_signal, temp_value + 1)
+      temp_value += 1
+
+    q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    q.submit(TestHCQ.d0)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    TestHCQ.d0.timeline_value += 1
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
+    assert val == 2000.0, f"got val {val}"
+
+  def test_run_1000_times(self):
+    temp_signal = TestHCQ.d0._alloc_signal(value=0)
+    q = TestHCQ.compute_queue()
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
+    q.signal(temp_signal, 2).wait(temp_signal, 2)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.prg.arg.global_size,
+           TestHCQ.prg.arg.local_size)
+    for _ in range(1000):
+      TestHCQ.d0._set_signal(temp_signal, 1)
+      q.submit(TestHCQ.d0)
+      TestHCQ.compute_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
+      TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+      TestHCQ.d0.timeline_value += 1
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
+    assert val == 2000.0, f"got val {val}"
+
+  def test_run_to_3(self):
+    temp_signal = TestHCQ.d0._alloc_signal(value=0)
+    q = TestHCQ.compute_queue()
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
+    q.signal(temp_signal, 1).wait(temp_signal, 1)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
+    q.signal(temp_signal, 2).wait(temp_signal, 2)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
+    q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    TestHCQ.d0.timeline_value += 1
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
+    assert val == 3.0, f"got val {val}"
+
+  def test_update_exec(self):
+    q = TestHCQ.compute_queue()
+    exec_cmd_idx = len(q)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
+    q.update_exec(exec_cmd_idx, (1,1,1), (1,1,1))
+    q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    TestHCQ.d0.timeline_value += 1
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
+    assert val == 1.0, f"got val {val}"
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
+    assert val == 0.0, f"got val {val}, should not be updated"
+
+  @unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind")
+  def test_bind_run(self):
+    temp_signal = TestHCQ.d0._alloc_signal(value=0)
+    q = TestHCQ.compute_queue()
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
+    q.signal(temp_signal, 2).wait(temp_signal, 2)
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.prg.arg.global_size,
+           TestHCQ.prg.arg.local_size)
+    q.bind(TestHCQ.d0)
+    for _ in range(1000):
+      TestHCQ.d0._set_signal(temp_signal, 1)
+      q.submit(TestHCQ.d0)
+      TestHCQ.compute_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
+      TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+      TestHCQ.d0.timeline_value += 1
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
+    assert val == 2000.0, f"got val {val}"
+
+  @unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind")
+  def test_update_exec_binded(self):
+    q = TestHCQ.compute_queue()
+    exec_ptr = q.ptr()
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
+    q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    q.bind(TestHCQ.d0)
+
+    q.update_exec(exec_ptr, (1,1,1), (1,1,1))
+    q.submit(TestHCQ.d0)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    TestHCQ.d0.timeline_value += 1
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
+    assert val == 1.0, f"got val {val}"
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
+    assert val == 0.0, f"got val {val}, should not be updated"
+
+  @unittest.skipIf(DEV.interface.startswith("MOCK"), "Can't handle async update on CPU")
+  def test_wait_signal(self):
+    temp_signal = TestHCQ.d0._alloc_signal(value=0)
+    TestHCQ.compute_queue().wait(temp_signal, value=1).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
+    with self.assertRaises(RuntimeError):
+      TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value, timeout=50)
+    # clean up
+    TestHCQ.d0._set_signal(temp_signal, 1)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value, timeout=100)
+    TestHCQ.d0.timeline_value += 1
+
+  @unittest.skipIf(DEV.interface.startswith("MOCK"), "Can't handle async update on CPU")
+  def test_wait_copy_signal(self):
+    temp_signal = TestHCQ.d0._alloc_signal(value=0)
+    TestHCQ.copy_queue().wait(temp_signal, value=1).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
+    with self.assertRaises(RuntimeError):
+      TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value, timeout=50)
+    # clean up
+    TestHCQ.d0._set_signal(temp_signal, 1)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value, timeout=100)
+    TestHCQ.d0.timeline_value += 1
+
+  def test_run_normal(self):
+    q = TestHCQ.compute_queue()
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
+    q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    TestHCQ.d0.timeline_value += 1
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
+    assert val == 1.0, f"got val {val}"
+
+  def test_submit_empty_queues(self):
+    TestHCQ.compute_queue().submit(TestHCQ.d0)
+    TestHCQ.copy_queue().submit(TestHCQ.d0)
+
+  def test_signal_timeout(self):
+    with self.assertRaises(RuntimeError):
+      TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value, timeout=50)
+      TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value + 122, timeout=50)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1, timeout=50)
+
+  def test_signal(self):
+    new_timeline_value = TestHCQ.d0.timeline_value + 0xff
+    TestHCQ.compute_queue().signal(TestHCQ.d0.timeline_signal, new_timeline_value).submit(TestHCQ.d0)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, new_timeline_value)
+    TestHCQ.d0.timeline_value = new_timeline_value + 1 # update to not break runtime
+
+  def test_copy_signal(self):
+    new_timeline_value = TestHCQ.d0.timeline_value + 0xff
+    TestHCQ.copy_queue().signal(TestHCQ.d0.timeline_signal, new_timeline_value).submit(TestHCQ.d0)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, new_timeline_value)
+    TestHCQ.d0.timeline_value = new_timeline_value + 1 # update to not break runtime
+
+  def test_run_signal(self):
+    q = TestHCQ.compute_queue()
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
+    q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    q.submit(TestHCQ.d0)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    TestHCQ.d0.timeline_value += 1
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
+    assert val == 1.0, f"got val {val}"
+
+  def test_copy_1000_times(self):
+    q = TestHCQ.copy_queue()
+    q.copy(TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf, 8)
+    q.copy(TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf, 8)
+    for _ in range(1000):
+      q.submit(TestHCQ.d0)
+      TestHCQ.copy_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
+      TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+      TestHCQ.d0.timeline_value += 1
+    # confirm the signal didn't exceed the put value
+    with self.assertRaises(RuntimeError):
+      TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value + 1, timeout=50)
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
+    assert val == 0.0, f"got val {val}"
+
+  def test_copy(self):
+    q = TestHCQ.copy_queue()
+    q.copy(TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf, 8)
+    q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    q.submit(TestHCQ.d0)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    TestHCQ.d0.timeline_value += 1
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
+    assert val == 1.0, f"got val {val}"
+
+  @unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind")
+  def test_bind_copy(self):
+    q = TestHCQ.copy_queue()
+    q.copy(TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf, 8)
+    q.copy(TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf, 8)
+    q.bind(TestHCQ.d0)
+    for _ in range(1000):
+      q.submit(TestHCQ.d0)
+      TestHCQ.copy_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
+      TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+      TestHCQ.d0.timeline_value += 1
+    # confirm the signal didn't exceed the put value
+    with self.assertRaises(RuntimeError):
+      TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value + 1, timeout=50)
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
+    assert val == 0.0, f"got val {val}"
+
+  def test_copy_bandwidth(self):
+    # THEORY: the bandwidth is low here because it's only using one SDMA queue. I suspect it's more stable like this at least.
+    SZ = 2_000_000_000
+    a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
+    b = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
+    q = TestHCQ.copy_queue()
+    q.copy(a._buf, b._buf, SZ)
+    et = _time_queue(q, TestHCQ.d0)
+    gb_s = (SZ/1e9)/et
+    print(f"same device copy:  {et*1e3:.2f} ms, {gb_s:.2f} GB/s")
+    assert 0.3 <= gb_s <= 1000
+
+  def test_cross_device_copy_bandwidth(self):
+    SZ = 2_000_000_000
+    b = Buffer(f"{Device.DEFAULT}:1", SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
+    a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
+    TestHCQ.d0._gpu_map(b._buf)
+    q = TestHCQ.copy_queue()
+    q.copy(a._buf, b._buf, SZ)
+    et = _time_queue(q, TestHCQ.d0)
+    gb_s = (SZ/1e9)/et
+    print(f"cross device copy: {et*1e3:.2f} ms, {gb_s:.2f} GB/s")
+    assert 0.3 <= gb_s <= 50
+
+  def test_interleave_compute_and_copy(self):
+    q = TestHCQ.compute_queue()
+    qc = TestHCQ.copy_queue()
+    q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)  # b = [1, 2]
+    q.signal(sig:=TestHCQ.d0._alloc_signal(value=0), value=1)
+    qc.wait(sig, value=1)
+    qc.copy(TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf, 8)
+    qc.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    qc.submit(TestHCQ.d0)
+    time.sleep(0.02) # give it time for the wait to fail
+    q.submit(TestHCQ.d0)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    TestHCQ.d0.timeline_value += 1
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
+    assert val == 1.0, f"got val {val}"
+
+  def test_cross_device_signal(self):
+    d1 = Device[f"{Device.DEFAULT}:1"]
+    q1 = TestHCQ.compute_queue()
+    q2 = TestHCQ.compute_queue()
+    q1.signal(sig:=TestHCQ.d0._alloc_signal(value=0), value=0xfff)
+    q2.wait(sig, value=0xfff)
+    q2.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    q2.submit(TestHCQ.d0)
+    q1.signal(d1.timeline_signal, d1.timeline_value)
+    q1.submit(d1)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+    TestHCQ.d0.timeline_value += 1
+    d1._wait_signal(d1.timeline_signal, d1.timeline_value)
+    d1.timeline_value += 1
+
+  def test_timeline_signal_rollover(self):
+    # NV 64bit, AMD 32bit
+    TestHCQ.d0.timeline_value = (1 << 64) - 20 if Device.DEFAULT == "NV" else (1 << 32) - 20 # close value to reset
+    TestHCQ.compute_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1).submit(TestHCQ.d0)
+    TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1)
+
+    for _ in range(40):
+      q = TestHCQ.compute_queue()
+      q.wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1)
+      q.exec(TestHCQ.runtime, TestHCQ.d0.kernargs_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size)
+      q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
+      TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
+      TestHCQ.d0.timeline_value += 1
+      val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
+      assert val == 1.0, f"got val {val}"
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/external/external_test_hcq_open.py

@@ -0,0 +1,20 @@
+import os
+if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
+
+import unittest, time
+from tinygrad import Device
+
+class TestOpen(unittest.TestCase):
+  def generate_test_open(n):
+    def test(self):
+      dev = Device[Device.DEFAULT]
+      for i in range(10):
+        dev.allocator.alloc(10 << 20)
+        time.sleep(0.5)
+    test.__name__ = f'test_open_{n}'
+    return test
+
+  for i in range(64): locals()[f'test_open_{i}'] = generate_test_open(i)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_hip_compile.py

@@ -0,0 +1,39 @@
+import time, unittest
+from tinygrad.runtime.support.hip_comgr import compile_hip
+from tinygrad import Tensor
+from tinygrad.device import Device
+from tinygrad.schedule import create_schedule
+from tinygrad.codegen.opt.kernel import Kernel
+
+class TestHIPCompileSpeed(unittest.TestCase):
+  @unittest.skipIf(Device.DEFAULT != "HIP", "only run on HIP")
+  def test_hip_compile(self):
+    a, b = Tensor([1,2,3,4,5]), Tensor([1,2,3,4,5])
+    out = a + b
+    lin = Kernel(create_schedule([out.uop])[-1].ast[0])
+    lin.to_program()
+
+    reference = """
+#include <hip/hip_common.h>
+    typedef long unsigned int size_t;
+    extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_local_id(unsigned int);
+    extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_group_id(unsigned int);
+    extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_local_size(unsigned int);
+    extern "C" __attribute__((global))void {name}(int* data0, const int* data1, const int* data2) {{
+      int gidx0 = __ockl_get_group_id(0); /* 5 */
+      int val0 = data1[gidx0];
+      int val1 = data2[gidx0];
+      data0[gidx0] = (val0+val1);
+    }}
+    """
+
+    def time_compile(code):
+      st = time.perf_counter()
+      compile_hip(code)
+      return (time.perf_counter() - st) * 1000
+
+    tinygrad_tm = min([time_compile(Device[Device.DEFAULT].renderer.render(f"test{i}", lin.uops)) for i in range(10)])
+    ref_tm = min([time_compile(reference.format(name=f"test{i}")) for i in range(10)])
+    print(f"tinygrad {tinygrad_tm:6.2f} ms")
+    print(f"reference {ref_tm:6.2f} ms")
+    assert (tinygrad_tm - ref_tm) <= 10

tinygrad_repo/test/external/external_test_hive_reset.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+import subprocess, sys
+from tinygrad.helpers import getenv
+
+LOOPS = getenv("LOOPS", 50)
+BROKEN = getenv("BROKEN", 0)
+ONLY_RESET = getenv("ONLY_RESET", 0)
+
+BROKEN_KERNEL_SCRIPT = """
+from tinygrad.device import Device
+from tinygrad.runtime.ops_amd import AMDProgram, AMDDevice
+from tinygrad.runtime.support.compiler_amd import compile_hip
+dev = Device["AMD"]
+assert isinstance(dev, AMDDevice) and dev.is_am(), "Need AM driver (not KFD)"
+broken_src = '''
+extern "C" __attribute__((global)) void broken(int* dummy) {
+  volatile int* bad_ptr = (volatile int*)0xDEAD00000000ULL;
+  *bad_ptr = 0x42;
+}
+'''
+broken_lib = compile_hip(broken_src, dev.arch)
+broken_prg = AMDProgram(dev, "broken", broken_lib)
+buf = dev.allocator.alloc(64)
+try:
+  broken_prg(buf, global_size=(1,1,1), local_size=(1,1,1), wait=True)
+  print("  ERROR: Kernel did not fault!")
+except RuntimeError as e:
+  print(f"  Got expected error: {e}")
+"""
+
+for i in range(LOOPS):
+  print(f"=== Running hive_reset.py ({i+1}/{LOOPS}) ===")
+  subprocess.run([sys.executable, "extra/amdpci/hive_reset.py"], check=True)
+  print("=== hive_reset complete ===")
+
+  if BROKEN:
+    print(f"=== Running broken kernel ({i+1}/{LOOPS}) ===")
+    ret = subprocess.run([sys.executable, "-c", BROKEN_KERNEL_SCRIPT])
+    print(f"=== broken kernel exited with code {ret.returncode} ===")
+  elif not ONLY_RESET:
+    print(f"=== Running test_tiny.py ({i+1}/{LOOPS}) ===")
+    ret = subprocess.run([sys.executable, "test/test_tiny.py", "TestTiny.test_plus"])
+    print(f"=== test_tiny.py exited with code {ret.returncode} ===")

tinygrad_repo/test/external/external_test_image.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python
+import os
+import unittest
+import numpy as np
+if 'IMAGE' not in os.environ:
+  os.environ['IMAGE'] = '2'
+os.environ['CL'] = '1'
+os.environ['OPT'] = '2'
+from tinygrad.tensor import Tensor
+from tinygrad.nn import Conv2d
+
+class TestImage(unittest.TestCase):
+  def test_create_image(self):
+    t = Tensor.ones(128, 128, 1)
+    t = t.reshape(128, 32, 4) + 3
+    t.realize()
+    np.testing.assert_array_equal(t.numpy(), np.ones((128,32,4))*4)
+
+  def test_sum_image(self):
+    t1 = Tensor.ones(16, 16, 1).reshape(16, 4, 4) + 3
+    t1.realize()
+    t1 = t1.sum()
+    t1.realize()
+    assert t1.numpy() == 16*4*4*4, f"got {t1.numpy()}"
+
+  def test_add_image(self):
+    t1 = Tensor.ones(16, 16, 1).reshape(16, 4, 4) + 3
+    t2 = Tensor.ones(16, 16, 1).reshape(16, 4, 4) + 4
+    t1.realize()
+    t2.realize()
+    t3 = t1 + t2
+    t3.realize()
+    np.testing.assert_array_equal(t3.numpy(), np.ones((16,4,4))*9)
+
+  def test_padded_conv(self):
+    bs, in_chans, out_chans = 1,12,32
+    tiny_conv = Conv2d(in_chans, out_chans, 3, bias=None, padding=1)
+    tiny_dat = Tensor.ones(bs, 12, 64, 128)
+    tiny_conv(tiny_dat).realize()
+
+  def test_op_conv(self):
+    bs, in_chans, out_chans = 1,12,32
+    tiny_conv = Conv2d(in_chans, out_chans, 3, bias=None, padding=1)
+    tiny_dconv = Conv2d(out_chans, out_chans, 1, bias=None, padding=0)
+    tiny_dat = Tensor.ones(bs, 12, 64, 128)
+    p2 = tiny_conv(tiny_dat).relu()
+    p2 = tiny_dconv(p2)
+    p2.realize()
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_jit_on_models.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python
+import unittest
+import numpy as np
+from tinygrad import Tensor, dtypes
+from tinygrad.engine.jit import TinyJit
+from test.helpers import derandomize_model
+
+from examples.llama import Transformer
+
+def helper_test_jitted_correctness(gen, train, train_jit):
+  nojit = train(*gen()).numpy()
+  for _ in range(5): jit = train_jit(*gen()).numpy()
+  np.testing.assert_allclose(nojit, jit, rtol=1e-3, atol=1e-5)
+
+class TestJittedModels(unittest.TestCase):
+  def test_jitted_tiny_llama(self):
+    old_float = dtypes.default_float
+    dtypes.default_float = dtypes.float16
+
+    args_tiny = {"dim": 1024, "hidden_dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
+    model = Transformer(**args_tiny)
+    derandomize_model(model)
+    def test(t): return model(t, 0).realize()
+
+    @TinyJit
+    def test_jit(t): return model(t, 0).realize()
+    helper_test_jitted_correctness(lambda: (Tensor([[1,]]),), test, test_jit)
+    dtypes.default_float = old_float
+
+  def test_jitted_stable_diffusion(self):
+    from examples.stable_diffusion import UNetModel, unet_params
+    model = UNetModel(**unet_params)
+    derandomize_model(model)
+    def test(t, t2): return model(t, 801, t2).realize()
+
+    @TinyJit
+    def test_jit(t, t2): return model(t, 801, t2).realize()
+    helper_test_jitted_correctness(lambda: (Tensor.randn(1, 4, 16, 16),Tensor.randn(1, 77, 768)), test, test_jit)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/external/external_test_keccak.py

@@ -0,0 +1,31 @@
+import unittest, zipfile, re
+from tinygrad import Tensor
+from tinygrad.helpers import fetch, tqdm
+
+SHA3_URL = "https://csrc.nist.gov/CSRC/media/Projects/Cryptographic-Algorithm-Validation-Program/documents/sha3/sha-3bytetestvectors.zip"
+SHAKE_URL = "https://csrc.nist.gov/CSRC/media/Projects/Cryptographic-Algorithm-Validation-Program/documents/sha3/shakebytetestvectors.zip"
+
+class TestExternalKeccak(unittest.TestCase):
+  def test_sha3_224(self): self.check_nist_vectors(SHA3_URL, ["SHA3_224LongMsg.rsp", "SHA3_224ShortMsg.rsp"], "sha3_224")
+  def test_sha3_256(self): self.check_nist_vectors(SHA3_URL, ["SHA3_256LongMsg.rsp", "SHA3_256ShortMsg.rsp"], "sha3_256")
+  def test_shake_128(self): self.check_nist_vectors(SHAKE_URL, ["SHAKE128LongMsg.rsp", "SHAKE128ShortMsg.rsp"], "shake_128")
+
+  def check_nist_vectors(self, url: str, filenames: list[str], preset: str):
+    pattern = r"Len\s*=\s*(?P<Len>\d+)\s+Msg\s*=\s*(?P<Msg>[0-9a-fA-F\s]+)\s+(MD|Output)\s*=\s*(?P<Output>[0-9a-fA-F]+)"
+    vecs_zip = fetch(url)
+
+    for filename in filenames:
+      vecs = zipfile.ZipFile(vecs_zip).open(filename).read().decode()
+
+      vectors = [ (l, bytes.fromhex(match["Msg"].lower()), bytes.fromhex(match["Output"].lower()))
+        for match in re.finditer(pattern, vecs) if (l:=int(match["Len"])) < 8192 ]
+
+      self.assertTrue(len(vectors) > 0)
+
+      print("file", filename)
+      for data_len, data, output in tqdm(vectors):
+        tinyout = bytes(Tensor(data[:data_len//8]).keccak(preset).data())
+        self.assertEqual(tinyout, output)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_llama3_ff.py

@@ -0,0 +1,12 @@
+#!/usr/bin/env python3
+from tinygrad import Tensor, TinyJit, nn
+from extra.models.llama import FeedForward
+
+if __name__ == "__main__":
+  model = FeedForward(4096, 14336)
+  for x in nn.state.get_parameters(model): x.replace(x.half()).realize()
+  jrun = TinyJit(model)
+  for i in range(5):
+    print(f"*** run {i}")
+    jrun(Tensor.rand(1, 4096))
+

tinygrad_repo/test/external/external_test_llama3_layer.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python3
+from tinygrad import Tensor, TinyJit, nn, dtypes
+from tinygrad.helpers import getenv
+from extra.models.llama import TransformerBlock, precompute_freqs_cis
+
+BS = getenv("BS", 1)
+SEQLEN = getenv("SEQLEN", 128)
+
+# DEFAULT_FLOAT=bfloat16 SEQLEN=8192 ASM_GEMM=1 HK_FLASH_ATTENTION=1 DEV=NULL:HIP:gfx950 DEBUG=2 VIZ=1 PYTHONPATH="."
+# python test/external/external_test_llama3_layer.py
+
+if __name__ == "__main__":
+  dim, hidden_dim, n_heads, n_kv_heads, norm_eps = 4096, 14336, 32, 8, 1e-5
+  layer = TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, max_context=0)
+  for x in nn.state.get_parameters(layer): x.replace(x.cast(dtypes.default_float)).realize()
+
+  freqs_cis = precompute_freqs_cis(dim // n_heads, SEQLEN, theta=500000.0).contiguous().realize()
+
+  @TinyJit
+  def run(t): return layer(t, 0, freqs_cis, None)
+
+  for i in range(5):
+    print(f"*** run {i}")
+    run(Tensor.rand(BS, SEQLEN, dim, dtype=dtypes.default_float).realize())

tinygrad_repo/test/external/external_test_lm_head.py

@@ -0,0 +1,10 @@
+from tinygrad import Tensor, nn
+
+if __name__ == "__main__":
+  vocab_size = 50257
+  n_embd = 768
+  lm_head = nn.Linear(n_embd, vocab_size, bias=False)
+  bs = 4
+  seq_len = 1024
+  x = Tensor.rand(bs, seq_len, n_embd)
+  ret = lm_head(x).realize()

tinygrad_repo/test/external/external_test_losses.py

@@ -0,0 +1,40 @@
+from tinygrad import Tensor
+from test.external.mlperf_retinanet.focal_loss import sigmoid_focal_loss as ref_sigmoid_focal_loss
+from test.external.mlperf_unet3d.dice import DiceCELoss
+from examples.mlperf.losses import dice_ce_loss, sigmoid_focal_loss, l1_loss
+
+import numpy as np
+import torch
+import unittest
+
+class ExternalTestLosses(unittest.TestCase):
+  def setUp(self):
+    np.random.seed(42)
+
+  def _assert_loss(self, pred, tgt, tinygrad_metrics, ref_metrics, rtol=1e-07, atol=0, **kwargs):
+    tinygrad_metrics_res = tinygrad_metrics(Tensor(pred), Tensor(tgt), **kwargs)
+    ref_metrics_res = ref_metrics(torch.from_numpy(pred), torch.from_numpy(tgt), **kwargs)
+    np.testing.assert_allclose(tinygrad_metrics_res.numpy(), ref_metrics_res.numpy(), rtol=rtol, atol=atol)
+
+  def test_dice_ce_loss(self):
+    pred, label = np.random.rand(1, 3, 128, 128, 128).astype(np.float32), np.ones((1, 1, 128, 128, 128)).astype(np.uint8)
+    tinygrad_metrics_res, ref_metrics_res = dice_ce_loss, DiceCELoss(True, True, "NCDHW", False)
+    self._assert_loss(pred, label, tinygrad_metrics_res, ref_metrics_res, atol=1e-4)
+
+  def test_sigmoid_focal_loss(self):
+    def _apply_logit(p): return np.log(p / (1 - p))
+    pred, tgt = _apply_logit(np.random.rand(5,2).astype(np.float32)), np.random.randint(0, 2, (5, 2)).astype(np.float32)
+    for reduction in ["mean", "sum", "none"]:
+      for alpha, gamma in zip([-1, 0.58], [0, 2]):
+        self._assert_loss(pred, tgt, sigmoid_focal_loss, ref_sigmoid_focal_loss, rtol=1e-4, alpha=alpha, gamma=gamma, reduction=reduction)
+
+  def test_l1_loss(self):
+    N, C, H, W = 3, 4, 5, 6
+    shapes = ((N, C), (N, C, H), (N, C, H, W))
+    for reduction in ["mean", "sum", "none"]:
+      for shape in shapes:
+        pred, tgt = np.random.randint(shape).astype(np.float32), np.random.randint(shape)
+        self._assert_loss(pred, tgt, l1_loss, torch.nn.functional.l1_loss, reduction=reduction)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_mamba.py

@@ -0,0 +1,24 @@
+import unittest
+from test.helpers import slow
+from examples.mamba import Mamba, generate
+from transformers import AutoTokenizer
+
+PROMPT = 'Why is gravity '
+TOKENIZER = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
+
+@slow
+class TestMamba(unittest.TestCase):
+  def test_mamba_130M(self):
+    OUT_130M = '''Why is gravity \nnot a good idea?\n\nA:'''
+    model = Mamba.from_pretrained('130m')
+    tinyoutput = generate(model, TOKENIZER, PROMPT, n_tokens_to_gen=10)
+    self.assertEqual(OUT_130M, tinyoutput)
+    del model
+  def test_mamba_370M(self):
+    OUT_370M = '''Why is gravity \nso important?\nBecause it's the only'''
+    model = Mamba.from_pretrained('370m')
+    tinyoutput = generate(model, TOKENIZER, PROMPT, n_tokens_to_gen=10)
+    self.assertEqual(OUT_370M, tinyoutput)
+    del model
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_metrics.py

@@ -0,0 +1,42 @@
+from tinygrad import Tensor
+from test.external.mlperf_unet3d.dice import DiceScore
+from examples.mlperf.metrics import dice_score, log_perplexity
+
+import numpy as np
+import torch
+import unittest, math
+
+class ExternalTestMetrics(unittest.TestCase):
+  def _test_metrics(self, tinygrad_metrics, orig_metrics, pred, label, atol=1e-8, rtol=1e-7):
+    tinygrad_metrics_res = tinygrad_metrics(Tensor(pred), Tensor(label)).squeeze().numpy()
+    orig_metrics_res = orig_metrics(torch.from_numpy(pred), torch.from_numpy(label)).numpy()
+    np.testing.assert_allclose(tinygrad_metrics_res, orig_metrics_res, atol=atol, rtol=rtol)
+
+  def test_dice(self):
+    pred, label = np.random.rand(1, 3, 128, 128, 128).astype(np.float32), np.ones((1, 1, 128, 128, 128)).astype(np.uint8)
+    self._test_metrics(dice_score, DiceScore(), pred, label)
+
+  def test_log_perplexity(self):
+    # equally likely
+    np.testing.assert_allclose(log_perplexity(Tensor([[[1.0, 1, 1, 1]]]), Tensor([[2]])).numpy(), math.log(4))
+    np.testing.assert_allclose(log_perplexity(Tensor([[[1.0]*256]*32]), Tensor([[2]*32])).numpy(), math.log(256), rtol=1e-6)
+    # pretty correct and incorrect
+    np.testing.assert_allclose(log_perplexity(Tensor([[[10000., 0, 0, 0]]]), Tensor([[0]])).numpy(), 0)
+    np.testing.assert_allclose(log_perplexity(Tensor([[[0.0, 10000, 10000, 10000]]]), Tensor([[0]])).numpy(), 10000, rtol=1e-3)
+    # higher logit -> lower loss
+    x = Tensor([[[4.0, 3, 2, 1]]])
+    for i in range(x.numel()-1): self.assertLess(log_perplexity(x, Tensor([[i]])).item(), log_perplexity(x, Tensor([[i+1]])).item())
+    # torch eval examples
+    np.testing.assert_allclose(
+      log_perplexity(Tensor([[[0.3659, 0.7025, 0.3104], [0.0097, 0.6577, 0.1947]]]), Tensor([[2, 1]])).exp().numpy(),
+      2.7593, rtol=1e-5)
+    np.testing.assert_allclose(
+      log_perplexity(Tensor([[[0.3, 0.7, 0.3, 0.1], [0.5, 0.4, 0.1, 0.4],[0.1, 0.1, 0.2, 0.5]],
+                             [[0.1, 0.6, 0.1, 0.5], [0.3, 0.7, 0.3, 0.4], [0.3, 0.7, 0.3, 0.4]]]), Tensor([[2, 1, 3],  [1, 0, 1]])).exp().numpy(),
+      3.6216, rtol=1e-5)
+    np.testing.assert_allclose(
+      log_perplexity(Tensor([[[0.3659, 0.7025, 0.3104], [0.0097, 0.6577, 0.1947]]]), Tensor([[2, 1]]), ignore_index=1).exp().numpy(),
+      3.5372, rtol=1e-4)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_mnist_data_select.py

@@ -0,0 +1,8 @@
+#!/bin/bash
+from tinygrad import Tensor
+from extra.datasets import fetch_mnist
+
+if __name__ == "__main__":
+  X_train, Y_train, X_test, Y_test = fetch_mnist(tensors=True)
+  samples = Tensor.randint(512, high=X_train.shape[0])
+  select = X_train[samples].realize()

tinygrad_repo/test/external/external_test_onnx_backend.py

@@ -0,0 +1,210 @@
+import unittest
+from typing import Any, Tuple
+from onnx.backend.base import Backend, BackendRep
+import onnx.backend.test
+import numpy as np
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.helpers import getenv, OSX
+from tinygrad.nn.onnx import OnnxRunner
+
+# pip3 install tabulate
+pytest_plugins = 'onnx.backend.test.report',
+
+class TinygradModel(BackendRep):
+  def __init__(self, run_onnx, input_names):
+    super().__init__()
+    self.fxn = run_onnx
+    self.input_names = input_names
+
+  def run(self, inputs: Any, **kwargs: Any) -> Tuple[Any, ...]:
+    real_inputs = dict(zip(self.input_names, inputs))
+    ret = self.fxn(real_inputs, debug=2)
+    return tuple(x.numpy() if isinstance(x, Tensor) else [i.numpy() for i in x] if isinstance(x, list) else np.array(x) for x in ret.values())
+
+class TinygradBackend(Backend):
+  @classmethod
+  def prepare(cls, model: onnx.ModelProto, device):
+    input_all = [x.name for x in model.graph.input]
+    input_initializer = [x.name for x in model.graph.initializer]
+    net_feed_input = [x for x in input_all if x not in input_initializer]
+    print("prepare", cls, device, net_feed_input)
+    model = Tensor(model.SerializeToString(), device="PYTHON")
+    run_onnx = OnnxRunner(model)
+    return TinygradModel(run_onnx, net_feed_input)
+
+  @classmethod
+  def supports_device(cls, device: str) -> bool:
+    # NOTE: this is onnx CPU
+    return device == "CPU"
+
+backend_test = onnx.backend.test.BackendTest(TinygradBackend, __name__)
+
+# BUG: buggy onnx tests
+backend_test.exclude('test_adam_multiple_cpu')
+
+# BUG: ORT fails these with runtime error
+backend_test.exclude('test_PReLU_1d_multiparam_cpu')
+backend_test.exclude('test_PReLU_2d_multiparam_cpu')
+backend_test.exclude('test_PReLU_3d_multiparam_cpu')
+
+# BUG: we don't match ORT here due to some div inaccuracy with floats
+backend_test.exclude('test_dynamicquantizelinear_cpu')
+backend_test.exclude('test_dynamicquantizelinear_expanded_cpu')
+
+# BUG: ORT fails these with numerical error but we match ORT numerically
+# see: https://onnx.ai/backend-scoreboard/onnxruntime_details_stable.html
+# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_qlinearmatmul_2D_int8_float16
+backend_test.exclude('test_qlinearmatmul_2D_int8_float16_cpu')
+# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_qlinearmatmul_3D_int8_float16
+backend_test.exclude('test_qlinearmatmul_3D_int8_float16_cpu')
+# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_qlinearmatmul_2D_int8_float32
+backend_test.exclude('test_qlinearmatmul_2D_int8_float32_cpu')
+# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_qlinearmatmul_3D_int8_float32
+backend_test.exclude('test_qlinearmatmul_3D_int8_float32_cpu')
+# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_maxunpool_export_with_output_shape
+backend_test.exclude('test_maxunpool_export_with_output_shape_cpu')
+# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_resize_downsample_scales_linear_align_corners
+backend_test.exclude('test_resize_downsample_scales_linear_align_corners_cpu')
+# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_resize_downsample_scales_cubic_align_corners
+backend_test.exclude('test_resize_downsample_scales_cubic_align_corners_cpu')
+
+# about different dtypes
+if dtypes.float64 not in Device[Device.DEFAULT].renderer.supported_dtypes():
+  backend_test.exclude('float64')
+  backend_test.exclude('DOUBLE')
+  # these have float64 inputs
+  backend_test.exclude('test_eyelike_with_dtype_cpu')
+  backend_test.exclude('test_reduce_log_sum_exp*')
+  backend_test.exclude('test_operator_add*')
+  backend_test.exclude('test_einsum_*')
+  backend_test.exclude('test_cumsum_*')
+
+if dtypes.float16 not in Device[Device.DEFAULT].renderer.supported_dtypes():
+  backend_test.exclude('float16')
+  backend_test.exclude('FLOAT16')
+
+# dtype cast
+backend_test.exclude('STRING')
+backend_test.exclude('FLOAT8')
+backend_test.exclude('INT4')
+backend_test.exclude('UINT4')
+backend_test.exclude('BFLOAT16')  # not supported in numpy
+backend_test.exclude('FLOAT4E2M1')
+
+backend_test.exclude('test_dequantizelinear_int4_cpu')
+backend_test.exclude('test_dequantizelinear_uint4_cpu')
+backend_test.exclude('test_quantizelinear_int4_cpu')
+backend_test.exclude('test_quantizelinear_uint4_cpu')
+
+# no support for FLOAT8
+backend_test.exclude('test_quantizelinear_e4m3fn_cpu')
+backend_test.exclude('test_quantizelinear_e5m2_cpu')
+backend_test.exclude('test_quantizelinear_e4m3fn_cpu')
+backend_test.exclude('test_quantizelinear_e5m2_cpu')
+backend_test.exclude('test_quantizelinear_float4e2m1_cpu')
+backend_test.exclude('test_dequantizelinear_e4m3fn_cpu')
+backend_test.exclude('test_dequantizelinear_e4m3fn_zero_point_cpu')
+backend_test.exclude('test_dequantizelinear_e4m3fn_float16_cpu')
+backend_test.exclude('test_dequantizelinear_e5m2_cpu')
+backend_test.exclude('test_dequantizelinear_float4e2m1_cpu')
+
+# we don't support indexes
+
+# no support for int pow
+backend_test.exclude('test_pow_types_int32_int32_cpu')
+backend_test.exclude('test_pow_types_int64_int64_cpu')
+
+# no boolean ops (2d, 3d, 4d)
+backend_test.exclude('test_bitshift_*')
+
+# no string ops
+backend_test.exclude('string')
+backend_test.exclude('test_strnorm_*')
+backend_test.exclude('test_regex_*')
+
+# no rnn
+backend_test.exclude('test_gru_*')
+backend_test.exclude('test_rnn_*')
+backend_test.exclude('test_lstm_*')
+backend_test.exclude('test_simple_rnn_*')
+
+# no control flow
+# control flow uses AttributeProto.GRAPH
+backend_test.exclude('test_loop*')
+backend_test.exclude('test_range_float_type_positive_delta_expanded_cpu') # requires loop
+backend_test.exclude('test_affine_grid_2d_align_corners_expanded_cpu')
+backend_test.exclude('test_affine_grid_2d_expanded_cpu')
+backend_test.exclude('test_affine_grid_3d_align_corners_expanded_cpu')
+backend_test.exclude('test_affine_grid_3d_expanded_cpu')
+backend_test.exclude('test_range_int32_type_negative_delta_expanded_cpu')
+
+# unsupported (strange) ops
+backend_test.exclude('test_bernoulli_*')
+backend_test.exclude('test_det_*')
+backend_test.exclude('test_col2im_*')
+backend_test.exclude('test_gridsample_*')
+backend_test.exclude('test_dft_*')
+backend_test.exclude('test_unique_*')
+backend_test.exclude('test_sequence_*')
+backend_test.exclude('test_nonmaxsuppression_*')
+backend_test.exclude('test_reversesequence_*')
+backend_test.exclude('test_roialign_*')
+backend_test.exclude('test_tfidfvectorizer_*')
+backend_test.exclude('test_stft_*')
+backend_test.exclude('test_melweightmatrix_*')
+
+# more strange ops
+backend_test.exclude('test_basic_deform_conv_*')
+backend_test.exclude('test_deform_conv_*')
+backend_test.exclude('test_lppool_*')
+backend_test.exclude('test_scan_*')
+backend_test.exclude('test_split_to_sequence_*')
+backend_test.exclude('test_ai_onnx_ml_tree_ensemble_*') # https://github.com/onnx/onnx/blob/main/onnx/reference/ops/aionnxml/op_tree_ensemble.py#L121
+
+backend_test.exclude('test_attention_4d_diff_heads_mask4d_padded_kv_cpu')  # needs nonpad_kv_seqlen handling
+backend_test.exclude('test_attention_4d_fp16_cpu')  # fp16 numerical issues
+backend_test.exclude('test_attention_4d_fp16_expanded_cpu')  # fp16 numerical issues
+backend_test.exclude('test_attention_4d_gqa_with_past_and_present_fp16_cpu')  # fp16 numerical issues
+backend_test.exclude('test_attention_4d_gqa_with_past_and_present_fp16_expanded_cpu')  # fp16 numerical issues
+
+
+# rest of the failing tests
+backend_test.exclude('test_resize_tf_crop_and_resize_cpu') # tf_crop_and_resize not implemented
+backend_test.exclude('test_resize_tf_crop_and_resize_axes_2_3_cpu') # tf_crop_and_resize not implemented
+backend_test.exclude('test_resize_tf_crop_and_resize_axes_3_2_cpu') # tf_crop_and_resize not implemented
+backend_test.exclude('test_resize_tf_crop_and_resize_extrapolation_value_cpu') # tf_crop_and_resize value not implemented
+backend_test.exclude('test_resize_downsample_scales_linear_antialias_cpu') # antialias not implemented
+backend_test.exclude('test_resize_downsample_sizes_linear_antialias_cpu') # antialias not implemented
+backend_test.exclude('test_resize_downsample_scales_cubic_antialias_cpu') # antialias not implemented
+backend_test.exclude('test_resize_downsample_sizes_cubic_antialias_cpu') # antialias not implemented
+backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_value_only_mapping_cpu') # bad data type string
+backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad data type string
+backend_test.exclude('test_if_opt_cpu') # ValueError: 13 is not a valid AttributeType
+backend_test.exclude('test_if_seq_cpu') # NotImplementedError: op='SequenceConstruct' is not supported
+
+if Device.DEFAULT == "METAL" or (OSX and Device.DEFAULT == "CL"):
+  # numerical inaccuracy
+  backend_test.exclude('test_mish_cpu')
+  backend_test.exclude('test_mish_expanded_cpu')
+
+# disable model tests for now since they are slow
+if not getenv("MODELTESTS"):
+  for x in backend_test.test_suite:
+    if 'OnnxBackendRealModelTest' in str(type(x)):
+      backend_test.exclude(str(x).split(" ")[0])
+else:
+  # model tests all pass!
+  backend_test.include('test_resnet50')
+  backend_test.include('test_inception_v1')
+  backend_test.include('test_inception_v2')
+  backend_test.include('test_densenet121')
+  backend_test.include('test_shufflenet')
+  backend_test.include('test_squeezenet')
+  backend_test.include('test_bvlc_alexnet')
+  backend_test.include('test_zfnet512')
+  backend_test.include('test_vgg19')
+
+globals().update(backend_test.enable_report().test_cases)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_onnx_ops.py

@@ -0,0 +1,622 @@
+# inputs, attributes, and outputs for tests are found here:
+# https://github.com/onnx/onnx/blob/main/docs/Operators.md
+# https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md
+
+from typing import Any
+import unittest, onnx, tempfile
+from tinygrad import dtypes, Tensor, Context
+from tinygrad.nn.onnx import OnnxRunner
+import numpy as np
+from extra.onnx_helpers import validate
+from onnx.defs import ONNX_DOMAIN, AI_ONNX_PREVIEW_TRAINING_DOMAIN
+MICROSOFT_CONTRIB_OPS_DOMAIN = "com.microsoft"
+# TODO: remove this once ORT supports 1.18.0
+from onnx.helper import VERSION_TABLE
+VERSION_MAP = {row[0]: row[1:] for row in VERSION_TABLE}
+IR_VERSION, ai_onnx, ai_onnx_ml, ai_onnx_training = VERSION_MAP["1.17.0"]
+
+
+class TestOnnxOps(unittest.TestCase):
+  DOMAIN = None
+  def helper_build_model(self, op:str, inps:dict[str, np.ndarray], opts:dict[str, Any], outs:list[str]):
+    onnx_inputs = [onnx.helper.make_tensor_value_info(name, onnx.helper.np_dtype_to_tensor_dtype(arr.dtype), arr.shape) for name, arr in inps.items()]
+    onnx_outputs = [onnx.helper.make_empty_tensor_value_info(name) for name in outs]
+    nodes = [onnx.helper.make_node(op, list(inps), list(outs), domain=self.DOMAIN, **opts)]
+    graph = onnx.helper.make_graph(nodes, f"test_{op.lower()}", onnx_inputs, onnx_outputs)
+    #model = onnx.helper.make_model(graph, producer_name=f"test_{op.lower()}")
+    # TODO: remove this once ORT supports 1.18.0
+    opset_id = None
+    if type(self).__name__ == "TestMainOnnxOps": opset_id = ai_onnx
+    if type(self).__name__ == "TestTrainingOnnxOps": opset_id = ai_onnx_training
+    if type(self).__name__ == "TestContribOnnxOps": opset_id = 1
+    model = onnx.helper.make_model(graph, producer_name=f"test_{op.lower()}", ir_version=IR_VERSION,
+                                    opset_imports=[onnx.helper.make_opsetid(self.DOMAIN, opset_id)])
+    return model
+
+  def helper_test_single_op(self, op:str, inps:dict[str, np.ndarray], opts:dict[str, Any], outs:list[str], rtol=1e-3, atol=1e-6):
+    model = self.helper_build_model(op, inps, opts, outs)
+    with tempfile.NamedTemporaryFile() as tmp:
+      onnx.save(model, tmp.name)
+      validate(tmp.name, inps, rtol, atol)
+
+class TestMainOnnxOps(TestOnnxOps):
+  DOMAIN = ONNX_DOMAIN
+  def test_reshape(self):
+    inputs = {"in": np.arange(6, dtype=np.float32), "shape": np.array([2,3], dtype=np.int64)}
+    attributes = {}
+    outputs = ["out"]
+    self.helper_test_single_op("Reshape", inputs, attributes, outputs)
+
+  def test_squeeze(self):
+    # axes is None
+    inputs = {"data": np.random.randn(1, 3, 1, 1).astype(np.float32)}
+    attributes = {}
+    outputs = ["squeezed"]
+    self.helper_test_single_op("Squeeze", inputs, attributes, outputs)
+
+  def test_conv(self):
+    # test VALID auto_pad
+    inputs = {
+      "x": np.random.randn(1, 3, 384, 384).astype(np.float32),
+      "w": np.random.randn(1152, 3, 14, 14).astype(np.float32),
+      "b": np.random.randn(1152).astype(np.float32)
+    }
+    attributes = {'auto_pad': 'VALID', 'dilations': (1, 1), 'group': 1, 'kernel_shape': (14, 14), 'strides': (14, 14)}
+    outputs = ["y"]
+    self.helper_test_single_op("Conv", inputs, attributes, outputs, atol=1e-4)
+
+  def test_pad_constant_value_zero(self):
+    from tinygrad.nn.onnx import onnx_ops
+    Pad = onnx_ops["Pad"]
+    x = Tensor.arange(4).reshape(1, 1, 2, 2).float()
+    pads = [0, 0, 1, 1, 0, 0, 1, 1]
+    out = Pad(x, pads, constant_value=0, value=3)
+    expected = x.pad((pads[3], pads[7], pads[2], pads[6], pads[1], pads[5], pads[0], pads[4]), value=0)
+    self.assertEqual(out.tolist(), expected.tolist())
+
+  def test_gather(self):
+    # test const negative indices
+    inputs = {
+      "input": np.random.randn(1, 3, 3).astype(np.float32),
+      "indices": np.array(-2, dtype=np.int64),
+    }
+    attributes = {'axis': 1}
+    outputs = ["y"]
+    self.helper_test_single_op("Gather", inputs, attributes, outputs)
+
+  def test_gather_jit_different_indices(self):
+    # Gather should not assume indices is const when it can change at runtime
+    from tinygrad import TinyJit
+    from tinygrad.nn.onnx import onnx_ops
+    Gather = onnx_ops["Gather"]
+
+    x = Tensor([10, 20, 30, 40, 50])
+    indices_list = [[0, 1], [2, 3], [4, 0]]
+    expected = [[10, 20], [30, 40], [50, 10]]
+
+    # without JIT: correct
+    self.assertEqual([Gather(x, Tensor(idx)).tolist() for idx in indices_list], expected)
+
+    @TinyJit
+    def gather_jit(x, indices): return Gather(x, indices)
+    self.assertEqual([gather_jit(x, Tensor(idx)).tolist() for idx in indices_list], expected)
+
+  def test_gather_jit_const_zero_index(self):
+    # Gather with const index=0 (falsy in Python) should work with JIT cache
+    from tinygrad import TinyJit
+    # Create model: y = Gather(x, 0) + x where 0 is from initializer
+    # The Add ensures there's a kernel to JIT
+    x_input = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, (5,))
+    y_output = onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, (5,))
+    idx_init = onnx.numpy_helper.from_array(np.array(0, dtype=np.int64), name="idx")
+    graph = onnx.helper.make_graph([
+      onnx.helper.make_node("Gather", ["x", "idx"], ["g"], axis=0),
+      onnx.helper.make_node("Add", ["g", "x"], ["y"])],
+      "test_gather_zero", [x_input], [y_output], [idx_init])
+    model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 13)])
+
+    with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as tmp:
+      onnx.save(model, tmp.name)
+      runner = OnnxRunner(tmp.name)
+
+      @TinyJit
+      def run_gather(x): return runner({"x": x})["y"]
+
+      # Run multiple times - JIT capture should use cached index=0 correctly
+      for val in [[10, 20, 30, 40, 50], [100, 200, 300, 400, 500], [1, 2, 3, 4, 5]]:
+        result = run_gather(Tensor(val, dtype=dtypes.float32))
+        np.testing.assert_equal(result.numpy(), np.array(val) + val[0])
+
+  # NOTE: resize OP is sensitive to numerical errors
+  def _test_resize_scales(self, scale_values, **kwargs):
+    for sc in scale_values:
+      for ct_mode in ["half_pixel", "align_corners", "asymmetric", "pytorch_half_pixel", "half_pixel_symmetric"]:
+        with self.subTest(coordinate_transformation_mode=ct_mode, scale=sc, **kwargs):
+          X = np.array([[[[1, 2, 3, 4],
+                          [5, 6, 7, 8],
+                          [9,10,11,12]]]], dtype=np.float32)
+          scales = np.array([1.0, 1.0, sc, sc], dtype=np.float32)
+          inputs = {"X": X, "roi": np.array([], dtype=np.float32), "scales": scales}
+          attributes = {"coordinate_transformation_mode": ct_mode, **kwargs}
+          outputs = ["out"]
+          self.helper_test_single_op("Resize", inputs, attributes, outputs)
+
+  def test_resize_linear_mode(self):
+    self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="linear")
+
+  def test_resize_nearest_mode(self):
+    # excluded 3.5 because some values divide into slight numerical differences, which when rounded gives wrong results
+    self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 20.0], mode="nearest")
+
+  def test_resize_cubic_mode(self):
+    self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="cubic", exclude_outside=1)
+    self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="cubic", exclude_outside=0)
+
+  def _test_if(self, then_value, else_value):
+    then_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, then_value.shape)
+    else_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, else_value.shape)
+
+    then_const_node = onnx.helper.make_node("Constant", inputs=[], outputs=["res"], value=onnx.numpy_helper.from_array(then_value))
+    else_const_node = onnx.helper.make_node("Constant", inputs=[], outputs=["res"], value=onnx.numpy_helper.from_array(else_value))
+
+    then_body = onnx.helper.make_graph([then_const_node], "then_body", [], [then_out])
+    else_body = onnx.helper.make_graph([else_const_node], "else_body", [], [else_out])
+
+    self.helper_test_single_op("If", {"cond": np.array(False).astype(bool)}, {"then_branch": then_body, "else_branch": else_body}, ["res"])
+    self.helper_test_single_op("If", {"cond": np.array(True).astype(bool)}, {"then_branch": then_body, "else_branch": else_body}, ["res"])
+
+  def test_if_different_shapes_broadcastable(self):
+    self._test_if(np.array([[1], [2]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
+
+  def test_if_different_shapes_not_broadcastable(self):
+    self._test_if(np.array([[1, 2, 3], [4, 5, 6]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
+
+  def test_if_jit_different_shapes(self):
+    # When shapes differ, Python selection evaluates condition at graph build time, breaking JIT
+    from tinygrad import TinyJit
+    from tinygrad.engine.jit import JitError
+    # then: x+1 shape (3,), else: x[:2]+1 shape (2,)
+    x_input = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, (3,))
+    then_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, (3,))
+    then_body = onnx.helper.make_graph([
+      onnx.helper.make_node("Constant", [], ["one"], value=onnx.numpy_helper.from_array(np.array(1, dtype=np.float32))),
+      onnx.helper.make_node("Add", ["x", "one"], ["res"])], "then_body", [x_input], [then_out])
+    else_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, (2,))
+    else_body = onnx.helper.make_graph([
+      onnx.helper.make_node("Constant", [], ["starts"], value=onnx.numpy_helper.from_array(np.array([0], dtype=np.int32))),
+      onnx.helper.make_node("Constant", [], ["ends"], value=onnx.numpy_helper.from_array(np.array([2], dtype=np.int32))),
+      onnx.helper.make_node("Constant", [], ["one"], value=onnx.numpy_helper.from_array(np.array(1, dtype=np.float32))),
+      onnx.helper.make_node("Slice", ["x", "starts", "ends"], ["x2"]),
+      onnx.helper.make_node("Add", ["x2", "one"], ["res"])], "else_body", [x_input], [else_out])
+
+    cond_input = onnx.helper.make_tensor_value_info("cond", onnx.TensorProto.BOOL, (1,))
+    main_x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, (3,))
+    graph = onnx.helper.make_graph([onnx.helper.make_node("If", ["cond"], ["res"], then_branch=then_body, else_branch=else_body)],
+      "test", [cond_input, main_x], [onnx.helper.make_empty_tensor_value_info("res")])
+    model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 22)])
+
+    with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as tmp:
+      onnx.save(model, tmp.name)
+      runner = OnnxRunner(tmp.name)
+
+      @TinyJit
+      def run_if(cond, x): return runner({"cond": cond, "x": x})["res"]
+
+      x = Tensor([1.0, 2.0, 3.0])
+      with self.assertRaises(JitError):
+        for _ in range(3):
+          run_if(Tensor([True]), x)
+
+  def test_resize_downsample_scales_linear_align_corners(self):
+    # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-131
+    X = np.array([[[[1, 2, 3, 4], [5, 6, 7, 8]]]], dtype=np.float32)
+    scales = np.array([1.0, 1.0, 0.6, 0.6], dtype=np.float32)
+    inputs = {"X": X, "roi": np.array([], dtype=np.float32), "scales": scales}
+    attributes = {"mode": "linear", "coordinate_transformation_mode": "align_corners"}
+    outputs = ["out"]
+    self.helper_test_single_op("Resize", inputs, attributes, outputs)
+
+  def test_resize_downsample_scales_cubic_align_corners(self):
+    # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-131
+    X = np.array([[[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]]], dtype=np.float32)
+    scales = np.array([1.0, 1.0, 0.8, 0.8], dtype=np.float32)
+    inputs = {"X": X, "roi": np.array([], dtype=np.float32), "scales": scales}
+    attributes = {"mode": "cubic", "coordinate_transformation_mode": "align_corners"}
+    outputs = ["out"]
+    self.helper_test_single_op("Resize", inputs, attributes, outputs)
+
+  def test_maxunpool_export_with_output_shape(self):
+    # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-91
+    xT = np.array([[[[5, 6], [7, 8]]]], dtype=np.float32)
+    xI = np.array([[[[5, 7], [13, 15]]]], dtype=np.int64)
+    output_shape = np.array((1, 1, 5, 5), dtype=np.int64)
+    inputs = {"x": xT, "indices": xI, "output_shape": output_shape}
+    attributes = {"kernel_shape": [2, 2], "strides": [2, 2]}
+    outputs = ["y"]
+    self.helper_test_single_op("MaxUnpool", inputs, attributes, outputs)
+
+  def test_averagepool_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_True(self):
+    # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-13
+    inputs = {"x": np.random.randn(1, 1, 32, 32, 32).astype(np.float32)}
+    attributes = {"kernel_shape": (5, 5, 5), "strides": (3, 3, 3), "dilations": (2, 2, 2), "count_include_pad": 1, "ceil_mode": True}
+    outputs = ["y"]
+    self.helper_test_single_op("AveragePool", inputs, attributes, outputs)
+
+  def test_isinf(self):
+    # https://github.com/onnx/onnx/blob/main/docs/Operators.md#isinf
+    # attributes are int but output expects bool
+    x = np.array([-1.2, np.nan, np.inf, 2.8, -np.inf, np.inf], dtype=np.float32)
+    inputs = {"x": x}
+    attributes = {"detect_negative":1, "detect_positive":1}
+    outputs = ["y"]
+    model = self.helper_build_model("IsInf", inputs, attributes, outputs)
+    runner = OnnxRunner(Tensor(model.SerializeToString(), device="PYTHON"))
+    outputs = runner(inputs)
+    assert outputs["y"].dtype is dtypes.bool
+
+  def test_quantize_linear(self):
+    test_cases = [
+      {"test_case": "round_half_to_even", "qdtype": np.int8, "qzero_point": 0, "x": [-1.5, -0.5, 0.5, 1.5], "scale": 1.0},
+      {"test_case": "round_to_even_before_add_zero_point", "qdtype": np.uint8, "qzero_point": 1, "x": [0.5, 1.5], "scale": 1.0},
+    ]
+    for case in test_cases:
+      with self.subTest(test_case=case["test_case"]):
+        inputs = {
+          "x": np.array([case["x"]], dtype=np.float32),
+          "y_scale": np.array(case["scale"], dtype=np.float32),
+          "y_zero_point": np.array(case["qzero_point"], dtype=case["qdtype"])
+        }
+        self.helper_test_single_op("QuantizeLinear", inputs, {}, ["y"])
+
+  def test_dynamic_quantize_linear(self):
+    test_cases = [
+      {"name": "round_half_to_even", "x": np.array([0, 0.5, 1.5, 255], dtype=np.float32)},
+      {"name": "round_zero_point_half_down_to_even", "x": np.array([-1, 509], dtype=np.float32)},
+      {"name": "round_zero_point_half_up_to_even", "x": np.array([-11, 499], dtype=np.float32)},
+      # other tests from https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-45
+      {"name": "max_adjusted", "x": np.array([-1.0, -2.1, -1.3, -2.5, -3.34, -4.0], dtype=np.float32)},
+      {"name": "min_adjusted", "x": np.array([1, 2.1, 1.3, 2.5, 3.34, 4.0, 1.5, 2.6, 3.9, 4.0, 3.0, 2.345], dtype=np.float32).reshape((3, 4))},
+    ]
+    for case in test_cases:
+      with self.subTest(test_case=case["name"]):
+        self.helper_test_single_op("DynamicQuantizeLinear", {"x": case["x"]}, {}, ["y", "y_scale", "y_zero_point"])
+
+  def test_qlinear_conv(self):
+    for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
+      for b in (np.ones([32], dtype=np.int32), np.zeros([32], dtype=np.int32)):
+        for channel_shape in [(), (32,)]:
+          with self.subTest(dtype=dtype, zero_point=zero_point, channel_shape=channel_shape):
+            dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
+            inputs = {
+              "x": np.random.randint(dtype_min, dtype_max + 1, [1, 3, 224, 224], dtype=dtype),
+              "x_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+              "x_zero_point": np.array(zero_point, dtype=dtype),
+              "w": np.random.randint(dtype_min, dtype_max + 1, [32, 3, 3, 3], dtype=dtype),
+              "w_scale": np.random.uniform(0.01, 0.1, channel_shape).astype(np.float32),
+              "w_zero_point": np.full(channel_shape, zero_point, dtype=dtype),
+              "y_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+              "y_zero_point": np.array(zero_point, dtype=dtype),
+              "b": b
+            }
+            attributes = {'auto_pad': 'NOTSET', 'dilations': (1, 1), 'group': 1, 'kernel_shape': (3, 3), 'pads': (1, 1, 1, 1), 'strides': (2, 2)}
+            outputs = ["out"]
+            self.helper_test_single_op("QLinearConv", inputs, attributes, outputs, atol=1) # occasionally inaccurate
+
+  def test_qlinear_matmul(self):
+    for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
+      with self.subTest(dtype=dtype, zero_point=zero_point):
+        dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
+        inputs = {
+          "A": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
+          "A_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+          "A_zero_point": np.array(zero_point, dtype=dtype),
+          "B": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
+          "B_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+          "B_zero_point": np.array(zero_point, dtype=dtype),
+          "Y_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+          "Y_zero_point": np.array(zero_point, dtype=dtype)
+        }
+        attributes = {}
+        outputs = ["Y"]
+        self.helper_test_single_op("QLinearMatMul", inputs, attributes, outputs)
+
+    for name,val in (("round_half_down_to_even", 1), ("round_half_up_to_even", 3)):
+      with self.subTest(test_case=name, val=val):
+        inputs = {
+          "A": np.array([val], dtype=np.int8),
+          "A_scale": np.array(0.5, dtype=np.float32),
+          "A_zero_point": np.array(0, dtype=np.int8),
+          "B": np.array([1], dtype=np.int8),
+          "B_scale": np.array(1, dtype=np.float32),
+          "B_zero_point": np.array(0, dtype=np.int8),
+          "Y_scale": np.array(1, dtype=np.float32),
+          "Y_zero_point": np.array(0, dtype=np.int8)
+        }
+        attributes = {}
+        outputs = ["Y"]
+        self.helper_test_single_op("QLinearMatMul", inputs, attributes, outputs)
+
+  def _run_qlinearmatmul_test(self, quant_type, dtype, dims):
+    # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-111
+    if dims == 2:
+      a = np.array([[208, 236, 0, 238], [3, 214, 255, 29]])
+      b = np.array([[152, 51, 244], [60, 26, 255], [0, 127, 246], [127, 254, 247]])
+    else:
+      a = np.array([[[208, 236, 0, 238], [3, 214, 255, 29]], [[208, 236, 0, 238], [3, 214, 255, 29]]])
+      b = np.array([[[152, 51, 244], [60, 26, 255], [0, 127, 246], [127, 254, 247]], [[152, 51, 244], [60, 26, 255], [0, 127, 246], [127, 254, 247]]])
+    a_zero_point = np.array([113])
+    b_zero_point = np.array([114])
+    y_zero_point = np.array([118])
+    if quant_type == np.int8:
+      a, b, a_zero_point, b_zero_point, y_zero_point = (x - 127 for x in (a, b, a_zero_point, b_zero_point, y_zero_point))
+    a, b, a_zero_point, b_zero_point, y_zero_point = (x.astype(quant_type) for x in (a, b, a_zero_point, b_zero_point, y_zero_point))
+    inputs = {
+      "a": a, "a_scale": np.array([0.0066], dtype=dtype), "a_zero_point": a_zero_point,
+      "b": b, "b_scale": np.array([0.00705], dtype=dtype), "b_zero_point": b_zero_point,
+      "y_scale": np.array([0.0107], dtype=dtype), "y_zero_point": y_zero_point
+    }
+    self.helper_test_single_op("QLinearMatMul", inputs, {}, ["y"],)
+  def test_qlinearmatmul_2D_int8_float16(self): self._run_qlinearmatmul_test(np.int8, np.float16, 2)
+  def test_qlinearmatmul_3D_int8_float16(self): self._run_qlinearmatmul_test(np.int8, np.float16, 3)
+  def test_qlinearmatmul_2D_int8_float32(self): self._run_qlinearmatmul_test(np.int8, np.float32, 2)
+  def test_qlinearmatmul_3D_int8_float32(self): self._run_qlinearmatmul_test(np.int8, np.float32, 3)
+
+  def test_reduce_l2_half(self):
+    inputs = {"data": np.random.randn(1, 1, 32, 32, 32).astype(np.half)*100}
+    self.helper_test_single_op("ReduceL2", inputs, {}, ["reduced"])
+
+  def test_same_device_as_input(self):
+    from tinygrad.nn.onnx import onnx_ops
+    EyeLike = onnx_ops["EyeLike"]
+    Shape = onnx_ops["Shape"]
+    Compress = onnx_ops["Compress"]
+    with Context(DEV="CPU"):
+      x = Tensor.arange(4, device="PYTHON").reshape(2,2)
+      self.assertEqual(EyeLike(x).device, x.device)
+      self.assertEqual(Shape(x).device, x.device)
+      out = Compress(x, [True, False, True, False])
+      self.assertEqual(out.device, x.device)
+      self.assertEqual(out.tolist(), [0, 2])
+
+class TestTrainingOnnxOps(TestOnnxOps):
+  # NOTE: ORT doesn't actually support training ops on cpu so we test using functions provided by onnx
+  DOMAIN = AI_ONNX_PREVIEW_TRAINING_DOMAIN
+  def _validate_training(self, op:str, onnx_fxn, inps:dict[str, np.ndarray], opts:dict[str, Any], outs:list[str]):
+    model = self.helper_build_model(op, inps, opts, outs)
+    if op == "Momentum": del opts['mode']
+    runner = OnnxRunner(Tensor(model.SerializeToString(), device="PYTHON"))
+    tiny_out = runner(inps)
+    onnx_out = onnx_fxn(**inps, **opts)
+    for (nm, t_out), o_out in  zip(tiny_out.items(), onnx_out):
+      np.testing.assert_allclose(t_out.numpy(), o_out, rtol=1e-6, atol=1e-6, err_msg=f"{nm} failed")
+
+  def test_adagrad_t(self):
+    from onnx.backend.test.case.node.adagrad import apply_adagrad
+    for t in [0, 1, 3, 100]:
+      inputs = {
+        "r": np.array(0.01, dtype=np.float32),
+        "t": np.array(t, dtype=np.int32),
+        "x": np.random.randn(3, 3).astype(np.float32),
+        "g": np.random.randn(3, 3).astype(np.float32),
+        "h": np.random.randn(3, 3).astype(np.float32),
+      }
+      attributes = {"decay_factor": 0.1, "epsilon": 1e-6, "norm_coefficient": 0.01}
+      outputs = ["X_out", "H_out"]
+      self._validate_training("Adagrad", apply_adagrad, inputs, attributes, outputs)
+
+  def test_momentum(self):
+    from onnx.backend.test.case.node.momentum import apply_momentum, apply_nesterov
+    for onnx_fxn, mode in ((apply_momentum, "standard"), (apply_nesterov, "nesterov")):
+      for t in [0, 1, 3, 100]:
+        inputs = {
+          "r": np.array(0.01, dtype=np.float32),
+          "t": np.array(t, dtype=np.int32),
+          "x": np.random.randn(3, 3).astype(np.float32),
+          "g": np.random.randn(3, 3).astype(np.float32),
+          "v": np.random.randn(3, 3).astype(np.float32),
+        }
+        attributes = {"alpha": 0.9, "beta": 0.1, "mode": mode, "norm_coefficient": 0.01}
+        outputs = ["X_out", "V_out"]
+        self._validate_training("Momentum", onnx_fxn, inputs, attributes, outputs)
+
+  def test_adam(self):
+    from onnx.backend.test.case.node.adam import apply_adam
+    for t in [0, 1, 3, 100]:
+      inputs = {
+        "r": np.array(0.01, dtype=np.float32),
+        "t": np.array(t, dtype=np.int32),
+        "x": np.random.randn(3, 3).astype(np.float32),
+        "g": np.random.randn(3, 3).astype(np.float32),
+        "v": np.random.randn(3, 3).astype(np.float32),
+        "h": np.random.randn(3, 3).astype(np.float32),
+      }
+      attributes = { "alpha": 0.9, "beta": 0.999, "epsilon": 1e-8, "norm_coefficient": 0.01, "norm_coefficient_post": 0.02 }
+      outputs = ["X_new", "V_new", "H_new"]
+      self._validate_training("Adam", apply_adam, inputs, attributes, outputs)
+
+class TestContribOnnxOps(TestOnnxOps):
+  DOMAIN = MICROSOFT_CONTRIB_OPS_DOMAIN
+  def test_attention(self):
+    batch_size, seq_len, input_hidden_size = 2, 8, 256
+    num_heads, head_size = 4, 64
+    hidden_size = num_heads * head_size
+    v_hidden_size = hidden_size
+
+    # for mask_index
+    right_padding_mask = np.random.randint(1, seq_len + 1, size=(batch_size,), dtype=np.int32)
+    end_positions = np.random.randint(1, seq_len + 1, size=(batch_size,), dtype=np.int32)
+    start_positions = np.array([np.random.randint(0, end) for end in end_positions], dtype=np.int32)
+    left_padding_mask = np.concatenate([end_positions, start_positions])
+
+    base_inps = {
+      "input": np.random.randn(batch_size, seq_len, input_hidden_size).astype(np.float32),
+      "weights": np.random.randn(input_hidden_size, hidden_size * 3).astype(np.float32),
+      # bias is required in ORT (segfaults otherwise), eventhough docs says it's optional
+      "bias": np.random.randn(hidden_size * 2 + v_hidden_size).astype(np.float32),
+    }
+    base_opts = {"num_heads": num_heads}
+
+    test_cases = [
+      ({}, {}),
+      ({}, {"scale": 0.1}),
+      ({}, {"scale": 1.0}),
+      ({}, {"unidirectional": 1}),
+      ({"mask_index": right_padding_mask}, {}),
+      ({"mask_index": left_padding_mask}, {}),
+      ({"mask_index": np.random.randint(0, seq_len, size=(batch_size, seq_len), dtype=np.int32)}, {"mask_filter_value": -5000.0}),
+      ({"mask_index": np.random.randint(0, seq_len, size=(batch_size, seq_len, seq_len), dtype=np.int32)}, {"mask_filter_value": -np.inf}),
+      # BUG: when `mask_index` is used with `unidirectional`, the first value must be True
+      # otherwise this will trigger a different ORT behavior where start consecutive Falses will be turned True
+      # e.g. mask_index = [[0, 0, 1, 0, 1, 1, 1, 1], [0, 0, 1, 0, 1, 1, 1, 1]]
+      # will need mask[:, :, 0:1, 0:1] = True
+      ({"mask_index": np.array([[1, 0, 1, 0, 1, 1, 1, 1], [1, 0, 1, 0, 1, 1, 1, 1]], dtype=np.int32)}, {"unidirectional": 1}),
+      ({ "weights": np.random.randn(input_hidden_size, hidden_size + hidden_size + 128).astype(np.float32),
+         "bias": np.random.randn(hidden_size + hidden_size + 128).astype(np.float32)},
+       {"qkv_hidden_sizes": [hidden_size, hidden_size, 128]}),
+      # TODO: past is not tested. ORT gives type error for input
+    ]
+
+    for i, (extra_inps, extra_opts) in enumerate(test_cases):
+      with self.subTest(f"test_attention_{i}"):
+        inps = {**base_inps, **extra_inps}
+        opts = {**base_opts, **extra_opts}
+        outputs = ["output", "present"] if "past" in inps else ["output"]
+        self.helper_test_single_op("Attention", inps, opts, outputs, atol=1e-4)
+
+  def test_skip_layer_normalization(self):
+    shape = (2, 8, 32)
+    for has_beta in [True, False]:
+      for has_bias in [True, False]:
+        with self.subTest(has_beta=has_beta, has_bias=has_bias):
+          hidden_size = shape[-1]
+          inputs = {
+            "input": np.random.randn(*shape).astype(np.float32),
+            "skip": np.random.randn(*shape).astype(np.float32),
+            "gamma": np.random.randn(hidden_size).astype(np.float32),
+          }
+          if has_beta: inputs["beta"] = np.random.randn(hidden_size).astype(np.float32)
+          if has_bias: inputs["bias"] = np.random.randn(hidden_size).astype(np.float32)
+          attributes = {"epsilon": 1e-12}
+          outputs = ["output", "mean", "inv_std_var", "input_skip_bias_sum"]
+          self.helper_test_single_op("SkipLayerNormalization", inputs, attributes, outputs)
+
+  def test_bias_gelu(self):
+    shape = (2,3,4)
+    inputs = {
+      "A": np.random.randn(*shape).astype(np.float32),
+      "B": np.random.randn(shape[-1]).astype(np.float32)
+    }
+    attributes = {}
+    outputs = ["C"]
+    self.helper_test_single_op("BiasGelu", inputs, attributes, outputs)
+
+  def test_qlinear_add(self):
+    for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
+      with self.subTest(dtype=dtype, zero_point=zero_point):
+        dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
+        inputs = {
+          "A": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
+          "A_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+          "A_zero_point": np.array(zero_point, dtype=dtype),
+          "B": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
+          "B_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+          "B_zero_point": np.array(zero_point, dtype=dtype),
+          "C_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+          "C_zero_point": np.array(zero_point, dtype=dtype)
+        }
+        attributes = {}
+        outputs = ["C"]
+        self.helper_test_single_op("QLinearAdd", inputs, attributes, outputs, atol=1) # TODO: look into why this is inaccurate
+
+  def test_qlinear_add_round_half_to_even(self):
+    with self.subTest(test_case="round_half_to_even"):
+      inputs = {
+        "A": np.array([1, 1, 1, 1], dtype=np.int8),
+        "A_scale": np.array(1, dtype=np.float32),
+        "A_zero_point": np.array(0, dtype=np.int8),
+        "B": np.array([1, 5, -3, -7], dtype=np.int8),
+        "B_scale": np.array(1, dtype=np.float32),
+        "B_zero_point": np.array(0, dtype=np.int8),
+        "C_scale": np.array(4, dtype=np.float32),
+        "C_zero_point": np.array(0, dtype=np.int8)
+      }
+      attributes = {}
+      outputs = ["C"]
+      self.helper_test_single_op("QLinearAdd", inputs, attributes, outputs, atol=1) # TODO: look into why this is inaccurate
+
+  def test_qlinear_mul(self):
+    for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
+      with self.subTest(dtype=dtype, zero_point=zero_point):
+        dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
+        inputs = {
+          "A": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
+          "A_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+          "A_zero_point": np.array(zero_point, dtype=dtype),
+          "B": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
+          "B_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+          "B_zero_point": np.array(zero_point, dtype=dtype),
+          "C_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+          "C_zero_point": np.array(zero_point, dtype=dtype)
+        }
+        attributes = {}
+        outputs = ["C"]
+        self.helper_test_single_op("QLinearMul", inputs, attributes, outputs)
+
+    with self.subTest(test_case="round_half_to_even"):
+      inputs = {
+        "A": np.array([1, 1, 1, 1], dtype=np.int8),
+        "A_scale": np.array(1, dtype=np.float32),
+        "A_zero_point": np.array(0, dtype=np.int8),
+        "B": np.array([2, 6, -2, -6], dtype=np.int8),
+        "B_scale": np.array(1, dtype=np.float32),
+        "B_zero_point": np.array(0, dtype=np.int8),
+        "C_scale": np.array(4, dtype=np.float32),
+        "C_zero_point": np.array(0, dtype=np.int8)
+      }
+      attributes = {}
+      outputs = ["C"]
+      self.helper_test_single_op("QLinearMul", inputs, attributes, outputs)
+
+  def test_qlinear_global_average_pool(self):
+    for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
+      for channels_last in [0, 1]:
+        with self.subTest(dtype=dtype, zero_point=zero_point, channels_last=channels_last):
+          dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
+          # NCHW for channels_last=0, NHWC for channels_last=1
+          shape = [1, 3, 32, 32] if channels_last == 0 else [1, 32, 32, 3]
+          inputs = {
+            "X": np.random.randint(dtype_min, dtype_max + 1, shape, dtype=dtype),
+            "x_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+            "x_zero_point": np.array(zero_point, dtype=dtype),
+            "y_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
+            "y_zero_point": np.array(zero_point, dtype=dtype)
+          }
+          attributes = {"channels_last": channels_last}
+          outputs = ["C"]
+          self.helper_test_single_op("QLinearGlobalAveragePool", inputs, attributes, outputs)
+
+  def test_same_device_as_input(self):
+    from tinygrad.nn.onnx import onnx_ops, OpSetId, Domain
+    EmbedLayerNormalization = onnx_ops["EmbedLayerNormalization"]
+    Attention = onnx_ops["Attention"]
+    with Context(DEV="CPU"):
+      input_ids = Tensor([[1, 2]], device="PYTHON", dtype=dtypes.int32)
+      segment_ids = Tensor([[0, 0]], device="PYTHON", dtype=dtypes.int32)
+      word = Tensor.ones(4, 3, device="PYTHON")
+      pos = Tensor.ones(5, 3, device="PYTHON")
+      seg = Tensor.ones(1, 3, device="PYTHON")
+      gamma, beta = Tensor.ones(3, device="PYTHON"), Tensor.zeros(3, device="PYTHON")
+      out, _, _ = EmbedLayerNormalization(input_ids, segment_ids, word, pos, seg, gamma, beta)
+      self.assertEqual(out.device, input_ids.device)
+      out.realize()
+
+      attn = Attention[OpSetId(Domain.MICROSOFT_CONTRIB_OPS, 1)]
+      x = Tensor.ones(1, 2, 4, device="PYTHON")
+      w = Tensor.ones(4, 12, device="PYTHON")
+      mask = Tensor([2, 0], device="PYTHON", dtype=dtypes.int32)
+      out, _ = attn(x, w, mask_index=mask, num_heads=1, unidirectional=1)
+      self.assertEqual(out.device, x.device)
+      out.realize()
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/external/external_test_onnx_runner.py

@@ -0,0 +1,165 @@
+import unittest, onnx, tempfile, pathlib
+import numpy as np
+from tinygrad import Tensor
+from tinygrad.uop.ops import Ops
+from typing import Any
+from tinygrad.nn.onnx import OnnxRunner, OnnxPBParser, OnnxDataType
+from hypothesis import given, strategies as st
+
+# copied from test_const_folding.py
+def _check_ast_count(desired_count:int, t:Tensor):
+  # NOTE: this has side effect because everything can be scheduled only once
+  linear = t.schedule_linear()
+  asts = [call for call in linear.src if call.src[0].op is Ops.SINK]
+  assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
+
+def build_onnx(nodes, from_disk:bool=True, **kwargs):
+  """Helper to build and return an OnnxRunner from ONNX nodes."""
+  graph = onnx.helper.make_graph(nodes, 'test', kwargs.get('inputs', []), kwargs.get('outputs', []), kwargs.get('initializers', []))
+  model = onnx.helper.make_model(graph)
+  if from_disk:
+    with tempfile.TemporaryDirectory() as tmpdir:
+      tmp_path = pathlib.Path(tmpdir)
+      model_path = tmp_path / "model.onnx"
+      onnx.save(model, model_path)
+      runner = OnnxRunner(model_path)
+  else:
+    # use the in-memory method
+    runner = OnnxRunner(Tensor(model.SerializeToString(), device="PYTHON"))
+  return runner
+
+class TestOnnxRunner(unittest.TestCase):
+  def _test_const_fold_unary_op(self, from_disk:bool):
+    runner = build_onnx(
+        nodes=[
+          onnx.helper.make_node('Expand', ['inp', 'shape'], ['expanded']),
+          onnx.helper.make_node('Exp', ['expanded'], ['output'])
+        ],
+        outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (5,))],
+        initializers=[
+          onnx.helper.make_tensor('inp', onnx.TensorProto.FLOAT, (), [1.0]),
+          onnx.helper.make_tensor('shape', onnx.TensorProto.INT64, (1,), [5])
+        ],
+        from_disk=from_disk)
+    output = runner({'inp': Tensor([1.0])})['output']
+    _check_ast_count(0, output)
+
+  def _test_const_fold_binary_op(self, from_disk:bool):
+    runner = build_onnx(
+        nodes=[onnx.helper.make_node('Add', ['inp', 'const'], ['output'])],
+        outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (4,))],
+        initializers=[
+          onnx.helper.make_tensor('inp', onnx.TensorProto.FLOAT, (4,), [1, 2, 3, 4]),
+          onnx.helper.make_tensor('const', onnx.TensorProto.FLOAT, (), [0])
+        ],
+        from_disk=from_disk)
+    output = runner({'inp': Tensor([1, 2, 3, 4])})['output']
+    _check_ast_count(0, output)
+
+  @unittest.skip("const folding is removed")
+  def test_const_fold_from_disk(self):
+    self._test_const_fold_unary_op(True)
+    self._test_const_fold_binary_op(True)
+
+  @unittest.skip("const folding is removed")
+  def test_const_fold_from_memory(self):
+    self._test_const_fold_unary_op(False)
+    # TODO: understand this and fix this, bitcast related
+    # self._test_const_fold_binary_op(False)
+
+  def test_external_data_loading(self):
+    weights = np.arange(4, dtype=np.float32)
+    tensor_with_data = onnx.helper.make_tensor('weights', onnx.TensorProto.FLOAT, weights.shape, weights.tobytes(), raw=True)
+    graph = onnx.helper.make_graph(
+        nodes=[onnx.helper.make_node('Add', ['inp', 'weights'], ['output'])],
+        name='test_external',
+        inputs=[onnx.helper.make_tensor_value_info('inp', onnx.TensorProto.FLOAT, (1,))],
+        outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, weights.shape)],
+        initializer=[tensor_with_data]
+    )
+    model = onnx.helper.make_model(graph)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+      tmp_path = pathlib.Path(tmpdir)
+      model_path = tmp_path / "model.onnx"
+      onnx.save_model(model, model_path, save_as_external_data=True, all_tensors_to_one_file=True, size_threshold=0, location="weights.onnx_data")
+      runner = OnnxRunner(model_path)
+      output = runner({'inp': Tensor([1])})['output']
+      np.testing.assert_equal(output.numpy(), weights + 1)
+
+all_dtypes = list(OnnxDataType)
+
+class TestOnnxRunnerDtypes(unittest.TestCase):
+  """
+  Internal tensors (initializers, attributes) fallback to default dtype if unsupported by device.
+  External tensors (inputs) preserve their original dtype - user must ensure compatibility with device.
+  """
+  def _get_expected_dtype(self, onnx_dtype: int, is_input: bool): return OnnxDataType(onnx_dtype).to_dtype()
+
+  @given(onnx_dtype=st.sampled_from(all_dtypes))
+  def test_input_dtype(self, onnx_dtype: int):
+    expected_dtype = self._get_expected_dtype(onnx_dtype, True)
+    runner = build_onnx(
+        nodes=[onnx.helper.make_node('Identity', ['input'], ['output'])],
+        inputs=[onnx.helper.make_tensor_value_info('input', onnx_dtype, ())],
+        outputs=[onnx.helper.make_tensor_value_info('output', onnx_dtype, ())],
+        from_disk=False)
+    self.assertEqual(runner.graph_inputs['input'].dtype, expected_dtype)
+
+  @given(onnx_dtype=st.sampled_from(all_dtypes))
+  def test_initializer_dtype(self, onnx_dtype: int):
+    expected_dtype = self._get_expected_dtype(onnx_dtype, False)
+    runner = build_onnx(
+        nodes=[onnx.helper.make_node('Identity', ['initializer'], ['output'])],
+        outputs=[onnx.helper.make_tensor_value_info('output', onnx_dtype, (2,))],
+        initializers=[onnx.helper.make_tensor('initializer', onnx_dtype, (2,), [1, 2])],
+        from_disk=False)
+    self.assertEqual(runner.graph_values['initializer'].dtype, expected_dtype)
+
+  @given(onnx_dtype=st.sampled_from(all_dtypes))
+  def test_node_attribute_dtype(self, onnx_dtype: int):
+    expected_dtype = self._get_expected_dtype(onnx_dtype, False)
+    value_tensor = onnx.helper.make_tensor('value', onnx_dtype, (2,), [1, 2])
+    runner = build_onnx(
+        nodes=[onnx.helper.make_node('Constant', [], ['output'], value=value_tensor)],
+        outputs=[onnx.helper.make_tensor_value_info('output', onnx_dtype, (2,))],
+        from_disk=False)
+    self.assertEqual(runner.graph_nodes[0].opts['value'].dtype, expected_dtype)
+
+# from openpilot selfdrive/modeld/get_model_metadata.py
+class MetadataOnnxPBParser(OnnxPBParser):
+  def _parse_ModelProto(self) -> dict:
+    obj: dict[str, Any] = {"graph": {"input": [], "output": []}, "metadata_props": []}
+    for fid, wire_type in self._parse_message(self.reader.len):
+      match fid:
+        case 7: obj["graph"] = self._parse_GraphProto()
+        case 14: obj["metadata_props"].append(self._parse_StringStringEntryProto())
+        case _: self.reader.skip_field(wire_type)
+    return obj
+
+class TestOnnxMetadata(unittest.TestCase):
+  def test_metadata_props(self):
+    graph = onnx.helper.make_graph(
+      nodes=[onnx.helper.make_node('Identity', ['input'], ['output'])],
+      name='test',
+      inputs=[onnx.helper.make_tensor_value_info('input', onnx.TensorProto.FLOAT, (1, 3))],
+      outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (1, 3))],
+    )
+    model = onnx.helper.make_model(graph)
+    model.metadata_props.append(onnx.StringStringEntryProto(key="model_checkpoint", value="v1.0"))
+    model.metadata_props.append(onnx.StringStringEntryProto(key="output_slices", value="dGVzdA=="))
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+      model_path = pathlib.Path(tmpdir) / "model.onnx"
+      onnx.save(model, model_path)
+      parsed = MetadataOnnxPBParser(model_path).parse()
+
+    # metadata_props should be accessible as dicts with "key" and "value"
+    self.assertEqual(len(parsed["metadata_props"]), 2)
+    self.assertEqual(parsed["metadata_props"][0]["key"], "model_checkpoint")
+    self.assertEqual(parsed["metadata_props"][0]["value"], "v1.0")
+    self.assertEqual(parsed["metadata_props"][1]["key"], "output_slices")
+    self.assertEqual(parsed["metadata_props"][1]["value"], "dGVzdA==")
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_opt.py

@@ -0,0 +1,240 @@
+#!/usr/bin/env python
+import gc, unittest
+import numpy as np
+import torch
+
+from tinygrad import GlobalCounters, Tensor, Device
+from tinygrad.helpers import getenv
+from tinygrad.nn.state import get_parameters
+from tinygrad.engine.realize import capturing, run_linear
+from tinygrad.tensor import _to_np_dtype
+
+class CLCache:
+  def __init__(self, allowed=None, strict=False, preclear=True, var_vals=None):
+    self.allowed, self.strict, self.preclear, self.var_vals = allowed, strict, preclear, var_vals if var_vals is not None else {}
+    self.count = 0
+  def add_linear(self, linear, var_vals):
+    self.count += len(linear.src)
+    run_linear(linear, var_vals)
+  def __enter__(self):
+    if self.preclear:
+      gc.collect()
+      for x in [x for x in gc.get_objects() if isinstance(x, Tensor)]:
+        x.realize()
+      GlobalCounters.reset()
+    capturing.append(self)
+    print("cache: entering")
+    return self
+  def __exit__(self, _type, value, traceback):
+    capturing.clear()
+    print(f"cache: exiting with size {self.count}", f"allowed {self.allowed}" if self.allowed is not None else "")
+    if self.allowed is not None:
+      assert self.count <= self.allowed, f"{self.count} > {self.allowed}"
+
+from extra.models.convnext import ConvNeXt
+from extra.models.efficientnet import EfficientNet
+from extra.models.resnet import ResNet18
+from extra.models.vit import ViT
+
+@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented")
+class TestInferenceMinKernels(unittest.TestCase):
+  def setUp(self):
+    self.training_old = Tensor.training
+    Tensor.training = False
+  def tearDown(self):
+    Tensor.training = self.training_old
+
+  def test_convnext(self):
+    model = ConvNeXt()
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
+    img = Tensor.randn(1, 3, 224, 224)
+    with CLCache(143):
+      model(img).realize()
+
+  def test_enet(self):
+    model = EfficientNet(getenv("ENET_NUM", 0), has_se=False)
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
+    img = Tensor.randn(1, 3, 224, 224)
+    with CLCache(51):
+      model.forward(img).realize()
+
+  def test_enet_se(self):
+    model = EfficientNet(getenv("ENET_NUM", 0), has_se=True)
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
+    img = Tensor.randn(1, 3, 224, 224)
+    # TODO: this seems very high
+    with CLCache(115):
+      model.forward(img).realize()
+
+  def test_resnet(self):
+    model = ResNet18()
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
+    img = Tensor.randn(1, 3, 224, 224)
+    with CLCache(23):
+      model.forward(img).realize()
+
+  def test_vit(self):
+    model = ViT(embed_dim=192, num_heads=3)
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
+    img = Tensor.randn(1, 3, 224, 224)
+    with CLCache(209) as cache: # NOTE: this is way too high
+      out = model.forward(img)
+      assert cache.count == 0, "ViT prerealized?"
+      out.realize()
+
+  @unittest.skip("llama is fp16 but CI does not have fp16")
+  def test_llama(self):
+    from examples.llama import Transformer
+    args_tiny = {"dim": 512, "hidden_dim": 1024, "n_heads": 8, "n_layers": 4, "norm_eps": 1e-05, "vocab_size": 1000}
+    model = Transformer(**args_tiny)
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
+    inp = Tensor([[1,2,3,4]])
+    with CLCache(100):
+      model(inp, 0).realize()
+
+@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented")
+class TestOptBinOp(unittest.TestCase):
+  def _test_no_binop_rerun(self, f1, f2=None, allowed=1):
+    a = Tensor.randn(16, 16)
+    b = Tensor.randn(16, 16)
+    with CLCache() as cache:
+      c = f1(a, b)
+      if f2 is not None: d = f2(a, b)
+      c.realize()
+      if f2 is not None: d.realize()
+      assert cache.count == allowed, "binop was rerun!"
+    if f2 is not None: np.testing.assert_allclose(c.numpy().ravel(), d.numpy().ravel(), rtol=1e-3, atol=1e-5)
+
+  def test_no_binop_rerun(self): return self._test_no_binop_rerun(lambda a,b: a*b, lambda a,b: (a*b).reshape(16, 16, 1))
+  def test_no_binop_rerun_alt(self): return self._test_no_binop_rerun(lambda a,b: (a*b).reshape(16, 16, 1), lambda a,b: a*b)
+  def test_no_binop_rerun_reduce_broadcast(self):
+    return self._test_no_binop_rerun(lambda a,b: a.sum()+b, lambda a,b: a.sum().reshape(1,1)+b, allowed=2)
+
+  @unittest.skip("this test started failing with the new change, based movementop issue")
+  def test_no_binop_rerun_transposed(self): return self._test_no_binop_rerun(lambda a,b: (a.T*b.T).T, lambda a,b: a*b)
+  def test_no_binop_rerun_mid_reshape(self): return self._test_no_binop_rerun(lambda a,b: (a*b).reshape(256)+a.reshape(256))
+
+  # currently non working tests
+  # def test_no_binop_rerun_preshape(self): return self._test_no_binop_rerun(lambda a,b: a.reshape(16, 16, 1)*b.reshape(16, 16, 1), lambda a,b: a*b)
+  #def test_no_binop_rerun_reduce(self): return self._test_no_binop_rerun(lambda a,b: (a*b).sum(), lambda a,b: (a*b).reshape(16, 16, 1).sum())
+  #def test_no_binop_rerun_reduce_alt(self): return self._test_no_binop_rerun(lambda a,b: a.sum(1)+b[0], lambda a,b: a.sum(1).reshape(1,16)+b[0])
+
+@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented")
+class TestOptReduceLoop(unittest.TestCase):
+  def test_loop_left(self):
+    a = Tensor.randn(16, 16)
+    b = Tensor.randn(16, 16)
+    with CLCache() as cache:
+      t = a.sum(0)
+      b = t.reshape(16,1).expand(16,16).sum(0)
+      c = (t+b)
+      c.realize()
+      assert cache.count == 2, "loop left fusion broken"
+
+  def test_loop_right(self):
+    a = Tensor.randn(16, 16)
+    b = Tensor.randn(16, 16)
+    with CLCache() as cache:
+      t = a.sum(0)
+      b = t.reshape(16,1).expand(16,16).sum(0)
+      c = (b+t)
+      c.realize()
+      assert cache.count == 2, "loop right fusion broken"
+
+@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented")
+class TestOptWChild(unittest.TestCase):
+  @unittest.skip("this no longer happens, use realize")
+  def test_unrealized_child(self):
+    a = Tensor.randn(16, 16)
+    b = Tensor.randn(16, 16)
+    with CLCache() as cache:
+      c = (a*b).sum()
+      d = c+1
+      e = c+2 # noqa: F841
+      d.realize()
+      assert cache.count == 2, "don't fuse if you have children"
+
+@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented")
+class TestOpt(unittest.TestCase):
+  def test_muladd(self):
+    a,b,c = [Tensor.randn(2,2).realize() for _ in range(3)]
+    na,nb,nc = a.numpy(),b.numpy(),c.numpy()
+    with CLCache(allowed=1):
+      d = a * b + c
+      d.realize()
+    np.testing.assert_allclose(d.numpy(), na*nb+nc, rtol=1e-5, atol=1e-7)
+
+  def test_permute_was_pushed(self):
+    a = Tensor.randn(16, 16, 16)
+    with CLCache(1):
+      c = a.sum(2)
+      d = c.permute(1,0).contiguous()
+      d.realize()
+    np.testing.assert_allclose(a.numpy().sum(2).transpose(1,0), d.numpy(), rtol=1e-3, atol=1e-5)
+
+  def test_permute_was_pushed_through_contract_reshape(self):
+    a = Tensor.randn(4, 4, 4, 4, 4)
+    with CLCache(1):
+      c = a.sum(-1)
+      d = c.reshape(16,16).permute(1,0).contiguous()
+      d.realize()
+    np.testing.assert_allclose(a.numpy().sum(-1).reshape(16,16).transpose(1,0), d.numpy(), rtol=1e-3, atol=1e-5)
+
+  def test_permute_was_pushed_through_contractw1s_reshape(self):
+    a = Tensor.randn(4, 4, 4, 4, 4)
+    with CLCache(1):
+      c = a.sum(-1)
+      d = c.reshape(16,1,16).permute(2,1,0).contiguous()
+      d.realize()
+    np.testing.assert_allclose(a.numpy().sum(-1).reshape(16,1,16).transpose(2,1,0), d.numpy(), rtol=1e-3, atol=1e-5)
+
+  def test_permute_was_pushed_through_expand_reshape(self):
+    a = Tensor.randn(16, 16, 16)
+    with CLCache(1):
+      c = a.sum(2)
+      d = c.reshape(4,4,4,4).permute(2,3,0,1).contiguous()
+      d.realize()
+    np.testing.assert_allclose(a.numpy().sum(2).transpose(1,0).reshape(4,4,4,4), d.numpy(), rtol=1e-3, atol=1e-5)
+
+  def test_no_reduceop_rerun(self):
+    a = Tensor.randn(16, 16, 16)
+    with CLCache() as cache:
+      c = a.sum(2)
+      d = a.sum(2).permute(1,0)
+      c.realize()
+      d.realize()
+      cache_len = cache.count
+    np.testing.assert_allclose(c.numpy().transpose(1,0), d.numpy(), rtol=1e-3, atol=1e-5)
+    assert cache_len == 1, "reduceop was rerun!"
+
+  def test_no_reduceop_rerun_alt(self):
+    a = Tensor.randn(16, 16, 16)
+    with CLCache() as cache:
+      c = a.sum(2).permute(1,0)
+      d = a.sum(2)
+      c.realize()
+      d.realize()
+      cache_len = cache.count
+    np.testing.assert_allclose(c.numpy(), d.numpy().transpose(1,0), rtol=1e-3, atol=1e-5)
+    assert cache_len == 1, "reduceop was rerun!"
+
+  def test_expand_reduce_is_folded_on_same_axis(self):
+    for axis in [0, 1]:
+      for n in [4, 8, 16]:
+        b = torch.ones(n, n).sum(axis).reshape(n, 1).expand(n, n).sum(axis)
+        with CLCache(allowed=3):
+          a = Tensor.ones(n, n).contiguous().sum(axis).reshape(n, 1).expand(n, n).sum(axis)
+          a.realize()
+        np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
+
+  def test_expand_reduce_is_folded_on_different_axes(self):
+    axis1, axis2 = 0, 1
+    for n in [4, 8, 16]:
+      b = torch.ones(n, n).sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
+      with CLCache(allowed=3):
+        a = Tensor.ones(n, n).contiguous().sum(axis1).reshape(n, 1).expand(n, n).sum(axis2)
+        a.realize()
+      np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-3, atol=1e-5)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_optim.py

@@ -0,0 +1,220 @@
+#!/usr/bin/env python
+import unittest, math
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.optimizers import Lamb
+from tensorflow.python.ops import math_ops
+from extra.lr_scheduler import LRSchedulerGroup
+
+from tinygrad.tensor import Tensor
+from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, AdamW
+
+from test.external.mlperf_resnet.lars_optimizer import LARSOptimizer
+
+from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup, CosineAnnealingLRWithWarmup, LambdaLR, LambdaLinearScheduler
+from test.external.mlperf_resnet.lars_util import PolynomialDecayWithWarmup as PolynomialDecayWithWarmup_tf
+
+np.random.seed(1337)
+x_init = np.random.randn(1,4).astype(np.float32)
+W_init = np.random.randn(4,4).astype(np.float32)
+m_init = np.random.randn(1,4).astype(np.float32)
+
+class TinyNet:
+  def __init__(self):
+    self.x = Tensor(x_init.copy())
+    self.W = Tensor(W_init.copy())
+    self.m = Tensor(m_init.copy())
+
+  def forward(self):
+    out = self.x.matmul(self.W).relu()
+    out = out.log_softmax(1)
+    out = out.mul(self.m).add(self.m).sum()
+    return out
+
+class TinyNetTF:
+  def __init__(self):
+    self.x = tf.Variable(x_init.copy(), trainable=True, name="x")
+    self.W = tf.Variable(W_init.copy(), trainable=True, name="W")
+    self.m = tf.constant(m_init.copy())
+
+  def forward(self):
+    out = tf.matmul(self.x, self.W)
+    out = tf.nn.relu(out)
+    out = tf.nn.log_softmax(out, axis=1)
+    out = tf.multiply(out, self.m) + self.m
+    out = tf.reduce_sum(out)
+    return out
+
+def step(optim, steps=1, kwargs={}, scheduler=None, schedopts=None, do_optim=True):
+  net = TinyNet()
+  optim = optim([net.x, net.W], **kwargs)
+  if scheduler is not None: scheduler = scheduler(optim, **schedopts)
+  lrs = []
+  for _ in range(steps):
+    if do_optim:
+      out = net.forward()
+      optim.zero_grad()
+      out.backward()
+    lrs.append(optim.lr.item() if not isinstance(optim, OptimizerGroup) else optim.optimizers[0].lr.item())
+    if do_optim: optim.step()
+    if scheduler is not None: scheduler.step()
+  return lrs, net.x.detach().numpy(), net.W.detach().numpy()
+
+def step_tf(optim, steps=1, kwargs={}, scheduler=None, schedopts=None, do_optim=True):
+  net = TinyNetTF()
+  if scheduler is not None: kwargs['lr'] = scheduler(**schedopts)
+  optim = optim(**kwargs)
+  lrs = []
+  for _ in range(steps):
+    if do_optim:
+      with tf.GradientTape() as tape:
+        out = net.forward()
+
+    lr_t = optim.learning_rate
+    # refer to test/external/mlperf_resnet/lars_optimizer.py:_prepare_local
+    if callable(lr_t): lr_t = lr_t(math_ops.cast(optim.iterations, tf.float32))
+    lrs.append(lr_t)
+
+    if do_optim:
+      grads = tape.gradient(out, [net.x, net.W])
+      optim.apply_gradients(zip(grads, [net.x, net.W]))
+      # optim calls scheduler in tf
+    else:
+      optim._iterations.assign_add(1)
+  return lrs, net.x.numpy(), net.W.numpy()
+
+# skip list is skipping W
+def create_tiny_lars(params, lr, skip_list=False):
+  if skip_list: return OptimizerGroup(LARS([params[0]], lr), SGD([params[1]], lr, classic=True, weight_decay=0., momentum=.9))
+  return LARS(params, lr)
+def create_tf_lars(lr, skip_list=False): return LARSOptimizer(lr, skip_list=["W"] if skip_list else None)
+def create_tf_lamb(lr=0.001, b1=0.9, b2=0.999, eps=1e-7, weight_decay=0.0):
+  return Lamb(learning_rate=float(lr), beta_1=b1, beta_2=b2, epsilon=eps, weight_decay=weight_decay)
+
+def create_tiny_polylr(optim, initial_lr, end_lr, train_steps, warmup, power=2, skip_list=False):
+  assert power == 2
+  if skip_list: return LRSchedulerGroup(
+    PolynomialDecayWithWarmup(optim[0], initial_lr, end_lr, train_steps, warmup, power),
+    PolynomialDecayWithWarmup(optim[1], initial_lr, end_lr, train_steps, warmup, power))
+  return PolynomialDecayWithWarmup(optim, initial_lr, end_lr, train_steps, warmup, power)
+def create_tf_polylr(initial_lr, end_lr, train_steps, warmup, power=2, skip_list=False):
+  assert power == 2
+  return PolynomialDecayWithWarmup_tf(1, 1, train_steps,
+                                      initial_learning_rate=initial_lr, end_learning_rate=end_lr, warmup_epochs=warmup)
+
+class ExternalTestOptim(unittest.TestCase):
+  def setUp(self):
+    self.old_training = Tensor.training
+    Tensor.training = True
+  def tearDown(self):
+    Tensor.training = self.old_training
+
+  def _test_optim(self, tinygrad_optim, tensorflow_optim, steps, opts, atol, rtol, tiny_sched=None, tf_sched=None, schedopts=None, do_optim=True):
+    for x,y in zip(step(tinygrad_optim, steps=steps, kwargs=opts, scheduler=tiny_sched, schedopts=schedopts, do_optim=do_optim),
+                   step_tf(tensorflow_optim, steps=steps, kwargs=opts, scheduler=tf_sched, schedopts=schedopts, do_optim=do_optim)):
+      np.testing.assert_allclose(x, y, atol=atol, rtol=rtol)
+
+  def _test_lamb(self, steps, opts, atol, rtol): self._test_optim(LAMB, create_tf_lamb, steps, opts, atol, rtol)
+  def _test_lars(self, steps, opts, atol, rtol): self._test_optim(create_tiny_lars, create_tf_lars, steps, opts, atol, rtol)
+  def _test_lars_polylr(self, steps, opts, schedopts, atol, rtol, do_optim=True):
+    self._test_optim(create_tiny_lars, create_tf_lars, steps, opts, atol, rtol,
+                     tiny_sched=create_tiny_polylr, tf_sched=create_tf_polylr, schedopts=schedopts, do_optim=do_optim)
+
+  def test_lamb(self): self._test_lamb(1, {'lr': 0.001}, 1e-5, 0)
+  def test_lamb_high_lr(self): self._test_lamb(1, {'lr': 10}, 1e-5, 1e-5)
+
+  def test_multistep_lamb(self): self._test_lamb(10, {'lr': 0.001}, 1e-5, 0)
+  def test_multistep_lamb_high_lr(self): self._test_lamb(10, {'lr': 10}, 1e-5, 3e-4)
+
+  def test_lars(self): self._test_lars(1, {'lr': 0.01}, 1e-5, 0)
+  def test_lars_high_lr(self): self._test_lars(1, {'lr': 10}, 1e-5, 1e-5)
+  def test_multistep_lars(self): self._test_lars(10, {'lr': 0.001}, 1e-5, 0)
+  def test_multistep_lars_high_lr(self): self._test_lars(10, {'lr': 10}, 1e-5, 3e-4)
+  def test_lars_skip(self): self._test_lars(10, {'lr': 10, 'skip_list': True}, 1e-5, 3e-4)
+  def test_lars_skip_high_lr(self): self._test_lars(1, {'lr': 10, 'skip_list': True}, 1e-5, 1e-5)
+  def test_lars_skip_multistep(self): self._test_lars(10, {'lr': 0.001, 'skip_list': True}, 1e-5, 0)
+  def test_lars_skip_multistep_high_lr(self): self._test_lars(10, {'lr': 10, 'skip_list': True}, 1e-5, 3e-4)
+
+  def test_lars_polylr(self):
+    self._test_lars_polylr(10, {'lr': 1.0}, {
+      'initial_lr': 1.0,
+      'end_lr': 1e-4,
+      'train_steps': 10,
+      'warmup': 3
+    }, 1e-5, 1e-5)
+  def test_lars_polylr_large(self):
+    self._test_lars_polylr(100, {'lr': 10.0}, {
+      'initial_lr': 10.0,
+      'end_lr': 1e-5,
+      'train_steps': 100,
+      'warmup': 43
+    }, 1e-5, 1e-5, do_optim=False)
+  def test_lars_polylr_skip(self):
+    self._test_lars_polylr(10, {'lr': 1.0, 'skip_list': True}, {
+      'initial_lr': 1.0,
+      'end_lr': 1e-4,
+      'train_steps': 10,
+      'warmup': 3,
+      'skip_list': True
+    }, 1e-5, 1e-5)
+
+  @unittest.skip("slow, but you can run this locally to check")
+  def test_lars_polylr_resnet(self):
+    train_files = 1_281_167
+    BS = 624
+    steps_per_epoch = train_files // BS
+    epochs = 45
+    warmup_epochs = 5
+    self._test_lars_polylr(steps_per_epoch * epochs, {'lr': 10.4}, {
+      'initial_lr': 10.4,
+      'end_lr': 1e-4,
+      # step counts for BS=624 EPOCHS=45 resnet
+      'train_steps': steps_per_epoch * epochs,
+      'warmup': steps_per_epoch * warmup_epochs,
+    }, 1e-5, 1e-5, do_optim=False)
+
+
+class TestCosineAnnealingLRWithWarmup(unittest.TestCase):
+  # only tests the lr
+  def _test_lr(self, base_lr, end_lr, warmup_steps, decay_steps):
+    net = TinyNet()
+    optim = AdamW([net.W], lr=0.0)
+    tiny_lr = CosineAnnealingLRWithWarmup(optim, base_lr, end_lr, warmup_steps, decay_steps)
+    lr = []
+    for _ in range(warmup_steps+decay_steps):
+      lr.append(optim.lr.item())
+      tiny_lr.step()
+    # reimplemented in python
+    expected = []
+    for i in range(warmup_steps): expected.append((i+1)/warmup_steps*base_lr)
+    for i in range(decay_steps): expected.append(end_lr+(base_lr-end_lr)*(1+math.cos((i+1)/decay_steps*math.pi))/2)
+    np.testing.assert_allclose(lr, expected, rtol=1e-5)
+
+  def test_lr_0(self): self._test_lr(3e-4, 8e-5, 3, 5)
+  def test_lr_1(self): self._test_lr(3e-4, 8e-5, 10, 20)
+  def test_lr_llama3(self): self._test_lr(8e-5, 8e-7, 20, 100)
+
+class TestLambdaLRLinearWarmup(unittest.TestCase):
+  def test_linear_lr_warmup(self):
+    BS, BASE_LR = 304, 2.5e-7
+    lr = BS * BASE_LR
+    # Use a dummy Tensor parameter for optimizer because the lr_scheduler only needs the optimizer's device and lr, the params aren't touched.
+    optimizer = AdamW([Tensor([1.])])
+    lambda_lr_callback = LambdaLinearScheduler(1000, 1.0, 1.0, 1e-06, 10000000000000).schedule
+    lr_scheduler = LambdaLR(optimizer, Tensor(lr, device=optimizer.device), lambda_lr_callback)
+    lrs = {}
+
+    # with above settings, optimizer.lr should warm up to lr over 1000 steps linearly
+    for i in range(1200):
+      lr_scheduler.step()
+      if i in {0, 499, 998, 999, 1000, 1199}:
+        lrs[i] = optimizer.lr.item()
+
+    np.testing.assert_allclose(lr, lrs[999], rtol=0, atol=1e-11)
+    np.testing.assert_equal(lrs[999], lrs[1000])
+    np.testing.assert_equal(lrs[999], lrs[1199])
+    np.testing.assert_allclose(lrs[999] / lrs[0], 1000, rtol=0, atol=1)
+    np.testing.assert_allclose(lrs[999] / lrs[499], 2, rtol=0, atol=1e-5)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_process_count.py

@@ -0,0 +1,57 @@
+import os, sys, time, multiprocessing
+
+N = int(os.environ.get("NPROC", str(os.cpu_count())))
+DEVICE = os.environ.get("DEV", "AMD")
+
+# this tests the total number of processes that can be running tinygrad at a time
+def proc(i, device, stop_evt):
+  from tinygrad import Tensor
+
+  try:
+    a = Tensor.ones(2, device=device).contiguous()
+    b = Tensor.ones(2, device=device).contiguous()
+    c = (a + b).realize()
+    assert c.tolist() == [2, 2]
+  except Exception as e:
+    # fail if it fails
+    print(f"[child {i:2d}] tinygrad op failed: {e}", file=sys.stderr)
+    # non-zero exit code propagated back to parent
+    sys.exit(1)
+
+  # TODO: wait here for global exit if success. fail if it fails
+  # -> We wait on a global Event shared from the parent.
+  print(f"[child {i:2d}] success")
+  stop_evt.wait()
+  # Normal successful exit
+  sys.exit(0)
+
+if __name__ == "__main__":
+  print(f"testing {N} concurrent tinygrad processes")
+
+  # global exit event, shared by all children
+  stop_evt = multiprocessing.Event()
+  procs = []
+
+  # launch n proc of proc 1 per 200 ms
+  for i in range(N):
+    p = multiprocessing.Process(target=proc, args=(i, DEVICE, stop_evt), name=f"tinygrad-proc-{i}")
+    p.start()
+    procs.append(p)
+    time.sleep(0.1)  # 100 ms between launches
+
+  # signal global exit
+  time.sleep(0.5)
+  stop_evt.set()
+
+  # join all children
+  for p in procs: p.join()
+
+  # check for failures
+  failed = [p for p in procs if p.exitcode != 0]
+  if failed:
+    print(f"{len(failed)} / {len(procs)} processes failed "
+          f"with exit codes: {[p.exitcode for p in failed]}", file=sys.stderr)
+    sys.exit(1)
+
+  print(f"All {len(procs)} tinygrad processes ran successfully")
+  sys.exit(0)

tinygrad_repo/test/external/external_test_schedule_scaling.py

@@ -0,0 +1,134 @@
+import unittest, time
+from tinygrad import Tensor
+
+class TestScheduleScaling(unittest.TestCase):
+  """Test that .schedule() scales linearly with graph size (no O(n^2) behavior)."""
+
+  def _assert_linear(self, fn, n_small=200, n_large=1000):
+    """Assert schedule time scales at most ~linearly: time(n_large)/time(n_small) should be close to n_large/n_small."""
+    fn(n_small).schedule_linear()  # warmup
+    t_small = min(self._time_schedule(fn, n) for n in [n_small]*3)
+    t_large = min(self._time_schedule(fn, n) for n in [n_large]*3)
+    size_ratio = n_large / n_small  # 5.0
+    time_ratio = t_large / t_small
+    # O(n) -> time_ratio ~ 5, O(n^2) -> time_ratio ~ 25. threshold at 10 catches n^2 with margin.
+    self.assertLess(time_ratio / size_ratio, 2.0,
+      f"schedule appears superlinear: n={n_small} {t_small*1e3:.1f}ms, n={n_large} {t_large*1e3:.1f}ms "
+      f"(time grew {time_ratio:.1f}x for {size_ratio:.0f}x size, per-node ratio {time_ratio/size_ratio:.2f})")
+
+  @staticmethod
+  def _time_schedule(fn, n) -> float:
+    st = time.perf_counter()
+    fn(n).schedule_linear()
+    return time.perf_counter() - st
+
+  # *** rangeify: ending_ranges accumulation and consumer merge ***
+
+  # ending_ranges accumulation via sum([], []) and nested scan in run_rangeify.
+  # this creates reduce ops whose ending_ranges lists grow with graph depth, causing O(n^2) list copies.
+  def test_multi_reduce_scaling(self):
+    def multi_reduce(n):
+      x = Tensor.empty(256, 256)
+      for _ in range(n):
+        s = x.sum(axis=-1, keepdim=True)
+        x = x + s + s
+      return x
+    self._assert_linear(multi_reduce)
+
+  # reduce+elementwise chain stresses ending_ranges propagation and post-rangeify rewrites
+  def test_wide_reduce_scaling(self):
+    def wide_reduce(n):
+      x = Tensor.empty(256, 256)
+      for _ in range(n):
+        x = x + x.sum(axis=-1, keepdim=True)
+      return x
+    self._assert_linear(wide_reduce)
+
+  # expand ops inject into ending_ranges via the EXPAND path in run_rangeify
+  def test_expand_reduce_scaling(self):
+    def expand_reduce(n):
+      x = Tensor.empty(256, 1)
+      for _ in range(n):
+        y = x.expand(256, 256)
+        x = (y + y).sum(axis=-1, keepdim=True)
+      return x
+    self._assert_linear(expand_reduce)
+
+  # *** graph_rewrite: multi-consumer DAG patterns ***
+
+  # multi-consumer diamond pattern (fan-out/fan-in) stresses consumer_rngs merge in run_rangeify
+  def test_diamond_scaling(self):
+    def diamond(n):
+      x = Tensor.empty(256, 256)
+      for _ in range(n):
+        a = x + 1
+        b = x + 2
+        x = a + b
+      return x
+    self._assert_linear(diamond)
+
+  # elementwise chain baseline — should be trivially O(n)
+  def test_chain_scaling(self):
+    def chain(n):
+      x = Tensor.empty(256, 256)
+      for _ in range(n): x = x + 1
+      return x
+    self._assert_linear(chain)
+
+  # softmax has multi-consumer structure (x used for max, exp, and sum), stresses graph_rewrite on DAGs
+  def test_softmax_scaling(self):
+    def softmax_chain(n):
+      x = Tensor.empty(64, 256)
+      for _ in range(n): x = x.softmax(axis=-1)
+      return x
+    self._assert_linear(softmax_chain)
+
+  # *** post-rangeify: symbolic rewrites, kernel splitting ***
+
+  # matmul chain stresses symbolic+reduce_collapse and split_store
+  def test_matmul_scaling(self):
+    def matmul_chain(n):
+      xs = [Tensor.empty(32, 32) for _ in range(n + 1)]
+      result = xs[0]
+      for i in range(n): result = result @ xs[i + 1]
+      return result
+    self._assert_linear(matmul_chain)
+
+  # contiguous chain stresses remove_bufferize callbacks (toposort per BUFFERIZE node)
+  def test_contiguous_scaling(self):
+    def contiguous_chain(n):
+      x = Tensor.empty(256, 256)
+      for _ in range(n): x = (x + 1).contiguous()
+      return x
+    self._assert_linear(contiguous_chain)
+
+  # *** schedule: AFTER handling, assign ***
+
+  # assign chain stresses AFTER cycle detection (toposort inside toposort loop in get_rangeify_map)
+  def test_assign_scaling(self):
+    def assign_chain(n):
+      x = Tensor.empty(256, 256).realize()
+      for _ in range(n): x.assign(x + 1)
+      return x
+    self._assert_linear(assign_chain)
+
+  # layernorm has multi-consumer reduces (mean reused in variance), stresses consumer_rngs merge and symbolic rewrites
+  def test_layernorm_scaling(self):
+    def layernorm_chain(n):
+      x = Tensor.empty(64, 256)
+      for _ in range(n):
+        mean = x.mean(axis=-1, keepdim=True)
+        var = ((x - mean) ** 2).mean(axis=-1, keepdim=True)
+        x = (x - mean) / (var + 1e-5).sqrt()
+      return x
+    self._assert_linear(layernorm_chain)
+
+  # concat chain stresses MSTACK/MSELECT handling and wide SINK construction
+  def test_concat_scaling(self):
+    def concat_chain(n):
+      parts = [Tensor.empty(4, 256) + i for i in range(n)]
+      return parts[0].cat(*parts[1:])
+    self._assert_linear(concat_chain)
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)

tinygrad_repo/test/external/external_test_simple_tokenizer.py

@@ -0,0 +1,50 @@
+import functools, multiprocessing
+from transformers import AutoTokenizer
+from datasets import load_dataset
+from tinygrad.llm.cli import SimpleTokenizer
+from tinygrad.helpers import tqdm, getenv, partition
+
+@functools.cache
+def get_tokenizers():
+  print("getting tokenizers")
+  base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
+  special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()), lambda e: e[1] in base_tokenizer.all_special_ids)
+  simple_tokenizer = SimpleTokenizer(dict(normal_tokens), dict(special_tokens))
+  return base_tokenizer, simple_tokenizer
+
+def test_tokenize(samp) -> bool:
+  base_tokenizer, simple_tokenizer = get_tokenizers()
+  idx, txt = samp
+  try: simple_tokens = tuple(simple_tokenizer.encode(txt))
+  except RuntimeError: simple_tokens = ()
+  base_tokens = tuple(base_tokenizer.encode(txt, add_special_tokens=False))
+  if simple_tokens != base_tokens:
+    print(f"tokens mismatch at index: {idx}.\n")
+    color_codes = [91, 92, 94, 93, 95]
+    def color_tokens(tids):
+      return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([t])}" for i, t in enumerate(tids)) + "\033[0m"
+    print("simple:  ", color_tokens(simple_tokens))
+    print("official:", color_tokens(base_tokens) + "\n")
+    return False
+  if simple_tokenizer.decode(simple_tokens) != txt:
+    print(f"decode mismatch at {idx}")
+    return False
+  return True
+
+# use ALLOW_FAILED=-1 to go over the entire dataset without printing.
+if __name__ == "__main__":
+  print("loading datasets")
+  ds = load_dataset("OpenAssistant/oasst1")
+  loaded_ds = [(idx, el["text"]) for idx, el in enumerate(ds["train"])]
+  print(f"loaded {len(loaded_ds)}")
+
+  allow_failed = getenv("ALLOW_FAILED", 10)
+  fail_count, total = 0, 0
+  with multiprocessing.Pool(16) as pool:
+    for good in tqdm(pool.imap_unordered(test_tokenize, loaded_ds), total=len(loaded_ds)):
+      total += 1
+      if not good:
+        fail_count += 1
+        allow_failed -= 1
+        if allow_failed == 0: break
+    print(f"{fail_count}/{total} samples are inconsistent with the official tokenizer.")

tinygrad_repo/test/external/external_test_speed_llama.py

@@ -0,0 +1,58 @@
+# NOTE: this only tests the speed of the LLaMA codegen, it doesn't actually run the net
+import unittest, time
+from examples.llama import Transformer, MODEL_PARAMS
+from tinygrad.tensor import Tensor
+from tinygrad import Device
+from tinygrad.nn.state import get_state_dict
+from tinygrad.device import Allocator, Compiled
+from tinygrad.codegen import to_program_cache
+from tinygrad.helpers import Profiling
+
+class FakeProgram:
+  def __init__(self, name:str, lib:bytes, *args, **kwargs): pass
+  def __call__(self, *bufs, global_size, local_size, vals=(), wait=False, **kw): pass
+
+class FakeAllocator(Allocator[Compiled]):
+  def _alloc(self, sz, options): return None
+  def _copyin(self, dest, src:memoryview): pass
+
+class TestLLaMASpeed(unittest.TestCase):
+  def test_llama_compile(self):
+    backup_program = Device[Device.DEFAULT].runtime
+    backup_allocator = Device[Device.DEFAULT].allocator
+    backup_compiler = Device[Device.DEFAULT].compiler.compile_cached
+    Device[Device.DEFAULT].runtime = FakeProgram
+    Device[Device.DEFAULT].allocator = FakeAllocator(Device.default)
+
+    print("testing llama python run time")
+    model = Transformer(**MODEL_PARAMS["1"]["7B"]["args"])
+    print("built model")
+    # assign fake tensors to the values
+    for v in get_state_dict(model).values(): v.assign(Tensor.empty(*v.shape, dtype=v.dtype))
+    print("assigned empty tensors, doing warmup")
+
+    def run_llama(st, empty_cache=True):
+      if empty_cache: to_program_cache.clear()
+      tms = [time.perf_counter()]
+      for i in range(5):
+        model(Tensor([[1,2,3,4]]), i).realize()
+        tms.append(time.perf_counter())
+      timings = [(tms[i+1]-tms[i])*1000 for i in range(len(tms)-1)]
+      print(f"{st:15s} mean runtime: {sum(timings)/len(timings):7.2f}ms, runs: ", ", ".join(f'{x:7.2f}' for x in timings))
+
+    run_llama("codegen(0)")
+    run_llama("codegen(1)")
+
+    # test no compiler use for this
+    Device[Device.DEFAULT].compiler.compile_cached = None
+    run_llama("methodcache", False)
+    with Profiling(sort='time', frac=0.1, fn="/tmp/llama.prof", ts=5):
+      run_llama("profile", False)
+
+    Device[Device.DEFAULT].runtime = backup_program
+    Device[Device.DEFAULT].allocator = backup_allocator
+    Device[Device.DEFAULT].compiler.compile_cached = backup_compiler
+
+if __name__ == '__main__':
+  TestLLaMASpeed().test_llama_compile()
+  #unittest.main()

tinygrad_repo/test/external/external_test_speed_theoretical.py

@@ -0,0 +1,48 @@
+import time, unittest
+from tinygrad import Tensor, TinyJit, Device, dtypes
+from tinygrad.helpers import getenv, GlobalCounters
+
+SZMAX = getenv("SZMAX", 10)
+SZMIN = min(SZMAX, getenv("SZMIN", 10))
+def _test(tcount, fxn, dtype=dtypes.float):
+  print(f"**** testing {fxn.__name__} {dtype}")
+  allgbs = []
+  for sz in range(SZMIN, SZMAX+1):
+    jfxn = TinyJit(fxn)
+    ts = [Tensor.zeros((2**sz)*1024*1024, dtype=dtype).contiguous().realize() for _ in range(tcount)]
+    tms = []
+    for _ in range(10):
+      ts = [(x+1).realize() for x in ts]
+      Device.default.synchronize()
+      GlobalCounters.global_ops = 0
+      GlobalCounters.global_mem = 0
+      st = time.perf_counter()
+      jfxn(*ts).nbytes()
+      Device.default.synchronize()
+      tms.append(time.perf_counter() - st)
+      ops, mem = GlobalCounters.global_ops, GlobalCounters.global_mem
+    gflops = ops*1e-9/min(tms)
+    gbs = mem*1e-9/min(tms)
+    print(f"{ts[0].nbytes()/(1024*1024):10.0f} MB, {min(tms)*1e3:6.2f} ms {gbs:10.2f} GB/s {gflops:10.2f} GFLOPS {str(ts[0].shape):20s}")
+    allgbs.append(gbs)
+  return max(allgbs)
+
+MEMBW = getenv("MEMBW", 10)
+class TestRamBandwidth(unittest.TestCase):
+  def test_add(self): self.assertGreater(_test(2, Tensor.add), MEMBW)
+  def test_exp(self): self.assertGreater(_test(1, Tensor.exp), MEMBW)
+  def test_sum(self): self.assertGreater(_test(1, Tensor.sum), MEMBW)
+
+# ratio between MEM and FLOPS < 1000
+# NOTE: On AMD, (x*x)+1 gets ~30 TFLOPS, (x*x)+3 gets ~60 TFLOPS
+def flopsmax(x):
+  for _ in range(500): x = (x*x)+3
+  return x
+
+class TestFlops(unittest.TestCase):
+  def test_flops_int8(self): _test(1, flopsmax, dtypes.int8)
+  def test_flops_fp16(self): _test(1, flopsmax, dtypes.half)
+  def test_flops_fp32(self): _test(1, flopsmax)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_tlsf.py

@@ -0,0 +1,78 @@
+import unittest
+from tinygrad.runtime.support.memory import TLSFAllocator
+
+class TestTLSFAllocator(unittest.TestCase):
+  def setUp(self):
+    self.allocator = TLSFAllocator(1024, block_size=16)
+
+  def test_basic_alloc_free(self):
+    addr1 = self.allocator.alloc(32)
+    self.assertEqual(addr1, 0)
+
+    addr2 = self.allocator.alloc(64)
+    self.assertEqual(addr2, 32)
+
+    self.allocator.free(addr1)
+    addr3 = self.allocator.alloc(32)
+    self.assertEqual(addr3, 0)
+
+  def test_block_size_alignment(self):
+    addr1 = self.allocator.alloc(20)
+    addr2 = self.allocator.alloc(35)
+
+    self.assertEqual(addr1 % 16, 0)
+    self.assertEqual(addr2 % 16, 0)
+
+  def test_merge_blocks(self):
+    addr1 = self.allocator.alloc(32)
+    addr2 = self.allocator.alloc(32)
+    self.allocator.alloc(32)
+
+    self.allocator.free(addr1)
+    self.allocator.free(addr2)
+    addr4 = self.allocator.alloc(64)
+    self.assertEqual(addr4, addr1)
+
+  def test_split_blocks(self):
+    addr1 = self.allocator.alloc(128)
+    self.allocator.free(addr1)
+
+    addr2 = self.allocator.alloc(32)
+    self.assertEqual(addr2, addr1)
+
+    addr3 = self.allocator.alloc(32)
+    self.assertEqual(addr3, addr1 + 32)
+
+  def test_out_of_memory(self):
+    with self.assertRaises(MemoryError):
+      self.allocator.alloc(2048)
+
+  def test_fragmentation_handling(self):
+    addrs = []
+    for _ in range(5):
+      addrs.append(self.allocator.alloc(32))
+
+    # Free alternate blocks
+    for i in range(0, len(addrs), 2):
+      self.allocator.free(addrs[i])
+
+  def test_custom_start_address(self):
+    allocator = TLSFAllocator(1024, start_addr=1000)
+    addr1 = allocator.alloc(32)
+    self.assertEqual(addr1, 1000)
+
+    addr2 = allocator.alloc(64)
+    self.assertEqual(addr2, 1032)
+
+  def test_block_tracking(self):
+    addr1 = self.allocator.alloc(32)
+    addr2 = self.allocator.alloc(64)
+
+    self.assertTrue(addr1 in [addr - self.allocator.start_addr for addr in self.allocator.blocks])
+    self.assertTrue(addr2 in [addr - self.allocator.start_addr for addr in self.allocator.blocks])
+
+    self.allocator.free(addr1)
+    self.assertTrue(addr1 in [addr - self.allocator.start_addr for addr in self.allocator.blocks])
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_usb_asm24.py

@@ -0,0 +1,73 @@
+import unittest, time
+from tinygrad.runtime.support.usb import ASM24Controller
+from tinygrad.helpers import Timing
+from tinygrad import Tensor, Device
+import numpy as np
+
+class TestASMController(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.ctrl = ASM24Controller()
+
+  def test_write_and_read(self):
+    base = 0xF000
+    data = b"hello!"
+    self.ctrl.write(base, data)
+    out = self.ctrl.read(base, len(data))
+    self.assertEqual(out, data)
+
+  def test_scsi_write_and_read_from_f000(self):
+    payload = bytes([0x5B]) * 4096
+    self.ctrl.scsi_write(payload, lba=0)
+    back = self.ctrl.read(0xF000, len(payload))
+    self.assertEqual(back, payload)
+
+  def test_scsi_write_speed_4k(self):
+    payload = bytes([0x5A]) * 4096
+    start = time.perf_counter()
+    self.ctrl.scsi_write(payload, lba=0)
+    dur_ms = (time.perf_counter() - start) * 1000
+    print(f"scsi_write 4K took {dur_ms:.3f} ms")
+
+  def test_read_speed_4k(self):
+    payload = bytes([0xA5]) * 4096
+    self.ctrl.write(0xF000, payload)
+    start = time.perf_counter()
+    out = self.ctrl.read(0xF000, 4096)
+    dur_ms = (time.perf_counter() - start) * 1000
+    print(f"read 4K took {dur_ms:.3f} ms")
+    self.assertEqual(out, payload)
+
+class TestDevCopySpeeds(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.sz = 512
+    cls.dev = Device["AMD"]
+    if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
+
+  def testCopyCPUtoDefault(self):
+    for _ in range(10):
+      t = Tensor.ones(self.sz, self.sz, device="CPU").contiguous().realize()
+      with Timing(f"copyin of {t.nbytes()/1e6:.2f} MB:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"): # noqa: F821
+        t.to(Device.DEFAULT).realize()
+        Device[Device.DEFAULT].synchronize()
+      del t
+
+  def testCopyDefaulttoCPU(self):
+    t = Tensor.ones(self.sz, self.sz).contiguous().realize()
+    for _ in range(10):
+      with Timing(f"copyout of {t.nbytes()/1e6:.2f} MB:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"):
+        t.to('CPU').realize()
+
+  def testValidateCopies(self):
+    t = Tensor.randn(self.sz, self.sz, device="CPU").contiguous().realize()
+    x = t.to(Device.DEFAULT).realize()
+    Device[Device.DEFAULT].synchronize()
+
+    y = x.to('CPU').realize()
+
+    np.testing.assert_equal(t.numpy(), y.numpy())
+    del x, y, t
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/external/external_test_whisper_librispeech.py

@@ -0,0 +1,83 @@
+import unittest
+import torch
+import tqdm
+import torchaudio
+import pathlib
+import jiwer
+import os
+import numpy as np
+from whisper.normalizers import EnglishTextNormalizer
+from examples.whisper import init_whisper, transcribe_waveform
+
+class TestWhisperLibriSpeech(unittest.TestCase):
+  # reference WERs determined by running https://github.com/openai/whisper/blob/main/notebooks/LibriSpeech.ipynb
+  # the values should be consistent with the paper D.1.1 https://cdn.openai.com/papers/whisper.pdf#page=22
+  # tinygrad WERs do not perfectly match due to what seem to be precision differences vs torch
+  def test_en_tiny(self):
+    run_evaluation("tiny.en", 0.056629001883239174, 0.05655609406528749)
+
+  def test_tiny(self):
+    run_evaluation("tiny", 0.0771121409407306, 0.07558413638335187)
+
+  def test_en_base(self):
+    run_evaluation("base.en", 0.041412520064205455, 0.04271408904897505)
+
+  def test_en_small(self):
+    run_evaluation("small.en", 0.03369011117172363, 0.030531615969223228)
+
+def run_evaluation(model_name, tinygrad_expected_wer, reference_wer):
+  dataset = LibriSpeech()
+  batch_size=16
+  loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
+
+  model, enc = init_whisper(model_name, batch_size=batch_size)
+
+  hypotheses = []
+  references = []
+
+  for audio, texts in tqdm.tqdm(loader):
+    transcriptions = transcribe_waveform(model, enc, audio.numpy(), truncate=True)
+    hypotheses.extend(transcriptions)
+    references.extend(texts)
+
+  normalizer = EnglishTextNormalizer()
+  normalized_hypotheses = [normalizer(text) for text in hypotheses]
+  normalized_references = [normalizer(text) for text in references]
+  wer = jiwer.wer(normalized_hypotheses, normalized_references)
+
+  np.testing.assert_almost_equal(wer, tinygrad_expected_wer)
+  print(f'tinygrad WER {wer} vs reference WER {reference_wer}')
+  del model, enc
+
+class LibriSpeech(torch.utils.data.Dataset):
+  def __init__(self):
+    folder = pathlib.Path(__file__).parent.parent.parent / "extra" / "datasets" / "librispeech"
+    if not os.path.exists(folder):
+      os.makedirs(folder)
+
+    self.dataset = torchaudio.datasets.LIBRISPEECH(
+      root=folder,
+      url="test-clean",
+      download=True,
+    )
+
+  def __len__(self):
+    return len(self.dataset)
+
+  def __getitem__(self, item):
+    audio, sample_rate, text, _, _, _ = self.dataset[item]
+    assert sample_rate == 16000
+    return pad_or_trim_tensor(audio[0]), text
+
+def pad_or_trim_tensor(tensor, target_len=480000):
+  curr_len = len(tensor)
+  if curr_len == target_len:
+    return tensor
+  elif curr_len < target_len:
+    return torch.cat((tensor, torch.zeros(target_len - curr_len)))
+  else:
+    return tensor[:target_len]
+
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_test_yolov8.py

@@ -0,0 +1,73 @@
+import numpy as np
+from examples.yolov8 import YOLOv8, get_variant_multiples, preprocess, label_predictions, postprocess
+from tinygrad import Tensor
+import unittest
+import io, cv2
+import onnxruntime as ort
+import ultralytics
+from tinygrad.nn.state import safe_load, load_state_dict
+from tinygrad.helpers import fetch
+
+class TestYOLOv8(unittest.TestCase):
+  def test_all_load_weights(self):
+    for variant in ['n', 's', 'm', 'l', 'x']:
+      depth, width, ratio = get_variant_multiples(variant)
+      TinyYolov8 = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)
+      state_dict = safe_load(fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{variant}.safetensors'))
+      load_state_dict(TinyYolov8, state_dict)
+      print(f'successfully loaded weights for yolov{variant}')
+
+  def test_predictions(self):
+    test_image_urls = ['https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/bus.jpg', 'https://www.aljazeera.com/wp-content/uploads/2022/10/2022-04-28T192650Z_1186456067_UP1EI4S1I0P14_RTRMADP_3_SOCCER-ENGLAND-MUN-CHE-REPORT.jpg']
+    variant = 'n'
+    depth, width, ratio = get_variant_multiples(variant)
+    TinyYolov8 = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)
+    state_dict = safe_load(fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{variant}.safetensors'))
+    load_state_dict(TinyYolov8, state_dict)
+
+    for i in range(len(test_image_urls)):
+      img = cv2.imdecode(np.frombuffer(fetch(test_image_urls[i]).read_bytes(), np.uint8), 1)
+      test_image = preprocess([img])
+      predictions = TinyYolov8(test_image)
+      labels = label_predictions(predictions.numpy())
+      assert labels == {5: 1, 0: 4, 11: 1} if i == 0 else labels == {0: 12, 29: 1, 32: 1}
+
+  def test_forward_pass_torch_onnx(self):
+    variant = 'n'
+    weights_location = fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{variant}.safetensors')
+    weights_location_pt = fetch(f'https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8{variant}.pt', name=f"yolov8{variant}.pt")   # it needs the pt extension  # noqa: E501
+    weights_location_onnx = weights_location_pt.parent / f"yolov8{variant}.onnx"
+
+    # the ultralytics export prints a lot of unneccesary things
+    if not weights_location_onnx.is_file():
+      model = ultralytics.YOLO(model=weights_location_pt, task='Detect')
+      model.export(format="onnx",imgsz=[640, 480])
+
+    depth, width, ratio = get_variant_multiples(variant)
+    TinyYolov8 = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)
+    state_dict = safe_load(weights_location)
+    load_state_dict(TinyYolov8, state_dict)
+
+    image_location = [np.frombuffer(io.BytesIO(fetch('https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/bus.jpg').read_bytes()).read(), np.uint8)]  # noqa: E501
+    orig_image = [cv2.imdecode(image_location[0], 1)]
+
+    input_image = preprocess(orig_image)
+
+    onnx_session = ort.InferenceSession(weights_location_onnx)
+    onnx_input_name = onnx_session.get_inputs()[0].name
+    onnx_output_name = onnx_session.get_outputs()[0].name
+    onnx_output = onnx_session.run([onnx_output_name], {onnx_input_name: input_image.numpy()})
+
+    tiny_output = TinyYolov8(input_image).numpy()
+    onnx_output = postprocess(Tensor(onnx_output[0])).numpy()
+    #invalid boxes are multiplied by zero in postprocess
+    onnx_output = onnx_output[onnx_output[:, 4] != 0]
+    tiny_output = tiny_output[tiny_output[:, 4] != 0]
+
+    # currently rtol is 0.025 because there is a 1-2% difference in our predictions
+    # because of the zero padding in SPPF module (line 280) maxpooling layers rather than the -infinity in torch.
+    # This difference does not make a difference "visually".
+    np.testing.assert_allclose(onnx_output, tiny_output, atol=5e-4, rtol=0.025)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/external_uop_gc.py

@@ -0,0 +1,85 @@
+import gc
+from tinygrad import Tensor, UOp, Device, nn
+from tinygrad.schedule import schedule_cache
+from tinygrad.codegen import to_program, to_program_cache
+from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape
+from tinygrad.uop.divandmod import fold_divmod_general
+from test.test_tiny import TestTiny
+
+def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
+def print_uops():
+  for x in gc.get_objects():
+    if isinstance(x, UOp): print(x)
+
+def start(): pass
+def single_tensor(): Tensor([2])
+def two_plus_two(): Tensor([2])+Tensor([2])
+def two_plus_two_schedule(): (Tensor([2])+Tensor([2])).schedule_linear()
+def two_plus_two_kernel():
+  linear = (Tensor([2])+Tensor([2])).schedule_linear()
+  to_program(linear.src[-1].src[0], Device.default.renderer)
+def two_plus_two_linearize():
+  linear = (Tensor([2])+Tensor([2])).schedule_linear()
+  to_program(linear.src[-1].src[0], Device.default.renderer)
+def two_plus_two_realize(): (Tensor([2])+Tensor([2])).realize()
+def two_plus_two_item(): (Tensor([2])+Tensor([2])).item()
+def gradient_test():
+  x = Tensor.eye(3)
+  y = Tensor([[2.0,0,-2.0]])
+  z = y.matmul(x).sum()
+  z.backward()
+def realized_eye():
+  Tensor.eye(3).clone().realize()
+def realized_list():
+  Tensor([[2.0,0,-2.0]]).realize()
+def kernel_matmul():
+  x = Tensor.eye(3)
+  y = Tensor([[2.0,0,-2.0]])
+  z = y.matmul(x)
+  linear = z.schedule_linear()
+  to_program(linear.src[-1].src[0], Device.default.renderer)
+def realized_matmul():
+  x = Tensor.eye(3)
+  y = Tensor([[2.0,0,-2.0]])
+  z = y.matmul(x)
+  Tensor.realize(z)
+def realized_gradient():
+  x = Tensor.eye(3)
+  y = Tensor([[2.0,0,-2.0]])
+  z = y.matmul(x).sum()
+  z.backward()
+  Tensor.realize(x, y, z, x.grad, y.grad)
+def nn_batchnorm(): nn.BatchNorm(64)
+def nn_conv2d(): nn.Conv2d(64, 64, 3)
+def plus(): TestTiny().test_plus()
+def mnist(): TestTiny().test_mnist()
+def mnist_backward(): TestTiny().test_mnist_backward()
+
+tests = [start, single_tensor, two_plus_two, two_plus_two_schedule, two_plus_two_kernel,
+         two_plus_two_linearize, two_plus_two_realize, two_plus_two_item, gradient_test,
+         realized_eye, realized_list, kernel_matmul, realized_matmul, realized_gradient,
+         nn_batchnorm, nn_conv2d, plus, mnist, mnist_backward]
+
+if __name__ == "__main__":
+  gc.disable()
+  start_uops = uops_allocated()
+  # there's a few consts created as default values
+  print_uops()
+  for t in tests:
+    t()
+
+    # these caches will keep uops alive
+    schedule_cache.clear()
+    to_program_cache.clear()
+    apply_movement_op.cache_clear()
+    _apply_reshape.cache_clear()
+    fold_divmod_general.cache_clear()
+    Tensor._device_seeds.clear()
+    Tensor._device_rng_counters.clear()
+
+    new_uops = uops_allocated()
+    gc.collect()
+    new_uops_gc = uops_allocated()
+    print(f"{t.__name__:30s}: {new_uops:3d} -> {new_uops_gc:3d}")
+    if new_uops != start_uops: print_uops()
+    assert new_uops == start_uops

tinygrad_repo/test/external/fuzz_fast_idiv.py

@@ -0,0 +1,25 @@
+import random
+import z3
+from tinygrad import dtypes, Device
+from tinygrad.uop.validate import uops_to_z3, z3_cdiv
+from tinygrad.uop.ops import UOp
+from tinygrad.uop.decompositions import fast_idiv
+random.seed(42)
+
+powers_of_two = [2**i for i in range(64)]
+if __name__ == "__main__":
+  for i in range(10_000):
+    if i % 1000 == 0:
+      print(f"Progress: {i}")
+    dt = random.choice(dtypes.ints + tuple(dt.vec(4) for dt in dtypes.ints))
+    u = UOp.variable('x', random.randint(dt.min, 0), random.randint(1, dt.max), dtype=dt)
+    d = random.randint(1, max(1, u.arg[2])*2)
+    if d in powers_of_two: continue
+    expr = fast_idiv(Device[Device.DEFAULT].renderer, u, d)
+    if expr is None: continue
+
+    solver = z3.Solver()
+    z3_expr, x =uops_to_z3(solver, expr, u)
+
+    if solver.check(z3_expr != z3_cdiv(x, d)) == z3.sat:
+      assert False, f"Failed: {expr.render()} != x//{d} at x={solver.model()}\nx={u}\nd={d}\n{z3_expr=}\n{x/d=}"

tinygrad_repo/test/external/fuzz_graph.py

@@ -0,0 +1,129 @@
+import random, ctypes
+import numpy as np
+from tinygrad.device import Buffer, Device
+from tinygrad.helpers import Context, getenv, from_mv
+from tinygrad.dtype import dtypes
+from tinygrad.tensor import Tensor, _to_np_dtype
+from tinygrad.engine.realize import BufferXfer, get_runner, ExecItem
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.engine.jit import apply_graph_to_jit
+
+BUF_LEN = getenv("BUF_LEN", 128)
+
+cached_prgs = {}
+def gen_prg(device, inputs_cnt):
+  if (device, inputs_cnt) in cached_prgs: return cached_prgs[(device, inputs_cnt)]
+
+  with Context(DEBUG=0):
+    fst = [Tensor.randn(BUF_LEN, dtype=dtypes.int).realize() for i in range(inputs_cnt)]
+    s = fst[0]
+    for i in range(1, inputs_cnt): s = s.bitwise_xor(fst[i])
+
+    linear = s.schedule_linear()
+    prg = get_runner(device, linear.src[-1].src[0])
+  cached_prgs[(device, inputs_cnt)] = prg
+  return prg
+
+def alloc_rawbuffer(device, fill=False):
+  rawbuf = Buffer(device, BUF_LEN, dtypes.int).ensure_allocated()
+  if fill:
+    with Context(DEBUG=0):
+      data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
+      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_memoryview())
+  return rawbuf
+
+def gen_kernel_ji(device, deps):
+  assert len(deps) >= 2
+  out = alloc_rawbuffer(device)
+  prg = gen_prg(device, len(deps))
+  return ExecItem(UOp(Ops.NOOP), [out] + deps, prg=prg)
+
+def gen_copy_ji(device, deps):
+  assert len(deps) == 1
+  out = alloc_rawbuffer(device)
+  prg = BufferXfer(deps[0].nbytes, device, deps[0].device)
+  return ExecItem(UOp(Ops.NOOP), [out] + deps, prg=prg)
+
+def gen_graph():
+  input_buffers = []
+  all_buffers = []
+  jis = []
+
+  last_n_deps = getenv("LAST_N_DEPS", 0)
+
+  kernel_count = random.randint(2, getenv("MAX_KERNELS", 128))
+  for i in range(kernel_count):
+    target_device_id = random.randint(0, getenv("MAX_DEVICES", 6) - 1)
+    target_device = f"{Device.DEFAULT}:{target_device_id}"
+    is_copy = random.randint(0, 10) < 3
+
+    if is_copy:
+      deps_pool = [buf for buf in all_buffers[-last_n_deps:] if buf.device != target_device]
+      if len(deps_pool) == 0: deps = []
+      else: deps = random.sample(deps_pool, 1)
+    else:
+      deps_pool = [buf for buf in all_buffers[-last_n_deps:] if buf.device == target_device]
+      deps_count = random.randint(0, min(getenv("MAX_DEPS_COUNT", 6), len(deps_pool)))
+      if deps_count == 0: deps = []
+      else: deps = random.sample(deps_pool, deps_count)
+
+    if len(deps) == 0 or (not is_copy and len(deps) < 2):
+      buf = alloc_rawbuffer(target_device, fill=True)
+      input_buffers.append(buf)
+      all_buffers.append(buf)
+    elif is_copy:
+      jis.append(gen_copy_ji(target_device, deps))
+      all_buffers.append(jis[-1].bufs[0])
+    else:
+      jis.append(gen_kernel_ji(target_device, deps))
+      all_buffers.append(jis[-1].bufs[0])
+
+  return jis, all_buffers, input_buffers
+
+def run_jit(jis, all_buffers, input_buffers, var_vals):
+  with Context(DEBUG=0):
+    for rawbuf in all_buffers:
+      if rawbuf in input_buffers: continue
+      mv = memoryview(bytearray(rawbuf.nbytes))
+      ctypes.memset(from_mv(mv), 0, len(mv))
+      rawbuf.copyin(mv)
+
+  for ei in jis: ei.run(var_vals, jit=True)
+
+  with Context(DEBUG=0):
+    res_buffers = []
+    for rawbuf in all_buffers: res_buffers.append(rawbuf.as_memoryview())
+    return res_buffers
+
+def fuzz_graph(jis, all_buffers, input_buffers):
+  ground_thruth_bufs = run_jit(jis, input_buffers, all_buffers, {})
+  ground_truth_np = [np.frombuffer(x, _to_np_dtype(all_buffers[i].dtype)) for i,x in enumerate(ground_thruth_bufs)]
+
+  for _ in range(getenv("FUZZ_GRAPH_SPLIT_RUNS", 64)):
+    max_split_points = len(jis) // 3
+    split_points = random.randint(0, min(max_split_points, getenv("FUZZ_GRAPH_MAX_SPLITS", 8)))
+    split = [0]
+    for i in range(split_points - 1):
+      split.append(random.randint(split[-1] + 2, len(jis) - 2 * (max_split_points - i)))
+    split.append(len(jis))
+
+    graphed_jit = []
+    for sp in range(len(split)-1):
+      graphed_jit += apply_graph_to_jit(jis[split[sp]:split[sp+1]], [], {})
+
+    for _ in range(getenv("FUZZ_GRAPH_SPLIT_RETRY_RUNS", 4)):
+      test_bufs = run_jit(graphed_jit, input_buffers, all_buffers, {})
+      test_bufs_np = [np.frombuffer(x, _to_np_dtype(all_buffers[i].dtype)) for i,x in enumerate(test_bufs)]
+      for i in range(len(ground_thruth_bufs)): np.testing.assert_equal(ground_truth_np[i], test_bufs_np[i])
+
+if __name__ == "__main__":
+  SEED = getenv("SEED", 42)
+  random.seed(SEED)
+  np.random.seed(SEED)
+
+  next_graph_id = 0
+  for i in range(getenv("ITERS", 1000)):
+    print("Running graph", next_graph_id)
+    jis, all_buffers, input_buffers = gen_graph()
+    fuzz_graph(jis, all_buffers, input_buffers)
+    next_graph_id += 1

tinygrad_repo/test/external/fuzz_kfd.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+import random
+from tqdm import trange
+from typing import List
+from tinygrad import Device
+from tinygrad.runtime.ops_amd import AMDDevice, HWQueue
+
+if __name__ == "__main__":
+  dev: List[AMDDevice] = [Device[f"KFD:{i}"] for i in range(6)]
+  print(f"got {len(dev)} devices")
+
+  buffers = [(rd:=random.choice(dev), rd.allocator.alloc(random.randint(1, 10000))) for i in range(100)]
+
+  for _ in trange(100000):
+    d1, b1 = random.choice(buffers)
+    d2, b2 = random.choice(buffers)
+    d1._gpu_map(b2)
+    q = HWQueue()
+    q.signal(sig:=AMDDevice._alloc_signal(10))
+    qc = HWQueue()
+    qc.wait(sig)
+    qc.copy(b1, b2, min(b1.size, b2.size))
+    d1.completion_signal.value = 1
+    qc.signal(d1.completion_signal)
+    qc.submit(d1)
+    q.wait(d1.completion_signal)
+    q.submit(d1)
+    AMDDevice._wait_on(d1.completion_signal.event_id)

tinygrad_repo/test/external/fuzz_shape_ops.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+import unittest
+from math import prod
+
+from hypothesis import assume, given, settings, strategies as st
+from hypothesis.extra import numpy as stn
+
+import numpy as np
+import torch
+from tinygrad import Tensor
+from tinygrad.helpers import getenv
+
+
+settings.register_profile(__file__, settings.default,
+                          max_examples=100, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+
+
+# torch wraparound for large numbers
+st_int32 = st.integers(-2147483648, 2147483647)
+
+@st.composite
+def st_shape(draw) -> tuple[int, ...]:
+  s = draw(stn.array_shapes(min_dims=0, max_dims=6,
+                            min_side=0, max_side=128))
+  assume(prod(s) <= 1024 ** 2)
+  assume(prod([d for d in s if d]) <= 1024 ** 4)
+  return s
+
+def tensors_for_shape(s:tuple[int, ...]) -> tuple[torch.tensor, Tensor]:
+  x = np.arange(prod(s)).reshape(s)
+  return torch.from_numpy(x), Tensor(x)
+
+def apply(tor, ten, tor_fn, ten_fn=None):
+  ok = True
+  try: tor = tor_fn(tor)
+  except: tor, ok = None, not ok  # noqa: E722
+  try: ten = ten_fn(ten) if ten_fn is not None else tor_fn(ten)
+  except: ten, ok = None, not ok  # noqa: E722
+  return tor, ten, ok
+
+class TestShapeOps(unittest.TestCase):
+  @settings.get_profile(__file__)
+  @given(st_shape(), st_int32, st.one_of(st_int32, st.lists(st_int32)))
+  def test_split(self, s:tuple[int, ...], dim:int, sizes:int|list[int]):
+    tor, ten = tensors_for_shape(s)
+    tor, ten, ok = apply(tor, ten, lambda t: t.split(sizes, dim))
+    assert ok
+    if tor is None and ten is None: return
+
+    assert len(tor) == len(ten)
+    assert all([np.array_equal(tor.numpy(), ten.numpy()) for (tor, ten) in zip(tor, ten)])
+
+  @settings.get_profile(__file__)
+  @given(st_shape(), st_int32, st_int32)
+  def test_chunk(self, s:tuple[int, ...], dim:int, num:int):
+    # chunking on a 0 dim is cloning and leads to OOM if done unbounded.
+    assume((0 <= (actual_dim := len(s)-dim if dim < 0 else dim) < len(s) and s[actual_dim] > 0) or
+           (num < 16))
+
+    tor, ten = tensors_for_shape(s)
+    tor, ten, ok = apply(tor, ten, lambda t: t.chunk(num, dim))
+    assert ok
+    if tor is None and ten is None: return
+
+    assert len(tor) == len(ten)
+    assert all([np.array_equal(tor.numpy(), ten.numpy()) for (tor, ten) in zip(tor, ten)])
+
+  @settings.get_profile(__file__)
+  @given(st_shape(), st_int32)
+  def test_squeeze(self, s:tuple[int, ...], dim:int):
+    tor, ten = tensors_for_shape(s)
+    tor, ten, ok = apply(tor, ten, lambda t: t.squeeze(dim))
+    assert ok
+    if tor is None and ten is None: return
+    assert np.array_equal(tor.numpy(), ten.numpy())
+
+  @settings.get_profile(__file__)
+  @given(st_shape(), st_int32)
+  def test_unsqueeze(self, s:tuple[int, ...], dim:int):
+    tor, ten = tensors_for_shape(s)
+    tor, ten, ok = apply(tor, ten, lambda t: t.unsqueeze(dim))
+    assert ok
+    if tor is None and ten is None: return
+    assert np.array_equal(tor.numpy(), ten.numpy())
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/external/fuzz_symbolic.py

@@ -0,0 +1,92 @@
+# NOTE: z3-solver 4.15.4 segfaults (exit code 139) when creating many z3.Context() with complex expressions.
+# Reproduces consistently with seed=74 around iteration 1767. Versions <=4.15.3 are fine.
+# Workaround: reuse a single z3.Context, or pin z3-solver<4.15.4 (see pyproject.toml).
+# To repro: pip install z3-solver==4.15.4.0 && python test/external/fuzz_symbolic.py 74
+import random, operator, sys
+import z3
+from tinygrad import Variable, dtypes
+from tinygrad.uop.ops import UOp
+from tinygrad.uop.validate import uops_to_z3
+from tinygrad.helpers import DEBUG
+
+seed = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
+print(f"Seed: {seed}", flush=True)
+random.seed(seed)
+
+unary_ops = [lambda a:a+random.randint(-4, 4), lambda a: a*random.randint(-4, 4),
+            lambda a: a//random.randint(1, 9), lambda a: a%random.randint(1, 9),
+            lambda a:a.maximum(random.randint(-10, 10)), lambda a:a.minimum(random.randint(-10, 10))]
+binary_ops = [lambda a,b: a+b, lambda a,b: a*b, lambda a,b:a.maximum(b), lambda a,b:a.minimum(b)]
+comp_ops = [operator.lt, operator.le, operator.gt, operator.ge]
+
+def random_or_sub_expression_int(depth, expr):
+  sub_expr = random.choice([e for e in expr.toposort() if e.dtype is not dtypes.bool])
+  return random.choice([random_int_expr(depth-1), sub_expr])
+
+def random_int_expr(depth=10):
+  if depth <= 0: return random.choice(v)
+  expr1 = random_int_expr(depth-1)
+
+  # we give more weight to arithmatic ops than to minimum and maximum
+  ops = [
+    lambda: random.choices(unary_ops, weights=[4, 4, 4, 4, 1, 1])[0](expr1),
+    # for the second operand its either another random exprssion or some subexpression of the first operand
+    lambda: random.choices(binary_ops, [8, 1, 1, 1])[0](expr1, random_or_sub_expression_int(depth-1, expr1)),
+    lambda: random_bool_expr(3, random_or_sub_expression_int(depth-1, expr1)).where(expr1, random_or_sub_expression_int(depth-1, expr1)),
+  ]
+  # we give weight proportional to the amount of ops in each branch
+  return random.choices(ops, weights=[6, 4, 1])[0]()
+
+def random_bool_expr(depth=10, expr1=None):
+  if depth == 0: return True
+  if expr1 is None: expr1 = random_int_expr(depth-1)
+  expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(dtypes.int, random.randint(-10, 10))])
+  return random.choice(comp_ops)(expr1, expr2)
+
+
+if __name__ == "__main__":
+  skipped = 0
+  for i in range(10000):
+    if i % 1000 == 0:
+      print(f"Running test {i}")
+    upper_bounds = [*list(range(1, 10)), 16, 32, 64, 128, 256]
+    u1 = Variable("v1", 0, random.choice(upper_bounds))
+    u2 = Variable("v2", 0, random.choice(upper_bounds))
+    u3 = Variable("v3", 0, random.choice(upper_bounds))
+    v = [u1,u2,u3]
+    expr = random_int_expr(6)
+
+    simplified_expr = expr.simplify()
+
+    solver = z3.Solver(ctx=z3.Context())
+    solver.set(timeout=5000)  # some expressions take very long verify, but its very unlikely they actually return sat
+    z3_expr, z3_simplified_expr, v1, v2, v3 = uops_to_z3(solver, expr, simplified_expr, u1, u2, u3)
+    check = solver.check(z3_simplified_expr != z3_expr)
+    if check == z3.unknown and DEBUG>=1:
+      skipped += 1
+      print("Skipped due to timeout or interrupt:\n" +
+            f"v1=Variable(\"{u1.arg[0]}\", {u1.arg[1]}, {u1.arg[2]})\n" +
+            f"v2=Variable(\"{u2.arg[0]}\", {u2.arg[1]}, {u2.arg[2]})\n" +
+            f"v3=Variable(\"{u3.arg[0]}\", {u3.arg[1]}, {u3.arg[2]})\n" +
+            f"expr = {expr.render(simplify=False)}\n")
+    elif check == z3.sat:
+      m = solver.model()
+      n1, n2, n3 = m[v1], m[v2], m[v3]
+      u1_val, u2_val, u3_val = u1.const_like(n1.as_long()), u2.const_like(n2.as_long()), u3.const_like(n3.as_long())
+      num = expr.simplify().substitute({u1:u1_val, u2:u2_val, u3:u3_val}).ssimplify()
+      rn = expr.substitute({u1:u1_val, u2:u2_val, u3:u3_val}).ssimplify()
+      if num==rn: print("z3 found a mismatch but the expressions are equal!!")
+      assert False, f"mismatched {expr.render()} at v1={m[v1]}; v2={m[v2]}; v3={m[v3]} = {num} != {rn}\n" +\
+            "Reproduce with:\n" +\
+            f"v1=Variable(\"{u1.arg[0]}\", {u1.arg[1]}, {u1.arg[2]})\n" +\
+            f"v2=Variable(\"{u2.arg[0]}\", {u2.arg[1]}, {u2.arg[2]})\n" +\
+            f"v3=Variable(\"{u3.arg[0]}\", {u3.arg[1]}, {u3.arg[2]})\n" +\
+            f"expr = {expr}\n" +\
+            f"v1_val, v2_val, v3_val = UOp.const(dtypes.int, {n1.as_long()}), UOp.const(dtypes.int, {n2.as_long()})," +\
+                f"UOp.const(dtypes.int, {n3.as_long()})\n" +\
+            "num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\
+            "rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\
+            "assert num==rn, f\"{num} != {rn}\"\n"
+
+    if DEBUG >= 2: print(f"validated {expr.render()}")
+  print(f"Skipped {skipped} expressions due to timeout")

tinygrad_repo/test/external/fuzz_symbolic_symbolic_div.py

@@ -0,0 +1,65 @@
+import random, sys
+import z3
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.uop.validate import uops_to_z3
+from tinygrad.helpers import DEBUG, colored
+
+seed = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
+print(f"Seed: {seed}", flush=True)
+random.seed(seed)
+
+def get_random_term(ranges, factors):
+  # 10% chance of nesting
+  if random.randint(0,9) == 0: return get_random_expr(ranges, factors)
+  return random.choice(ranges)*random.choice(factors)*random.choice([1, 1, 1, -1])
+
+def get_random_expr(ranges, factors):
+  num_terms = random.randint(2,4)
+  x = UOp.usum(*[get_random_term(ranges, factors) for _ in range(num_terms)])
+  return x.alu(random.choice([Ops.CDIV, Ops.CMOD]), x.ufix(random.choice(factors)*random.choice([1, 1, 1, -1])))
+
+if __name__ == "__main__":
+  skipped = 0
+  for i in range(700):
+    if i % 100 == 0:
+      print(f"Running test {i}")
+    upper_bounds = [*list(range(1, 4)), 16, 33, 53, 64, 256]
+    variable_names = ["i", "j", "k"]
+    variables = [UOp.variable(s, 1, random.choice(upper_bounds)) for s in variable_names]
+    factors = variables+upper_bounds
+    # add some products
+    for _ in range(2): factors.append(random.choice(variables)*random.choice(variables))
+    # add some adds
+    for _ in range(2): factors.append(random.choice(variables)+random.choice(factors))
+    num_ranges = 4
+    ranges = [UOp.range(random.choice(factors), i) for i in range(num_ranges)]
+    variable_names += [f"r{i}" for i in range(num_ranges)]
+    expr = get_random_expr(ranges, factors)
+
+    simplified_expr = expr.simplify()
+
+    if DEBUG>=1:
+      print(expr.render(simplify=False), "  -->  ", simplified_expr.render(simplify=False))
+
+    solver = z3.Solver()
+    solver.set(timeout=1000)  # some expressions take very long verify, but its very unlikely they actually return sat
+    z3_expr, z3_simplified_expr, *z3_vars = uops_to_z3(solver, expr, simplified_expr, *variables, *ranges)
+    check = solver.check(z3_simplified_expr != z3_expr)
+    if check == z3.unknown:
+      skipped += 1
+      if DEBUG>=1: print("skipped z3 verification due to timeout")
+    elif check == z3.sat:
+      print(colored("simplify INCORRECT!", "red"))
+      print(solver.model())
+      var_vals = {s:solver.model()[z] for s,z in zip(variable_names, z3_vars)}
+      print("reproduce with:")
+      print("var_vals = ", var_vals)
+      print("globals = var_vals|{'cdiv':cdiv,'cmod':cmod}")
+      print("expr = ast.simplify()")
+      print("assert eval(ast.render(pm=renderer_infer, simplify=False),globals) == eval(expr.render(pm=renderer_infer, simplify=False),globals)")
+      print()
+
+      assert False
+
+    if DEBUG >= 2: print(f"validated {expr.render()}")
+  print(f"Skipped {skipped} expressions due to timeout")

tinygrad_repo/test/external/mlperf_bert/external_benchmark_bert.py

@@ -0,0 +1,102 @@
+import unittest, time
+
+from tinygrad import Tensor, TinyJit, GlobalCounters, Device
+from tinygrad.helpers import getenv, Context
+from tinygrad.nn.optim import LAMB
+from tinygrad.nn.state import get_parameters
+from tinygrad.engine.realize import run_linear
+
+from extra.models import bert
+
+bs = getenv("BS", 16)
+seq_len = getenv("SEQ_LEN", 512)
+
+class BenchmarkBertTrain(unittest.TestCase):
+  def _get_layer(self, layer_id):
+    if not hasattr(self, "model"):
+      dropout_prob = 0.0 if getenv("DISABLE_DROPOUT") else 0.1
+      self.model = bert.BertForPretraining(attention_probs_dropout_prob=dropout_prob, hidden_dropout_prob=dropout_prob)
+    hidden_size = self.model.bert.embeddings.word_embeddings.embed_sz
+    intermediate_size = self.model.bert.encoder.layer[0].intermediate.dense.weight.shape[0]
+
+    layer_map = {
+      "embedding": self.model.bert.embeddings,
+      "attention_self": self.model.bert.encoder.layer[0].attention.self,
+      "attention_output": self.model.bert.encoder.layer[0].attention.output,
+      "intermediate": self.model.bert.encoder.layer[0].intermediate,
+      "output": self.model.bert.encoder.layer[0].output
+    }
+    input_shapes = {
+      "embedding": [(bs, seq_len), (bs, seq_len)],
+      "attention_self": [(bs, seq_len, hidden_size), (bs, 1, 1, seq_len)],
+      "attention_output": [(bs, seq_len, hidden_size), (bs, seq_len, 1)],
+      "intermediate": [(bs, seq_len, hidden_size)],
+      "output": [(bs, seq_len, intermediate_size), (bs, seq_len, 1)]
+    }.get(layer_id)
+
+    return f"{layer_id}-layer, Input: {input_shapes}", layer_map.get(layer_id), input_shapes
+
+  def _test_layer(self, name, layer, input_shapes):
+    optim = LAMB(get_parameters(layer))
+    with Context(TRACK_MATCH_STATS=0): Tensor.realize(*[t.assign(t.detach().contiguous()) for t in get_parameters(optim)])
+
+    JITCNT = getenv("JITCNT", 1)
+    Tensor.training = True
+    @TinyJit
+    def step(inputs):
+      optim.zero_grad()
+      for i in inputs: i.grad = None
+
+      y = layer(*inputs).contiguous().contiguous_backward()
+      y.sum().backward()
+      if getenv("ASSIGN", 1): linear, var_vals = Tensor.linear_with_vars(y, *list(inputs), *optim.schedule_step())
+      else: linear, var_vals = Tensor.linear_with_vars(y, *list(inputs), *[t.grad for t in optim.params])
+
+      for _ in range(JITCNT):
+        run_linear(linear, var_vals)
+
+    CNT = getenv("CNT", 5)
+    best_tm = None
+    flops, mem_used, mem, kernels = None, None, None, None
+    for _ in range(CNT):
+      with Context(TRACK_MATCH_STATS=0): inputs = [Tensor.randn(*shape).realize() for shape in input_shapes]
+      GlobalCounters.reset()
+
+      st = time.perf_counter()
+      step(inputs)
+      Device[Device.DEFAULT].synchronize()
+      et = time.perf_counter()
+
+      flops = GlobalCounters.global_ops / JITCNT
+      mem_used = GlobalCounters.mem_used
+      mem = GlobalCounters.global_mem / JITCNT
+      if kernels is None: kernels = GlobalCounters.kernel_count // JITCNT
+      tm = (et-st) / JITCNT
+      if best_tm is None or tm < best_tm: best_tm = tm
+    print(f"\r{name:70s}: {best_tm * 1000:>9.2f} ms, {flops / 10**12 / best_tm:>6.2f} TFLOPS, {mem / 10**9 / best_tm:>5.0f} GB/s, "
+          f"{mem_used / 10**9: 6.2f} GB used, {kernels:>5d} kernels")
+    return best_tm, flops, mem, kernels
+
+  def test_embedding_layer(self): self._est(*self._test_layer(*self._get_layer("embedding")), 1)
+  def test_attention_self_layer(self): self._est(*self._test_layer(*self._get_layer("attention_self")), 24) # Assumes BERT-large
+  def test_attention_output_layer(self): self._est(*self._test_layer(*self._get_layer("attention_output")), 24)
+  def test_intermediate_layer(self): self._est(*self._test_layer(*self._get_layer("intermediate")), 24)
+  def test_output_layer(self): self._est(*self._test_layer(*self._get_layer("output")), 24)
+
+  est_tm, est_flops, est_mem, est_kernels = 0, 0, 0, 0
+
+  @classmethod
+  def _est(cls, tm, flops, mem, kernels, mult):
+    cls.est_tm += tm * mult
+    cls.est_flops += flops * mult
+    cls.est_mem += mem * mult
+    cls.est_kernels += kernels * mult
+
+  @classmethod
+  def tearDownClass(cls):
+    print(f"\restimated step tm: {cls.est_tm * 1000.0:.2f} ms, {cls.est_flops / 10 ** 12 / cls.est_tm:.3f} tflops, "
+          f"{cls.est_mem / 10 ** 9 / cls.est_tm:.2f} GB/s, {cls.est_kernels} kernels")
+
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/external/mlperf_bert/external_test_checkpoint_loading.py

@@ -0,0 +1,55 @@
+# Test whether pretrained weights from first BERT pretraining phase have been loaded correctly
+# Usage:
+# 1. Download the BERT checkoints with `wikipedia_download.py`
+#    Command: BASEDIR=/path/to/wiki python3 wikipedia_download.py
+# 2. Run this script. (Adjust EVAL_BS and GPUS as needed)
+#    Command: EVAL_BEAM=4 DEFAULT_FLOAT=half GPUS=6 BASEDIR=/path/to/wiki python3 test/external/mlperf_bert/external_test_checkpoint_loading.py
+import os
+from tqdm import tqdm
+
+from tinygrad.tensor import Tensor
+from tinygrad.device import Device
+from tinygrad.helpers import getenv
+from tinygrad.nn.state import get_state_dict
+
+from examples.mlperf.helpers import get_mlperf_bert_model, get_data_bert
+from examples.mlperf.dataloader import batch_load_val_bert
+from examples.mlperf.model_train import eval_step_bert
+
+if __name__ == "__main__":
+  BASEDIR = os.environ["BASEDIR"] = getenv("BASEDIR", "/raid/datasets/wiki")
+  INIT_CKPT_DIR = getenv("INIT_CKPT_DIR", BASEDIR)
+  GPUS = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 1))]
+  EVAL_BS = getenv("EVAL_BS", 4 * len(GPUS))
+  max_eval_steps = (10000 + EVAL_BS - 1) // EVAL_BS
+
+  for i in range(10):
+    assert os.path.exists(os.path.join(BASEDIR, "eval", f"{i}.pkl")), \
+      f"File {i}.pkl does not exist in {os.path.join(BASEDIR, 'eval')}"
+
+  required_files = ["checkpoint", "model.ckpt-28252.data-00000-of-00001", "model.ckpt-28252.index", "model.ckpt-28252.meta"]
+  assert all(os.path.exists(os.path.join(INIT_CKPT_DIR, f)) for f in required_files), \
+    f"Missing checkpoint files in INIT_CKPT_DIR: {required_files}"
+
+  Tensor.training = False
+
+  model = get_mlperf_bert_model(INIT_CKPT_DIR)
+
+  for _, x in get_state_dict(model).items():
+    x.realize().to_(GPUS)
+
+  eval_accuracy = []
+  eval_it = iter(batch_load_val_bert(EVAL_BS))
+
+  for _ in tqdm(range(max_eval_steps), desc="Evaluating", total=max_eval_steps):
+    eval_data = get_data_bert(GPUS, eval_it)
+    eval_result: dict[str, Tensor] = eval_step_bert(model, eval_data["input_ids"], eval_data["segment_ids"], eval_data["input_mask"], \
+                                               eval_data["masked_lm_positions"], eval_data["masked_lm_ids"], \
+                                               eval_data["masked_lm_weights"], eval_data["next_sentence_labels"])
+
+    mlm_accuracy = eval_result["masked_lm_accuracy"].numpy().item()
+    eval_accuracy.append(mlm_accuracy)
+
+  total_lm_accuracy = sum(eval_accuracy) / len(eval_accuracy)
+  assert total_lm_accuracy >= 0.34, "Checkpoint loaded incorrectly. Accuracy should be very close to 0.34085 as per MLPerf BERT README."
+  print(f"Checkpoint loaded correctly. Accuracy of {total_lm_accuracy*100:.3f}% achieved. (Reference: 34.085%)")

tinygrad_repo/test/external/mlperf_bert/preprocessing/create_pretraining_data.py

@@ -0,0 +1,435 @@
+# https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/cleanup_scripts/create_pretraining_data.py
+# NOTE: This is a direct copy of the original script
+# NOTE: With python 3.7.12, pip install tensorflow=1.15.5
+"""Create masked LM/next sentence masked_lm TF examples for BERT."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python' # NOTE: This is a workaround for protobuf issue
+
+import collections
+import random
+import tokenization
+import tensorflow as tf
+
+
+flags = tf.flags
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("input_file", None,
+                    "Input raw text file (or comma-separated list of files).")
+
+flags.DEFINE_string(
+    "output_file", None,
+    "Output TF example file (or comma-separated list of files).")
+
+flags.DEFINE_string("vocab_file", None,
+                    "The vocabulary file that the BERT model was trained on.")
+
+flags.DEFINE_bool(
+    "do_lower_case", True,
+    "Whether to lower case the input text. Should be True for uncased "
+    "models and False for cased models.")
+
+flags.DEFINE_integer("max_seq_length", 128, "Maximum sequence length.")
+
+flags.DEFINE_integer("max_predictions_per_seq", 20,
+                     "Maximum number of masked LM predictions per sequence.")
+
+flags.DEFINE_integer("random_seed", 12345, "Random seed for data generation.")
+
+flags.DEFINE_integer(
+    "dupe_factor", 10,
+    "Number of times to duplicate the input data (with different masks).")
+
+flags.DEFINE_float("masked_lm_prob", 0.15, "Masked LM probability.")
+
+flags.DEFINE_float(
+    "short_seq_prob", 0.1,
+    "Probability of creating sequences which are shorter than the "
+    "maximum length.")
+
+
+class TrainingInstance(object):
+  """A single training instance (sentence pair)."""
+
+  def __init__(self, tokens, segment_ids, masked_lm_positions, masked_lm_labels,
+               is_random_next):
+    self.tokens = tokens
+    self.segment_ids = segment_ids
+    self.is_random_next = is_random_next
+    self.masked_lm_positions = masked_lm_positions
+    self.masked_lm_labels = masked_lm_labels
+
+  def __str__(self):
+    s = ""
+    s += "tokens: %s\n" % (" ".join(
+        [tokenization.printable_text(x) for x in self.tokens]))
+    s += "segment_ids: %s\n" % (" ".join([str(x) for x in self.segment_ids]))
+    s += "is_random_next: %s\n" % self.is_random_next
+    s += "masked_lm_positions: %s\n" % (" ".join(
+        [str(x) for x in self.masked_lm_positions]))
+    s += "masked_lm_labels: %s\n" % (" ".join(
+        [tokenization.printable_text(x) for x in self.masked_lm_labels]))
+    s += "\n"
+    return s
+
+  def __repr__(self):
+    return self.__str__()
+
+
+def write_instance_to_example_files(instances, tokenizer, max_seq_length,
+                                    max_predictions_per_seq, output_files):
+  """Create TF example files from `TrainingInstance`s."""
+  writers = []
+  for output_file in output_files:
+    writers.append(tf.python_io.TFRecordWriter(output_file))
+
+  writer_index = 0
+
+  total_written = 0
+  for (inst_index, instance) in enumerate(instances):
+    input_ids = tokenizer.convert_tokens_to_ids(instance.tokens)
+    input_mask = [1] * len(input_ids)
+    segment_ids = list(instance.segment_ids)
+    assert len(input_ids) <= max_seq_length
+
+    while len(input_ids) < max_seq_length:
+      input_ids.append(0)
+      input_mask.append(0)
+      segment_ids.append(0)
+
+    assert len(input_ids) == max_seq_length
+    assert len(input_mask) == max_seq_length
+    assert len(segment_ids) == max_seq_length
+
+    masked_lm_positions = list(instance.masked_lm_positions)
+    masked_lm_ids = tokenizer.convert_tokens_to_ids(instance.masked_lm_labels)
+    masked_lm_weights = [1.0] * len(masked_lm_ids)
+
+    while len(masked_lm_positions) < max_predictions_per_seq:
+      masked_lm_positions.append(0)
+      masked_lm_ids.append(0)
+      masked_lm_weights.append(0.0)
+
+    next_sentence_label = 1 if instance.is_random_next else 0
+
+    features = collections.OrderedDict()
+    features["input_ids"] = create_int_feature(input_ids)
+    features["input_mask"] = create_int_feature(input_mask)
+    features["segment_ids"] = create_int_feature(segment_ids)
+    features["masked_lm_positions"] = create_int_feature(masked_lm_positions)
+    features["masked_lm_ids"] = create_int_feature(masked_lm_ids)
+    features["masked_lm_weights"] = create_float_feature(masked_lm_weights)
+    features["next_sentence_labels"] = create_int_feature([next_sentence_label])
+
+    tf_example = tf.train.Example(features=tf.train.Features(feature=features))
+
+    writers[writer_index].write(tf_example.SerializeToString())
+    writer_index = (writer_index + 1) % len(writers)
+
+    total_written += 1
+
+    if inst_index < 20:
+      tf.logging.info("*** Example ***")
+      tf.logging.info("tokens: %s" % " ".join(
+          [tokenization.printable_text(x) for x in instance.tokens]))
+
+      for feature_name in features.keys():
+        feature = features[feature_name]
+        values = []
+        if feature.int64_list.value:
+          values = feature.int64_list.value
+        elif feature.float_list.value:
+          values = feature.float_list.value
+        tf.logging.info(
+            "%s: %s" % (feature_name, " ".join([str(x) for x in values])))
+
+  for writer in writers:
+    writer.close()
+
+  tf.logging.info("Wrote %d total instances", total_written)
+
+
+def create_int_feature(values):
+  feature = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
+  return feature
+
+
+def create_float_feature(values):
+  feature = tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))
+  return feature
+
+
+def create_training_instances(input_files, tokenizer, max_seq_length,
+                              dupe_factor, short_seq_prob, masked_lm_prob,
+                              max_predictions_per_seq, rng):
+  """Create `TrainingInstance`s from raw text."""
+  all_documents = [[]]
+
+  # Input file format:
+  # (1) One sentence per line. These should ideally be actual sentences, not
+  # entire paragraphs or arbitrary spans of text. (Because we use the
+  # sentence boundaries for the "next sentence prediction" task).
+  # (2) Blank lines between documents. Document boundaries are needed so
+  # that the "next sentence prediction" task doesn't span between documents.
+  for input_file in input_files:
+    with tf.gfile.GFile(input_file, "r") as reader:
+      while True:
+        line = tokenization.convert_to_unicode(reader.readline())
+        if not line:
+          break
+        line = line.strip()
+
+        # Empty lines are used as document delimiters
+        if not line:
+          all_documents.append([])
+        tokens = tokenizer.tokenize(line)
+        if tokens:
+          all_documents[-1].append(tokens)
+
+  # Remove empty documents
+  all_documents = [x for x in all_documents if x]
+  rng.shuffle(all_documents)
+
+  vocab_words = list(tokenizer.vocab.keys())
+  instances = []
+  for _ in range(dupe_factor):
+    for document_index in range(len(all_documents)):
+      instances.extend(
+          create_instances_from_document(
+              all_documents, document_index, max_seq_length, short_seq_prob,
+              masked_lm_prob, max_predictions_per_seq, vocab_words, rng))
+
+  rng.shuffle(instances)
+  return instances
+
+
+def create_instances_from_document(
+    all_documents, document_index, max_seq_length, short_seq_prob,
+    masked_lm_prob, max_predictions_per_seq, vocab_words, rng):
+  """Creates `TrainingInstance`s for a single document."""
+  document = all_documents[document_index]
+
+  # Account for [CLS], [SEP], [SEP]
+  max_num_tokens = max_seq_length - 3
+
+  # We *usually* want to fill up the entire sequence since we are padding
+  # to `max_seq_length` anyways, so short sequences are generally wasted
+  # computation. However, we *sometimes*
+  # (i.e., short_seq_prob == 0.1 == 10% of the time) want to use shorter
+  # sequences to minimize the mismatch between pre-training and fine-tuning.
+  # The `target_seq_length` is just a rough target however, whereas
+  # `max_seq_length` is a hard limit.
+  target_seq_length = max_num_tokens
+  if rng.random() < short_seq_prob:
+    target_seq_length = rng.randint(2, max_num_tokens)
+
+  # We DON'T just concatenate all of the tokens from a document into a long
+  # sequence and choose an arbitrary split point because this would make the
+  # next sentence prediction task too easy. Instead, we split the input into
+  # segments "A" and "B" based on the actual "sentences" provided by the user
+  # input.
+  instances = []
+  current_chunk = []
+  current_length = 0
+  i = 0
+  while i < len(document):
+    segment = document[i]
+    current_chunk.append(segment)
+    current_length += len(segment)
+    if i == len(document) - 1 or current_length >= target_seq_length:
+      if current_chunk:
+        # `a_end` is how many segments from `current_chunk` go into the `A`
+        # (first) sentence.
+        a_end = 1
+        if len(current_chunk) >= 2:
+          a_end = rng.randint(1, len(current_chunk) - 1)
+
+        tokens_a = []
+        for j in range(a_end):
+          tokens_a.extend(current_chunk[j])
+
+        tokens_b = []
+        # Random next
+        is_random_next = False
+        if len(current_chunk) == 1 or rng.random() < 0.5:
+          is_random_next = True
+          target_b_length = target_seq_length - len(tokens_a)
+
+          # This should rarely go for more than one iteration for large
+          # corpora. However, just to be careful, we try to make sure that
+          # the random document is not the same as the document
+          # we're processing.
+          for _ in range(10):
+            random_document_index = rng.randint(0, len(all_documents) - 1)
+            if random_document_index != document_index:
+              break
+
+          random_document = all_documents[random_document_index]
+          random_start = rng.randint(0, len(random_document) - 1)
+          for j in range(random_start, len(random_document)):
+            tokens_b.extend(random_document[j])
+            if len(tokens_b) >= target_b_length:
+              break
+          # We didn't actually use these segments so we "put them back" so
+          # they don't go to waste.
+          num_unused_segments = len(current_chunk) - a_end
+          i -= num_unused_segments
+        # Actual next
+        else:
+          is_random_next = False
+          for j in range(a_end, len(current_chunk)):
+            tokens_b.extend(current_chunk[j])
+        truncate_seq_pair(tokens_a, tokens_b, max_num_tokens, rng)
+
+        assert len(tokens_a) >= 1
+        assert len(tokens_b) >= 1
+
+        tokens = []
+        segment_ids = []
+        tokens.append("[CLS]")
+        segment_ids.append(0)
+        for token in tokens_a:
+          tokens.append(token)
+          segment_ids.append(0)
+
+        tokens.append("[SEP]")
+        segment_ids.append(0)
+
+        for token in tokens_b:
+          tokens.append(token)
+          segment_ids.append(1)
+        tokens.append("[SEP]")
+        segment_ids.append(1)
+
+        (tokens, masked_lm_positions,
+         masked_lm_labels) = create_masked_lm_predictions(
+             tokens, masked_lm_prob, max_predictions_per_seq, vocab_words, rng)
+        instance = TrainingInstance(
+            tokens=tokens,
+            segment_ids=segment_ids,
+            is_random_next=is_random_next,
+            masked_lm_positions=masked_lm_positions,
+            masked_lm_labels=masked_lm_labels)
+        instances.append(instance)
+      current_chunk = []
+      current_length = 0
+    i += 1
+
+  return instances
+
+
+MaskedLmInstance = collections.namedtuple("MaskedLmInstance",
+                                          ["index", "label"])
+
+
+def create_masked_lm_predictions(tokens, masked_lm_prob,
+                                 max_predictions_per_seq, vocab_words, rng):
+  """Creates the predictions for the masked LM objective."""
+
+  cand_indexes = []
+  for (i, token) in enumerate(tokens):
+    if token == "[CLS]" or token == "[SEP]":
+      continue
+    cand_indexes.append(i)
+
+  rng.shuffle(cand_indexes)
+
+  output_tokens = list(tokens)
+
+  num_to_predict = min(max_predictions_per_seq,
+                       max(1, int(round(len(tokens) * masked_lm_prob))))
+
+  masked_lms = []
+  covered_indexes = set()
+  for index in cand_indexes:
+    if len(masked_lms) >= num_to_predict:
+      break
+    if index in covered_indexes:
+      continue
+    covered_indexes.add(index)
+
+    masked_token = None
+    # 80% of the time, replace with [MASK]
+    if rng.random() < 0.8:
+      masked_token = "[MASK]"
+    else:
+      # 10% of the time, keep original
+      if rng.random() < 0.5:
+        masked_token = tokens[index]
+      # 10% of the time, replace with random word
+      else:
+        masked_token = vocab_words[rng.randint(0, len(vocab_words) - 1)]
+
+    output_tokens[index] = masked_token
+
+    masked_lms.append(MaskedLmInstance(index=index, label=tokens[index]))
+
+  masked_lms = sorted(masked_lms, key=lambda x: x.index)
+
+  masked_lm_positions = []
+  masked_lm_labels = []
+  for p in masked_lms:
+    masked_lm_positions.append(p.index)
+    masked_lm_labels.append(p.label)
+
+  return (output_tokens, masked_lm_positions, masked_lm_labels)
+
+
+def truncate_seq_pair(tokens_a, tokens_b, max_num_tokens, rng):
+  """Truncates a pair of sequences to a maximum sequence length."""
+  while True:
+    total_length = len(tokens_a) + len(tokens_b)
+    if total_length <= max_num_tokens:
+      break
+
+    trunc_tokens = tokens_a if len(tokens_a) > len(tokens_b) else tokens_b
+    assert len(trunc_tokens) >= 1
+
+    # We want to sometimes truncate from the front and sometimes from the
+    # back to add more randomness and avoid biases.
+    if rng.random() < 0.5:
+      del trunc_tokens[0]
+    else:
+      trunc_tokens.pop()
+
+
+def main(_):
+  tf.logging.set_verbosity(tf.logging.INFO)
+
+  tokenizer = tokenization.FullTokenizer(
+      vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)
+
+  input_files = []
+  for input_pattern in FLAGS.input_file.split(","):
+    input_files.extend(tf.gfile.Glob(input_pattern))
+
+  tf.logging.info("*** Reading from input files ***")
+  for input_file in input_files:
+    tf.logging.info("  %s", input_file)
+
+  rng = random.Random(FLAGS.random_seed)
+  instances = create_training_instances(
+      input_files, tokenizer, FLAGS.max_seq_length, FLAGS.dupe_factor,
+      FLAGS.short_seq_prob, FLAGS.masked_lm_prob, FLAGS.max_predictions_per_seq,
+      rng)
+
+  output_files = FLAGS.output_file.split(",")
+  tf.logging.info("*** Writing to output files ***")
+  for output_file in output_files:
+    tf.logging.info("  %s", output_file)
+
+  write_instance_to_example_files(instances, tokenizer, FLAGS.max_seq_length,
+                                  FLAGS.max_predictions_per_seq, output_files)
+
+
+if __name__ == "__main__":
+  flags.mark_flag_as_required("input_file")
+  flags.mark_flag_as_required("output_file")
+  flags.mark_flag_as_required("vocab_file")
+  tf.app.run()

tinygrad_repo/test/external/mlperf_bert/preprocessing/external_test_preprocessing_part.py

@@ -0,0 +1,79 @@
+# USAGE:
+# 1.  Download raw text data with `wikipedia_download.py`
+
+# 2.  Install python==3.7.12 and tensorflow==1.15.5
+#     Run `create_pretraining_data.py` to create TFRecords on specific part (This will take some time)
+#     Command: python3 create_pretraining_data.py --input_file=/path/to/part-00XXX-of-00500 --vocab_file=/path/to/vocab.txt \
+#                              --output_file=/path/to/output.tfrecord --max_seq_length=512 --max_predictions_per_seq=76
+#
+# 2.1 For eval: --input_file=/path/to/eval.txt and
+#     Command: python3 pick_eval_samples.py --input_tfrecord=/path/to/eval.tfrecord --output_tfrecord=/path/to/output_eval.tfrecord
+
+# 3.  Run `wikipedia.py` to preprocess the data with tinygrad (Use python > 3.7)
+#     Command: BASEDIR=/path/to/basedir python3 wikipedia.py pre-train X (NOTE: part number needs to match part of step 2)
+#     This will output to /path/to/basedir/train/X.pkl
+#
+# 3.1 For eval:
+#     Command: BASEDIR=/path/to/basedir python3 wikipedia.py pre-eval
+#     This will output to /path/to/basedir/eval.pkl
+
+# 4.  Run this script to verify the correctness of the preprocessing script for specific part
+#     Command: python3 external_test_preprocessing_part.py --preprocessed_part=/path/to/basedir/train/X.pkl --tf_records=/path/to/output.tfrecord
+import os, argparse, pickle
+from tqdm import tqdm
+
+# This is a workaround for protobuf issue
+os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+
+import tensorflow as tf
+
+tf.compat.v1.enable_eager_execution()
+tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+
+def _parse_function(proto, max_seq_length, max_predictions_per_seq):
+  feature_description = {
+      'input_ids': tf.io.FixedLenFeature([max_seq_length], tf.int64),
+      'input_mask': tf.io.FixedLenFeature([max_seq_length], tf.int64),
+      'segment_ids': tf.io.FixedLenFeature([max_seq_length], tf.int64),
+      'masked_lm_positions': tf.io.FixedLenFeature([max_predictions_per_seq], tf.int64),
+      'masked_lm_ids': tf.io.FixedLenFeature([max_predictions_per_seq], tf.int64),
+      'masked_lm_weights': tf.io.FixedLenFeature([max_predictions_per_seq], tf.float32),
+      'next_sentence_labels': tf.io.FixedLenFeature([1], tf.int64),
+  }
+  return tf.io.parse_single_example(proto, feature_description)
+
+def load_dataset(file_path, max_seq_length=512, max_predictions_per_seq=76):
+  dataset = tf.data.TFRecordDataset(file_path)
+  parse_function = lambda proto: _parse_function(proto, max_seq_length, max_predictions_per_seq) # noqa: E731
+  return dataset.map(parse_function)
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(description="Verify the correctness of the preprocessing script for specific part",
+                                   formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  parser.add_argument("--preprocessed_part", type=str, default=None,
+                      help="Path to preprocessed samples file from `wikipedia.py`")
+  parser.add_argument("--tf_records", type=str, default=None,
+                      help="Path to TFRecords file from `create_pretraining_data.py` (Reference implementation)")
+  parser.add_argument("--max_seq_length", type=int, default=512, help="Max sequence length. For MLPerf keep it at 512")
+  parser.add_argument("--max_predictions_per_seq", type=int, default=76, help="Max predictions per sequence. For MLPerf keep it at 76")
+  parser.add_argument("--is_eval", type=bool, default=False, help="Whether to run eval or train preprocessing")
+  args = parser.parse_args()
+
+  assert os.path.isfile(args.preprocessed_part), f"The specified file {args.preprocessed_part} does not exist."
+  assert os.path.isfile(args.tf_records), f"The specified TFRecords file {args.tf_records} does not exist."
+
+  with open(args.preprocessed_part, 'rb') as f:
+    preprocessed_samples = pickle.load(f)
+
+  dataset = load_dataset(args.tf_records, args.max_seq_length, args.max_predictions_per_seq)
+  tf_record_count = sum(1 for _ in dataset)
+  assert tf_record_count == len(preprocessed_samples), f"Samples in reference: {tf_record_count} != Preprocessed samples: {len(preprocessed_samples)}"
+  print(f"Total samples in the part: {tf_record_count}")
+
+  for i, (reference_example, preprocessed_sample) in tqdm(enumerate(zip(dataset, preprocessed_samples)), desc="Checking samples", total=len(preprocessed_samples)): # noqa: E501
+    feature_keys = ["input_ids", "input_mask", "segment_ids", "masked_lm_positions", "masked_lm_ids", "masked_lm_weights", "next_sentence_labels"]
+    for key in feature_keys:
+      reference_example_feature = reference_example[key].numpy()
+      assert (reference_example_feature == preprocessed_sample[key]).all(), \
+      f"{key} are not equal at index {i}\nReference: {reference_example_feature}\nPreprocessed: {preprocessed_sample[key]}"

tinygrad_repo/test/external/mlperf_bert/preprocessing/pick_eval_samples.py

@@ -0,0 +1,127 @@
+# https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/cleanup_scripts/pick_eval_samples.py
+# NOTE: This is a direct copy of the original script
+"""Script for picking certain number of sampels.
+"""
+
+import argparse
+import time
+import logging
+import collections
+import tensorflow as tf
+
+parser = argparse.ArgumentParser(
+    description="Eval sample picker for BERT.")
+parser.add_argument(
+    '--input_tfrecord',
+    type=str,
+    default='',
+    help='Input tfrecord path')
+parser.add_argument(
+    '--output_tfrecord',
+    type=str,
+    default='',
+    help='Output tfrecord path')
+parser.add_argument(
+    '--num_examples_to_pick',
+    type=int,
+    default=10000,
+    help='Number of examples to pick')
+parser.add_argument(
+    '--max_seq_length',
+    type=int,
+    default=512,
+    help='The maximum number of tokens within a sequence.')
+parser.add_argument(
+    '--max_predictions_per_seq',
+    type=int,
+    default=76,
+    help='The maximum number of predictions within a sequence.')
+args = parser.parse_args()
+
+max_seq_length = args.max_seq_length
+max_predictions_per_seq = args.max_predictions_per_seq
+logging.basicConfig(level=logging.INFO)
+
+def decode_record(record):
+  """Decodes a record to a TensorFlow example."""
+  name_to_features = {
+      "input_ids":
+          tf.FixedLenFeature([max_seq_length], tf.int64),
+      "input_mask":
+          tf.FixedLenFeature([max_seq_length], tf.int64),
+      "segment_ids":
+          tf.FixedLenFeature([max_seq_length], tf.int64),
+      "masked_lm_positions":
+          tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
+      "masked_lm_ids":
+          tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
+      "masked_lm_weights":
+          tf.FixedLenFeature([max_predictions_per_seq], tf.float32),
+      "next_sentence_labels":
+          tf.FixedLenFeature([1], tf.int64),
+  }
+
+  example = tf.parse_single_example(record, name_to_features)
+
+  # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
+  # So cast all int64 to int32.
+  for name in list(example.keys()):
+    t = example[name]
+    if t.dtype == tf.int64:
+      t = tf.to_int32(t)
+    example[name] = t
+
+  return example
+
+
+def create_int_feature(values):
+  feature = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
+  return feature
+
+
+def create_float_feature(values):
+  feature = tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))
+  return feature
+
+
+if __name__ == '__main__':
+  tic = time.time()
+  tf.enable_eager_execution()
+
+  d = tf.data.TFRecordDataset(args.input_tfrecord)
+  num_examples = 0
+  records = []
+  for record in d:
+    records.append(record)
+    num_examples += 1
+
+  writer = tf.python_io.TFRecordWriter(args.output_tfrecord)
+  i = 0
+  pick_ratio = num_examples / args.num_examples_to_pick
+  num_examples_picked = 0
+  for i in range(args.num_examples_to_pick):
+    example = decode_record(records[int(i * pick_ratio)])
+    features = collections.OrderedDict()
+    features["input_ids"] = create_int_feature(
+        example["input_ids"].numpy())
+    features["input_mask"] = create_int_feature(
+        example["input_mask"].numpy())
+    features["segment_ids"] = create_int_feature(
+        example["segment_ids"].numpy())
+    features["masked_lm_positions"] = create_int_feature(
+        example["masked_lm_positions"].numpy())
+    features["masked_lm_ids"] = create_int_feature(
+        example["masked_lm_ids"].numpy())
+    features["masked_lm_weights"] = create_float_feature(
+        example["masked_lm_weights"].numpy())
+    features["next_sentence_labels"] = create_int_feature(
+        example["next_sentence_labels"].numpy())
+
+    tf_example = tf.train.Example(features=tf.train.Features(feature=features))
+    writer.write(tf_example.SerializeToString())
+    num_examples_picked += 1
+
+  writer.close()
+  toc = time.time()
+  logging.info("Picked %d examples out of %d samples in %.2f sec",
+               num_examples_picked, num_examples, toc - tic)

tinygrad_repo/test/external/mlperf_bert/preprocessing/tokenization.py

@@ -0,0 +1,415 @@
+# https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/cleanup_scripts/tokenization.py
+# NOTE: This is a direct copy of the original script
+"""Tokenization classes."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+import re
+import unicodedata
+
+from absl import flags
+import six
+import tensorflow.compat.v1 as tf
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_bool(
+    "preserve_unused_tokens", False,
+    "If True, Wordpiece tokenization will not be applied to words in the vocab."
+)
+
+_UNUSED_TOKEN_RE = re.compile("^\\[unused\\d+\\]$")
+
+
+def preserve_token(token, vocab):
+  """Returns True if the token should forgo tokenization and be preserved."""
+  if not FLAGS.preserve_unused_tokens:
+    return False
+  if token not in vocab:
+    return False
+  return bool(_UNUSED_TOKEN_RE.search(token))
+
+
+def validate_case_matches_checkpoint(do_lower_case, init_checkpoint):
+  """Checks whether the casing config is consistent with the checkpoint name."""
+
+  # The casing has to be passed in by the user and there is no explicit check
+  # as to whether it matches the checkpoint. The casing information probably
+  # should have been stored in the bert_config.json file, but it's not, so
+  # we have to heuristically detect it to validate.
+
+  if not init_checkpoint:
+    return
+
+  m = re.match("^.*?([A-Za-z0-9_-]+)/bert_model.ckpt", init_checkpoint)
+  if m is None:
+    return
+
+  model_name = m.group(1)
+
+  lower_models = [
+      "uncased_L-24_H-1024_A-16", "uncased_L-12_H-768_A-12",
+      "multilingual_L-12_H-768_A-12", "chinese_L-12_H-768_A-12"
+  ]
+
+  cased_models = [
+      "cased_L-12_H-768_A-12", "cased_L-24_H-1024_A-16",
+      "multi_cased_L-12_H-768_A-12"
+  ]
+
+  is_bad_config = False
+  if model_name in lower_models and not do_lower_case:
+    is_bad_config = True
+    actual_flag = "False"
+    case_name = "lowercased"
+    opposite_flag = "True"
+
+  if model_name in cased_models and do_lower_case:
+    is_bad_config = True
+    actual_flag = "True"
+    case_name = "cased"
+    opposite_flag = "False"
+
+  if is_bad_config:
+    raise ValueError(
+        "You passed in `--do_lower_case=%s` with `--init_checkpoint=%s`. "
+        "However, `%s` seems to be a %s model, so you "
+        "should pass in `--do_lower_case=%s` so that the fine-tuning matches "
+        "how the model was pre-training. If this error is wrong, please "
+        "just comment out this check." % (actual_flag, init_checkpoint,
+                                          model_name, case_name, opposite_flag))
+
+
+def convert_to_unicode(text):
+  """Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
+  if six.PY3:
+    if isinstance(text, str):
+      return text
+    elif isinstance(text, bytes):
+      return text.decode("utf-8", "ignore")
+    else:
+      raise ValueError("Unsupported string type: %s" % (type(text)))
+  elif six.PY2:
+    if isinstance(text, str):
+      return text.decode("utf-8", "ignore")
+    elif isinstance(text, unicode): # noqa: F821
+      return text
+    else:
+      raise ValueError("Unsupported string type: %s" % (type(text)))
+  else:
+    raise ValueError("Not running on Python2 or Python 3?")
+
+
+def printable_text(text):
+  """Returns text encoded in a way suitable for print or `tf.logging`."""
+
+  # These functions want `str` for both Python2 and Python3, but in one case
+  # it's a Unicode string and in the other it's a byte string.
+  if six.PY3:
+    if isinstance(text, str):
+      return text
+    elif isinstance(text, bytes):
+      return text.decode("utf-8", "ignore")
+    else:
+      raise ValueError("Unsupported string type: %s" % (type(text)))
+  elif six.PY2:
+    if isinstance(text, str):
+      return text
+    elif isinstance(text, unicode): # noqa: F821
+      return text.encode("utf-8")
+    else:
+      raise ValueError("Unsupported string type: %s" % (type(text)))
+  else:
+    raise ValueError("Not running on Python2 or Python 3?")
+
+
+def load_vocab(vocab_file):
+  """Loads a vocabulary file into a dictionary."""
+  vocab = collections.OrderedDict()
+  with tf.gfile.GFile(vocab_file, "r") as reader:
+    while True:
+      token = convert_to_unicode(reader.readline())
+      if not token:
+        break
+      token = token.strip()
+      if token not in vocab:
+        vocab[token] = len(vocab)
+  return vocab
+
+
+def convert_by_vocab(vocab, items):
+  """Converts a sequence of [tokens|ids] using the vocab."""
+  output = []
+  for item in items:
+    output.append(vocab[item])
+  return output
+
+
+def convert_tokens_to_ids(vocab, tokens):
+  return convert_by_vocab(vocab, tokens)
+
+
+def convert_ids_to_tokens(inv_vocab, ids):
+  return convert_by_vocab(inv_vocab, ids)
+
+
+def whitespace_tokenize(text):
+  """Runs basic whitespace cleaning and splitting on a piece of text."""
+  text = text.strip()
+  if not text:
+    return []
+  tokens = text.split()
+  return tokens
+
+
+class FullTokenizer(object):
+  """Runs end-to-end tokenziation."""
+
+  def __init__(self, vocab_file, do_lower_case=True):
+    self.vocab = load_vocab(vocab_file)
+    self.inv_vocab = {v: k for k, v in self.vocab.items()}
+    self.basic_tokenizer = BasicTokenizer(
+        do_lower_case=do_lower_case, vocab=self.vocab)
+    self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab)
+
+  def tokenize(self, text):
+    split_tokens = []
+    for token in self.basic_tokenizer.tokenize(text):
+      if preserve_token(token, self.vocab):
+        split_tokens.append(token)
+        continue
+      for sub_token in self.wordpiece_tokenizer.tokenize(token):
+        split_tokens.append(sub_token)
+
+    return split_tokens
+
+  def convert_tokens_to_ids(self, tokens):
+    return convert_by_vocab(self.vocab, tokens)
+
+  def convert_ids_to_tokens(self, ids):
+    return convert_by_vocab(self.inv_vocab, ids)
+
+
+class BasicTokenizer(object):
+  """Runs basic tokenization (punctuation splitting, lower casing, etc.)."""
+
+  def __init__(self, do_lower_case=True, vocab=tuple()):
+    """Constructs a BasicTokenizer.
+
+    Args:
+      do_lower_case: Whether to lower case the input.
+      vocab: A container of tokens to not mutate during tokenization.
+    """
+    self.do_lower_case = do_lower_case
+    self.vocab = vocab
+
+  def tokenize(self, text):
+    """Tokenizes a piece of text."""
+    text = convert_to_unicode(text)
+    text = self._clean_text(text)
+
+    # This was added on November 1st, 2018 for the multilingual and Chinese
+    # models. This is also applied to the English models now, but it doesn't
+    # matter since the English models were not trained on any Chinese data
+    # and generally don't have any Chinese data in them (there are Chinese
+    # characters in the vocabulary because Wikipedia does have some Chinese
+    # words in the English Wikipedia.).
+    text = self._tokenize_chinese_chars(text)
+
+    orig_tokens = whitespace_tokenize(text)
+    split_tokens = []
+    for token in orig_tokens:
+      if preserve_token(token, self.vocab):
+        split_tokens.append(token)
+        continue
+      if self.do_lower_case:
+        token = token.lower()
+        token = self._run_strip_accents(token)
+      split_tokens.extend(self._run_split_on_punc(token))
+
+    output_tokens = whitespace_tokenize(" ".join(split_tokens))
+    return output_tokens
+
+  def _run_strip_accents(self, text):
+    """Strips accents from a piece of text."""
+    text = unicodedata.normalize("NFD", text)
+    output = []
+    for char in text:
+      cat = unicodedata.category(char)
+      if cat == "Mn":
+        continue
+      output.append(char)
+    return "".join(output)
+
+  def _run_split_on_punc(self, text):
+    """Splits punctuation on a piece of text."""
+    chars = list(text)
+    i = 0
+    start_new_word = True
+    output = []
+    while i < len(chars):
+      char = chars[i]
+      if _is_punctuation(char):
+        output.append([char])
+        start_new_word = True
+      else:
+        if start_new_word:
+          output.append([])
+        start_new_word = False
+        output[-1].append(char)
+      i += 1
+
+    return ["".join(x) for x in output]
+
+  def _tokenize_chinese_chars(self, text):
+    """Adds whitespace around any CJK character."""
+    output = []
+    for char in text:
+      cp = ord(char)
+      if self._is_chinese_char(cp):
+        output.append(" ")
+        output.append(char)
+        output.append(" ")
+      else:
+        output.append(char)
+    return "".join(output)
+
+  def _is_chinese_char(self, cp):
+    """Checks whether CP is the codepoint of a CJK character."""
+    # This defines a "chinese character" as anything in the CJK Unicode block:
+    #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+    #
+    # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+    # despite its name. The modern Korean Hangul alphabet is a different block,
+    # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+    # space-separated words, so they are not treated specially and handled
+    # like all of the other languages.
+    if ((cp >= 0x4E00 and cp <= 0x9FFF) or  #
+        (cp >= 0x3400 and cp <= 0x4DBF) or  #
+        (cp >= 0x20000 and cp <= 0x2A6DF) or  #
+        (cp >= 0x2A700 and cp <= 0x2B73F) or  #
+        (cp >= 0x2B740 and cp <= 0x2B81F) or  #
+        (cp >= 0x2B820 and cp <= 0x2CEAF) or
+        (cp >= 0xF900 and cp <= 0xFAFF) or  #
+        (cp >= 0x2F800 and cp <= 0x2FA1F)):  #
+      return True
+
+    return False
+
+  def _clean_text(self, text):
+    """Performs invalid character removal and whitespace cleanup on text."""
+    output = []
+    for char in text:
+      cp = ord(char)
+      if cp == 0 or cp == 0xfffd or _is_control(char):
+        continue
+      if _is_whitespace(char):
+        output.append(" ")
+      else:
+        output.append(char)
+    return "".join(output)
+
+
+class WordpieceTokenizer(object):
+  """Runs WordPiece tokenziation."""
+
+  def __init__(self, vocab, unk_token="[UNK]", max_input_chars_per_word=200):
+    self.vocab = vocab
+    self.unk_token = unk_token
+    self.max_input_chars_per_word = max_input_chars_per_word
+
+  def tokenize(self, text):
+    """Tokenizes a piece of text into its word pieces.
+
+    This uses a greedy longest-match-first algorithm to perform tokenization
+    using the given vocabulary.
+
+    For example:
+      input = "unaffable"
+      output = ["un", "##aff", "##able"]
+
+    Args:
+      text: A single token or whitespace separated tokens. This should have
+        already been passed through `BasicTokenizer.
+
+    Returns:
+      A list of wordpiece tokens.
+    """
+
+    text = convert_to_unicode(text)
+
+    output_tokens = []
+    for token in whitespace_tokenize(text):
+      chars = list(token)
+      if len(chars) > self.max_input_chars_per_word:
+        output_tokens.append(self.unk_token)
+        continue
+
+      is_bad = False
+      start = 0
+      sub_tokens = []
+      while start < len(chars):
+        end = len(chars)
+        cur_substr = None
+        while start < end:
+          substr = "".join(chars[start:end])
+          if start > 0:
+            substr = "##" + substr
+          if substr in self.vocab:
+            cur_substr = substr
+            break
+          end -= 1
+        if cur_substr is None:
+          is_bad = True
+          break
+        sub_tokens.append(cur_substr)
+        start = end
+
+      if is_bad:
+        output_tokens.append(self.unk_token)
+      else:
+        output_tokens.extend(sub_tokens)
+    return output_tokens
+
+
+def _is_whitespace(char):
+  """Checks whether `chars` is a whitespace character."""
+  # \t, \n, and \r are technically control characters but we treat them
+  # as whitespace since they are generally considered as such.
+  if char == " " or char == "\t" or char == "\n" or char == "\r":
+    return True
+  cat = unicodedata.category(char)
+  if cat == "Zs":
+    return True
+  return False
+
+
+def _is_control(char):
+  """Checks whether `chars` is a control character."""
+  # These are technically control characters but we count them as whitespace
+  # characters.
+  if char == "\t" or char == "\n" or char == "\r":
+    return False
+  cat = unicodedata.category(char)
+  if cat in ("Cc", "Cf"):
+    return True
+  return False
+
+
+def _is_punctuation(char):
+  """Checks whether `chars` is a punctuation character."""
+  cp = ord(char)
+  # We treat all non-letter/number ASCII as punctuation.
+  # Characters such as "^", "$", and "`" are not in the Unicode
+  # Punctuation class but we treat them as punctuation anyways, for
+  # consistency.
+  if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or
+      (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)):
+    return True
+  cat = unicodedata.category(char)
+  if cat.startswith("P"):
+    return True
+  return False

tinygrad_repo/test/external/mlperf_resnet/lars_optimizer.py

@@ -0,0 +1,219 @@
+# https://github.com/mlcommons/training/blob/e3769c8dcf88cd21e1001dd2f894b40a1513ec5d/image_classification/tensorflow2/lars_optimizer.py
+# changes: don't call lr_t if it's not a schedule
+
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Layer-wise Adaptive Rate Scaling optimizer for large-batch training."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+# from tf2_common.training import optimizer_v2modified
+from tensorflow.python.framework import ops
+from tensorflow.python.keras import backend_config
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import linalg_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import state_ops
+
+
+# class LARSOptimizer(optimizer_v2modified.OptimizerV2Modified):
+class LARSOptimizer(optimizer_v2.OptimizerV2):
+  """Layer-wise Adaptive Rate Scaling for large batch training.
+
+  Introduced by "Large Batch Training of Convolutional Networks" by Y. You,
+  I. Gitman, and B. Ginsburg. (https://arxiv.org/abs/1708.03888)
+
+  Implements the LARS learning rate scheme presented in the paper above. This
+  optimizer is useful when scaling the batch size to up to 32K without
+  significant performance degradation. It is recommended to use the optimizer
+  in conjunction with:
+      - Gradual learning rate warm-up
+      - Linear learning rate scaling
+      - Poly rule learning rate decay
+
+  Note, LARS scaling is currently only enabled for dense tensors. Sparse tensors
+  use the default momentum optimizer.
+  """
+
+  def __init__(
+      self,
+      learning_rate,
+      momentum=0.9,
+      weight_decay=0.0001,
+      # The LARS coefficient is a hyperparameter
+      eeta=0.001,
+      epsilon=0.0,
+      name="LARSOptimizer",
+      # Enable skipping variables from LARS scaling.
+      # TODO(sameerkm): Enable a direct mechanism to pass a
+      # subset of variables to the optimizer.
+      skip_list=None,
+      use_nesterov=False,
+      **kwargs):
+    """Construct a new LARS Optimizer.
+
+    Args:
+      learning_rate: A `Tensor`, floating point value, or a schedule that is a
+        `tf.keras.optimizers.schedules.LearningRateSchedule`, or a callable
+        that takes no arguments and returns the actual value to use. The
+        learning rate.
+      momentum: A floating point value. Momentum hyperparameter.
+      weight_decay: A floating point value. Weight decay hyperparameter.
+      eeta: LARS coefficient as used in the paper. Dfault set to LARS
+        coefficient from the paper. (eeta / weight_decay) determines the highest
+        scaling factor in LARS.
+      epsilon: Optional epsilon parameter to be set in models that have very
+        small gradients. Default set to 0.0.
+      name: Optional name prefix for variables and ops created by LARSOptimizer.
+      skip_list: List of strings to enable skipping variables from LARS scaling.
+        If any of the strings in skip_list is a subset of var.name, variable
+        'var' is skipped from LARS scaling. For a typical classification model
+        with batch normalization, the skip_list is ['batch_normalization',
+        'bias']
+      use_nesterov: when set to True, nesterov momentum will be enabled
+      **kwargs: keyword arguments.
+
+    Raises:
+      ValueError: If a hyperparameter is set to a non-sensical value.
+    """
+    if momentum < 0.0:
+      raise ValueError("momentum should be positive: %s" % momentum)
+    if weight_decay < 0.0:
+      raise ValueError("weight_decay should be positive: %s" % weight_decay)
+    super(LARSOptimizer, self).__init__(name=name, **kwargs)
+
+    self._set_hyper("learning_rate", learning_rate)
+
+    # When directly using class members, instead of
+    # _set_hyper and _get_hyper (such as learning_rate above),
+    # the values are fixed after __init(), and not being
+    # updated during the training process.
+    # This provides better performance but less flexibility.
+    self.momentum = momentum
+    self.weight_decay = weight_decay
+    self.eeta = eeta
+    self.epsilon = epsilon or backend_config.epsilon()
+    self._skip_list = skip_list
+    self.use_nesterov = use_nesterov
+
+  def _prepare_local(self, var_device, var_dtype, apply_state):
+    lr_t = self._get_hyper("learning_rate", var_dtype)
+    local_step = math_ops.cast(self.iterations, var_dtype)
+    if callable(lr_t): lr_t = math_ops.cast(lr_t(local_step), var_dtype)
+    learning_rate_t = array_ops.identity(lr_t)
+
+    apply_state[(var_device, var_dtype)].update(
+        dict(
+            learning_rate=learning_rate_t,
+            ))
+
+  def _create_slots(self, var_list):
+    for v in var_list:
+      self.add_slot(v, "momentum")
+
+  def compute_lr(self, grad, var, coefficients):
+    scaled_lr = coefficients["learning_rate"]
+    if self._skip_list is None or not any(v in var.name
+                                          for v in self._skip_list):
+      w_norm = linalg_ops.norm(var, ord=2)
+      g_norm = linalg_ops.norm(grad, ord=2)
+      trust_ratio = array_ops.where(
+          math_ops.greater(w_norm, 0),
+          array_ops.where(
+              math_ops.greater(g_norm, 0),
+              (self.eeta * w_norm /
+               (g_norm + self.weight_decay * w_norm + self.epsilon)), 1.0), 1.0)
+
+      scaled_lr = coefficients["learning_rate"] * trust_ratio
+      # Add the weight regularization gradient
+      grad = grad + self.weight_decay * var
+    return scaled_lr, grad
+
+  def _apply_dense(self, grad, var, apply_state=None):
+    return self._resource_apply_dense(grad, var, apply_state)
+
+  def _resource_apply_dense(self, grad, var, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    scaled_lr, grad = self.compute_lr(grad, var, coefficients)
+    mom = self.get_slot(var, "momentum")
+    # Use ApplyKerasMomentum instead of ApplyMomentum
+    # training_ops.resource_apply_keras_momentum(
+    #     var.handle,
+    #     mom.handle,
+    #     scaled_lr,
+    #     grad,
+    #     coefficients["momentum"],
+    #     use_locking=False,
+    #     use_nesterov=self.use_nesterov)
+
+    mom_t = mom * self.momentum - grad * scaled_lr
+    mom_t = state_ops.assign(mom, mom_t, use_locking=False)
+    if self.use_nesterov:
+      var_t = var + mom_t * self.momentum - grad * scaled_lr
+    else:
+      var_t = var + mom_t
+    return state_ops.assign(var, var_t, use_locking=False).op
+
+  # Fallback to momentum optimizer for sparse tensors
+  def _apply_sparse(self, grad, var, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    mom = self.get_slot(var, "momentum")
+    return tf.raw_ops.SparseApplyMomentum(
+        var=var,
+        accum=mom,
+        lr=coefficients["learning_rate"],
+        grad=grad.values,
+        indices=grad.indices,
+        momentum=self.momentum,
+        use_locking=False,
+        use_nesterov=self.use_nesterov)
+
+  def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    mom = self.get_slot(var, "momentum")
+    return tf.raw_ops.ResourceSparseApplyKerasMomentum(
+        var=var.handle,
+        accum=mom.handle,
+        lr=coefficients["learning_rate"],
+        grad=grad,
+        indices=indices,
+        momentum=self.momentum,
+        use_locking=False,
+        use_nesterov=self.use_nesterov)
+
+  def get_config(self):
+    config = super(LARSOptimizer, self).get_config()
+    config.update({
+        "learning_rate": self._serialize_hyperparameter("learning_rate"),
+        "momentum": self.momentum,
+        "weight_decay": self.weight_decay,
+        "eeta": self.eeta,
+        "epsilon": self.epsilon,
+        "use_nesterov": self.use_nesterov,
+    })
+    return config

tinygrad_repo/test/external/mlperf_resnet/lars_util.py

@@ -0,0 +1,179 @@
+# https://github.com/mlcommons/training/blob/e237206991d10449d9675d95606459a3cb6c21ad/image_classification/tensorflow2/lars_util.py
+# changes: commented out logging
+# changes: convert_to_tensor_v2 -> convert_to_tensor
+# changes: extend from tf.python.keras.optimizer_v2.learning_rate_schedule.LearningRateScheduler
+
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Enable Layer-wise Adaptive Rate Scaling optimizer in ResNet."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl import flags
+import tensorflow as tf
+
+#from tf2_common.utils.mlp_log import mlp_log
+from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.keras.optimizer_v2 import learning_rate_schedule
+
+FLAGS = flags.FLAGS
+
+
+def define_lars_flags():
+  """Defines flags needed by LARS optimizer."""
+
+  flags.DEFINE_float(
+      'end_learning_rate', default=None,
+      help=('Polynomial decay end learning rate.'))
+
+  flags.DEFINE_float(
+      'lars_epsilon', default=0.0,
+      help=('Override autoselected LARS epsilon.'))
+
+  flags.DEFINE_float(
+      'warmup_epochs', default=None,
+      help=('Override autoselected polynomial decay warmup epochs.'))
+
+  flags.DEFINE_float(
+      'momentum',
+      default=0.9,
+      help=('Momentum parameter used in the MomentumOptimizer.'))
+
+
+class PolynomialDecayWithWarmup(learning_rate_schedule.LearningRateSchedule):
+  """A LearningRateSchedule that uses a polynomial decay with warmup."""
+
+  def __init__(
+      self,
+      batch_size,
+      steps_per_epoch,
+      train_steps,
+      initial_learning_rate=None,
+      end_learning_rate=None,
+      warmup_epochs=None,
+      compute_lr_on_cpu=False,
+      name=None):
+    """Applies a polynomial decay to the learning rate with warmup."""
+    super(PolynomialDecayWithWarmup, self).__init__()
+
+    self.batch_size = batch_size
+    self.steps_per_epoch = steps_per_epoch
+    self.train_steps = train_steps
+    self.name = name
+    self.learning_rate_ops_cache = {}
+    self.compute_lr_on_cpu = compute_lr_on_cpu
+
+    if batch_size < 16384:
+      self.initial_learning_rate = 10.0
+      warmup_epochs_ = 5
+    elif batch_size < 32768:
+      self.initial_learning_rate = 25.0
+      warmup_epochs_ = 5
+    else:
+      self.initial_learning_rate = 31.2
+      warmup_epochs_ = 25
+
+    # Override default poly learning rate and warmup epochs
+    if initial_learning_rate:
+      self.initial_learning_rate = initial_learning_rate
+
+    if end_learning_rate:
+      self.end_learning_rate = end_learning_rate
+    else:
+      self.end_learning_rate = 0.0001
+
+    if warmup_epochs is not None:
+      warmup_epochs_ = warmup_epochs
+    self.warmup_epochs = warmup_epochs_
+
+    """
+    opt_name = FLAGS.optimizer.lower()
+    mlp_log.mlperf_print('opt_name', opt_name)
+    if opt_name == 'lars':
+      mlp_log.mlperf_print('{}_epsilon'.format(opt_name), FLAGS.lars_epsilon)
+    mlp_log.mlperf_print('{}_opt_weight_decay'.format(opt_name),
+                         FLAGS.weight_decay)
+    mlp_log.mlperf_print('{}_opt_base_learning_rate'.format(opt_name),
+                         self.initial_learning_rate)
+    mlp_log.mlperf_print('{}_opt_learning_rate_warmup_epochs'.format(opt_name),
+                         warmup_epochs_)
+    mlp_log.mlperf_print('{}_opt_end_learning_rate'.format(opt_name),
+                         self.end_learning_rate)
+    """
+    warmup_steps = warmup_epochs_ * steps_per_epoch
+    self.warmup_steps = tf.cast(warmup_steps, tf.float32)
+    self.decay_steps = train_steps - warmup_steps + 1
+    """
+    mlp_log.mlperf_print('{}_opt_learning_rate_decay_steps'.format(opt_name),
+                         int(self.decay_steps))
+    mlp_log.mlperf_print(
+        '{}_opt_learning_rate_decay_poly_power'.format(opt_name), 2.0)
+    mlp_log.mlperf_print('{}_opt_momentum'.format(opt_name), FLAGS.momentum)
+    """
+
+    self.poly_rate_scheduler = tf.keras.optimizers.schedules.PolynomialDecay(
+        initial_learning_rate=self.initial_learning_rate,
+        decay_steps=self.decay_steps,
+        end_learning_rate=self.end_learning_rate,
+        power=2.0)
+
+  def __call__(self, step):
+    if tf.executing_eagerly():
+      return self._get_learning_rate(step)
+
+    # In an eager function or graph, the current implementation of optimizer
+    # repeatedly call and thus create ops for the learning rate schedule. To
+    # avoid this, we cache the ops if not executing eagerly.
+    graph = tf.compat.v1.get_default_graph()
+    if graph not in self.learning_rate_ops_cache:
+      if self.compute_lr_on_cpu:
+        with tf.device('/device:CPU:0'):
+          self.learning_rate_ops_cache[graph] = self._get_learning_rate(step)
+      else:
+        self.learning_rate_ops_cache[graph] = self._get_learning_rate(step)
+    return self.learning_rate_ops_cache[graph]
+
+  def _get_learning_rate(self, step):
+    with ops.name_scope_v2(self.name or 'PolynomialDecayWithWarmup') as name:
+
+      initial_learning_rate = ops.convert_to_tensor(
+          self.initial_learning_rate, name='initial_learning_rate')
+      warmup_steps = ops.convert_to_tensor(
+          self.warmup_steps, name='warmup_steps')
+
+      warmup_rate = (
+          initial_learning_rate * step / warmup_steps)
+
+      poly_steps = math_ops.subtract(step, warmup_steps)
+      poly_rate = self.poly_rate_scheduler(poly_steps)
+
+      decay_rate = tf.where(step <= warmup_steps,
+                            warmup_rate, poly_rate, name=name)
+      return decay_rate
+
+  def get_config(self):
+    return {
+        'batch_size': self.batch_size,
+        'steps_per_epoch': self.steps_per_epoch,
+        'train_steps': self.train_steps,
+        'initial_learning_rate': self.initial_learning_rate,
+        'end_learning_rate': self.end_learning_rate,
+        'warmup_epochs': self.warmup_epochs,
+        'name': self.name,
+    }

tinygrad_repo/test/external/mlperf_retinanet/coco_utils.py

@@ -0,0 +1,89 @@
+# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/coco_utils.py
+
+import os
+import torch
+import torchvision
+
+from test.external.mlperf_retinanet import transforms as T
+
+class ConvertCocoPolysToMask(object):
+    def __init__(self, filter_iscrowd=True):
+        self.filter_iscrowd = filter_iscrowd
+
+    def __call__(self, image, target):
+        w, h = image.size
+
+        image_id = target["image_id"]
+        image_id = torch.tensor([image_id])
+
+        anno = target["annotations"]
+
+        if self.filter_iscrowd:
+            anno = [obj for obj in anno if obj['iscrowd'] == 0]
+
+        boxes = [obj["bbox"] for obj in anno]
+        # guard against no boxes via resizing
+        boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
+        boxes[:, 2:] += boxes[:, :2]
+        boxes[:, 0::2].clamp_(min=0, max=w)
+        boxes[:, 1::2].clamp_(min=0, max=h)
+
+        classes = [obj["category_id"] for obj in anno]
+        classes = torch.tensor(classes, dtype=torch.int64)
+
+        keypoints = None
+        if anno and "keypoints" in anno[0]:
+            keypoints = [obj["keypoints"] for obj in anno]
+            keypoints = torch.as_tensor(keypoints, dtype=torch.float32)
+            num_keypoints = keypoints.shape[0]
+            if num_keypoints:
+                keypoints = keypoints.view(num_keypoints, -1, 3)
+
+        keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
+        boxes = boxes[keep]
+        classes = classes[keep]
+
+        target = {}
+        target["boxes"] = boxes
+        target["labels"] = classes
+        target["image_id"] = image_id
+
+        # for conversion to coco api
+        area = torch.tensor([obj["area"] for obj in anno])
+        iscrowd = torch.tensor([obj["iscrowd"] for obj in anno])
+        target["area"] = area
+        target["iscrowd"] = iscrowd
+
+        return image, target
+
+class CocoDetection(torchvision.datasets.CocoDetection):
+    def __init__(self, img_folder, ann_file, transforms):
+        super(CocoDetection, self).__init__(img_folder, ann_file)
+        self._transforms = transforms
+
+    def __getitem__(self, idx):
+        img, target = super(CocoDetection, self).__getitem__(idx)
+        image_id = self.ids[idx]
+        target = dict(image_id=image_id, annotations=target)
+        if self._transforms is not None:
+            img, target = self._transforms(img, target)
+        return img, target
+
+def get_openimages(name, root, image_set, transforms):
+    PATHS = {
+        "train": os.path.join(root, "train"),
+        "val":   os.path.join(root, "validation"),
+    }
+
+    t = [ConvertCocoPolysToMask(filter_iscrowd=False)]
+
+    if transforms is not None:
+        t.append(transforms)
+    transforms = T.Compose(t)
+
+    img_folder = os.path.join(PATHS[image_set], "data")
+    ann_file = os.path.join(PATHS[image_set], "labels", f"{name}.json")
+
+    dataset = CocoDetection(img_folder, ann_file, transforms=transforms)
+
+    return dataset

tinygrad_repo/test/external/mlperf_retinanet/focal_loss.py

@@ -0,0 +1,51 @@
+# Copied from https://github.com/mlcommons/training/blob/cdd928d4596c142c15a7d86b2eeadbac718c8da2/single_stage_detector/ssd/model/focal_loss.py
+
+import torch
+import torch.nn.functional as F
+
+
+def sigmoid_focal_loss(
+    inputs: torch.Tensor,
+    targets: torch.Tensor,
+    alpha: float = 0.25,
+    gamma: float = 2,
+    reduction: str = "none",
+):
+    """
+    Original implementation from https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/focal_loss.py .
+    Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.
+
+    Args:
+        inputs: A float tensor of arbitrary shape.
+                The predictions for each example.
+        targets: A float tensor with the same shape as inputs. Stores the binary
+                classification label for each element in inputs
+                (0 for the negative class and 1 for the positive class).
+        alpha: (optional) Weighting factor in range (0,1) to balance
+                positive vs negative examples or -1 for ignore. Default = 0.25
+        gamma: Exponent of the modulating factor (1 - p_t) to
+               balance easy vs hard examples.
+        reduction: 'none' | 'mean' | 'sum'
+                 'none': No reduction will be applied to the output.
+                 'mean': The output will be averaged.
+                 'sum': The output will be summed.
+    Returns:
+        Loss tensor with the reduction option applied.
+    """
+    p = torch.sigmoid(inputs)
+    ce_loss = F.binary_cross_entropy_with_logits(
+        inputs, targets, reduction="none"
+    )
+    p_t = p * targets + (1 - p) * (1 - targets)
+    loss = ce_loss * ((1 - p_t) ** gamma)
+
+    if alpha >= 0:
+        alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
+        loss = alpha_t * loss
+
+    if reduction == "mean":
+        loss = loss.mean()
+    elif reduction == "sum":
+        loss = loss.sum()
+
+    return loss

tinygrad_repo/test/external/mlperf_retinanet/model/boxes.py

@@ -0,0 +1,65 @@
+# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/model/boxes.py
+
+import torch
+from torch import Tensor
+from typing import Tuple
+
+def _upcast(t: Tensor) -> Tensor:
+    # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type
+    if t.is_floating_point():
+        return t if t.dtype in (torch.float32, torch.float64) else t.float()
+    else:
+        return t if t.dtype in (torch.int32, torch.int64) else t.int()
+
+
+def box_area(boxes: Tensor) -> Tensor:
+    """
+    Computes the area of a set of bounding boxes, which are specified by their
+    (x1, y1, x2, y2) coordinates.
+
+    Args:
+        boxes (Tensor[N, 4]): boxes for which the area will be computed. They
+            are expected to be in (x1, y1, x2, y2) format with
+            ``0 <= x1 < x2`` and ``0 <= y1 < y2``.
+
+    Returns:
+        Tensor[N]: the area for each box
+    """
+    boxes = _upcast(boxes)
+    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
+
+
+# implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py
+# with slight modifications
+def _box_inter_union(boxes1: Tensor, boxes2: Tensor) -> Tuple[Tensor, Tensor]:
+    area1 = box_area(boxes1)
+    area2 = box_area(boxes2)
+
+    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
+    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]
+
+    wh = _upcast(rb - lt).clamp(min=0)  # [N,M,2]
+    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]
+
+    union = area1[:, None] + area2 - inter
+
+    return inter, union
+
+
+def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
+    """
+    Return intersection-over-union (Jaccard index) between two sets of boxes.
+
+    Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with
+    ``0 <= x1 < x2`` and ``0 <= y1 < y2``.
+
+    Args:
+        boxes1 (Tensor[N, 4]): first set of boxes
+        boxes2 (Tensor[M, 4]): second set of boxes
+
+    Returns:
+        Tensor[N, M]: the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2
+    """
+    inter, union = _box_inter_union(boxes1, boxes2)
+    iou = inter / union
+    return iou

tinygrad_repo/test/external/mlperf_retinanet/model/image_list.py

@@ -0,0 +1,27 @@
+# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/model/image_list.py
+
+import torch
+from torch import Tensor
+from typing import List, Tuple
+
+
+class ImageList(object):
+    """
+    Structure that holds a list of images (of possibly
+    varying sizes) as a single tensor.
+    This works by padding the images to the same size,
+    and storing in a field the original sizes of each image
+    """
+
+    def __init__(self, tensors: Tensor, image_sizes: List[Tuple[int, int]]):
+        """
+        Args:
+            tensors (tensor)
+            image_sizes (list[tuple[int, int]])
+        """
+        self.tensors = tensors
+        self.image_sizes = image_sizes
+
+    def to(self, device: torch.device) -> 'ImageList':
+        cast_tensor = self.tensors.to(device)
+        return ImageList(cast_tensor, self.image_sizes)

tinygrad_repo/test/external/mlperf_retinanet/model/transform.py

@@ -0,0 +1,163 @@
+# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/model/transform.py
+
+import torch
+
+from torch import nn, Tensor
+from typing import List, Tuple, Dict, Optional
+
+from test.external.mlperf_retinanet.model.image_list import ImageList
+
+@torch.jit.unused
+def _get_shape_onnx(image: Tensor) -> Tensor:
+    from torch.onnx import operators
+    return operators.shape_as_tensor(image)[-2:]
+
+def _resize_image_and_masks(image: Tensor,
+                            target: Optional[Dict[str, Tensor]] = None,
+                            image_size: Optional[Tuple[int, int]] = None,
+                            ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+    image = torch.nn.functional.interpolate(image[None], size=image_size, scale_factor=None, mode='bilinear',
+                                            recompute_scale_factor=None, align_corners=False)[0]
+
+    if target is None:
+        return image, target
+
+    if "masks" in target:
+        mask = target["masks"]
+        mask = torch.nn.functional.interpolate(mask[:, None].float(), size=image_size, scale_factor=None,
+                                               recompute_scale_factor=None)[:, 0].byte()
+        target["masks"] = mask
+    return image, target
+
+class GeneralizedRCNNTransform(nn.Module):
+    """
+    Performs input / target transformation before feeding the data to a GeneralizedRCNN
+    model.
+
+    The transformations it perform are:
+        - input normalization (mean subtraction and std division)
+        - input / target resizing to match image_size
+
+    It returns a ImageList for the inputs, and a List[Dict[Tensor]] for the targets
+    """
+
+    def __init__(self, image_size: Optional[Tuple[int, int]],
+                 image_mean: List[float], image_std: List[float],):
+        super(GeneralizedRCNNTransform, self).__init__()
+        self.image_size = image_size
+        self.image_mean = image_mean
+        self.image_std = image_std
+
+    def forward(self,
+                images: List[Tensor],
+                targets: Optional[List[Dict[str, Tensor]]] = None
+                ) -> Tuple[ImageList, Optional[List[Dict[str, Tensor]]]]:
+        images = list(img for img in images)
+        if targets is not None:
+            # make a copy of targets to avoid modifying it in-place
+            # once torchscript supports dict comprehension
+            # this can be simplified as follows
+            # targets = [{k: v for k,v in t.items()} for t in targets]
+            targets_copy: List[Dict[str, Tensor]] = []
+            for t in targets:
+                data: Dict[str, Tensor] = {}
+                for k, v in t.items():
+                    data[k] = v
+                targets_copy.append(data)
+            targets = targets_copy
+        for i in range(len(images)):
+            image = images[i]
+            target_index = targets[i] if targets is not None else None
+
+            if image.dim() != 3:
+                raise ValueError("images is expected to be a list of 3d tensors "
+                                 "of shape [C, H, W], got {}".format(image.shape))
+            image = self.normalize(image)
+            image, target_index = self.resize(image, target_index)
+            images[i] = image
+            if targets is not None and target_index is not None:
+                targets[i] = target_index
+
+        image_sizes = [img.shape[-2:] for img in images]
+        images = torch.stack(images)
+        image_sizes_list: List[Tuple[int, int]] = []
+        for image_size in image_sizes:
+            assert len(image_size) == 2
+            image_sizes_list.append((image_size[0], image_size[1]))
+
+        image_list = ImageList(images, image_sizes_list)
+        return image_list, targets
+
+    def normalize(self, image: Tensor) -> Tensor:
+        if not image.is_floating_point():
+            raise TypeError(
+                f"Expected input images to be of floating type (in range [0, 1]), "
+                f"but found type {image.dtype} instead"
+            )
+        dtype, device = image.dtype, image.device
+        mean = torch.as_tensor(self.image_mean, dtype=dtype, device=device)
+        std = torch.as_tensor(self.image_std, dtype=dtype, device=device)
+        return (image - mean[:, None, None]) / std[:, None, None]
+
+    def torch_choice(self, k: List[int]) -> int:
+        """
+        Implements `random.choice` via torch ops so it can be compiled with
+        TorchScript. Remove if https://github.com/pytorch/pytorch/issues/25803
+        is fixed.
+        """
+        index = int(torch.empty(1).uniform_(0., float(len(k))).item())
+        return k[index]
+
+    def resize(self,
+               image: Tensor,
+               target: Optional[Dict[str, Tensor]] = None,
+               ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        h, w = image.shape[-2:]
+        image, target = _resize_image_and_masks(image, target, self.image_size)
+
+        if target is None:
+            return image, target
+
+        bbox = target["boxes"]
+        bbox = resize_boxes(bbox, (h, w), image.shape[-2:])
+        target["boxes"] = bbox
+
+        return image, target
+
+    def postprocess(self,
+                    result: List[Dict[str, Tensor]],
+                    image_shapes: List[Tuple[int, int]],
+                    original_image_sizes: List[Tuple[int, int]]
+                    ) -> List[Dict[str, Tensor]]:
+        if self.training:
+            return result
+        for i, (pred, im_s, o_im_s) in enumerate(zip(result, image_shapes, original_image_sizes)):
+            boxes = pred["boxes"]
+            boxes = resize_boxes(boxes, im_s, o_im_s)
+            result[i]["boxes"] = boxes
+        return result
+
+    def __repr__(self) -> str:
+        format_string = self.__class__.__name__ + '('
+        _indent = '\n    '
+        format_string += "{0}Normalize(mean={1}, std={2})".format(_indent, self.image_mean, self.image_std)
+        format_string += "{0}Resize(height={1}, width={2}, mode='bilinear')".format(_indent, self.image_size[0],
+                                                                                         self.image_size[1])
+        format_string += '\n)'
+        return format_string
+
+def resize_boxes(boxes: Tensor, original_size: List[int], new_size: List[int]) -> Tensor:
+    ratios = [
+        torch.tensor(s, dtype=torch.float32, device=boxes.device) /
+        torch.tensor(s_orig, dtype=torch.float32, device=boxes.device)
+        for s, s_orig in zip(new_size, original_size)
+    ]
+    ratio_height, ratio_width = ratios
+    xmin, ymin, xmax, ymax = boxes.unbind(1)
+
+    xmin = xmin * ratio_width
+    xmax = xmax * ratio_width
+    ymin = ymin * ratio_height
+    ymax = ymax * ratio_height
+    res = torch.stack((xmin, ymin, xmax, ymax), dim=1)
+    return res

tinygrad_repo/test/external/mlperf_retinanet/model/utils.py

@@ -0,0 +1,123 @@
+# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/model/utils.py
+
+import torch
+
+class Matcher(object):
+    """
+    This class assigns to each predicted "element" (e.g., a box) a ground-truth
+    element. Each predicted element will have exactly zero or one matches; each
+    ground-truth element may be assigned to zero or more predicted elements.
+
+    Matching is based on the MxN match_quality_matrix, that characterizes how well
+    each (ground-truth, predicted)-pair match. For example, if the elements are
+    boxes, the matrix may contain box IoU overlap values.
+
+    The matcher returns a tensor of size N containing the index of the ground-truth
+    element m that matches to prediction n. If there is no match, a negative value
+    is returned.
+    """
+
+    BELOW_LOW_THRESHOLD = -1
+    BETWEEN_THRESHOLDS = -2
+
+    __annotations__ = {
+        'BELOW_LOW_THRESHOLD': int,
+        'BETWEEN_THRESHOLDS': int,
+    }
+
+    def __init__(self, high_threshold, low_threshold, allow_low_quality_matches=False):
+        # type: (float, float, bool) -> None
+        """
+        Args:
+            high_threshold (float): quality values greater than or equal to
+                this value are candidate matches.
+            low_threshold (float): a lower quality threshold used to stratify
+                matches into three levels:
+                1) matches >= high_threshold
+                2) BETWEEN_THRESHOLDS matches in [low_threshold, high_threshold)
+                3) BELOW_LOW_THRESHOLD matches in [0, low_threshold)
+            allow_low_quality_matches (bool): if True, produce additional matches
+                for predictions that have only low-quality match candidates. See
+                set_low_quality_matches_ for more details.
+        """
+        self.BELOW_LOW_THRESHOLD = -1
+        self.BETWEEN_THRESHOLDS = -2
+        assert low_threshold <= high_threshold
+        self.high_threshold = high_threshold
+        self.low_threshold = low_threshold
+        self.allow_low_quality_matches = allow_low_quality_matches
+
+    def __call__(self, match_quality_matrix):
+        """
+        Args:
+            match_quality_matrix (Tensor[float]): an MxN tensor, containing the
+            pairwise quality between M ground-truth elements and N predicted elements.
+
+        Returns:
+            matches (Tensor[int64]): an N tensor where N[i] is a matched gt in
+            [0, M - 1] or a negative value indicating that prediction i could not
+            be matched.
+        """
+        if match_quality_matrix.numel() == 0:
+            # empty targets or proposals not supported during training
+            if match_quality_matrix.shape[0] == 0:
+                raise ValueError(
+                    "No ground-truth boxes available for one of the images "
+                    "during training")
+
+            raise ValueError(
+                "No proposal boxes available for one of the images "
+                "during training")
+
+        # match_quality_matrix is M (gt) x N (predicted)
+        # Max over gt elements (dim 0) to find best gt candidate for each prediction
+        matched_vals, matches = match_quality_matrix.max(dim=0)
+        if self.allow_low_quality_matches:
+            all_matches = matches.clone()
+        else:
+            all_matches = None
+
+        # Assign candidate matches with low quality to negative (unassigned) values
+        below_low_threshold = matched_vals < self.low_threshold
+        between_thresholds = (matched_vals >= self.low_threshold) & (
+            matched_vals < self.high_threshold
+        )
+        matches[below_low_threshold] = self.BELOW_LOW_THRESHOLD
+        matches[between_thresholds] = self.BETWEEN_THRESHOLDS
+
+        if self.allow_low_quality_matches:
+            assert all_matches is not None
+            self.set_low_quality_matches_(matches, all_matches, match_quality_matrix)
+
+        return matches
+
+    def set_low_quality_matches_(self, matches, all_matches, match_quality_matrix):
+        """
+        Produce additional matches for predictions that have only low-quality matches.
+        Specifically, for each ground-truth find the set of predictions that have
+        maximum overlap with it (including ties); for each prediction in that set, if
+        it is unmatched, then match it to the ground-truth with which it has the highest
+        quality value.
+        """
+        # For each gt, find the prediction with which it has highest quality
+        highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)
+        # Find highest quality match available, even if it is low, including ties
+        gt_pred_pairs_of_highest_quality = torch.where(
+            match_quality_matrix == highest_quality_foreach_gt[:, None]
+        )
+        # Example gt_pred_pairs_of_highest_quality:
+        #   tensor([[    0, 39796],
+        #           [    1, 32055],
+        #           [    1, 32070],
+        #           [    2, 39190],
+        #           [    2, 40255],
+        #           [    3, 40390],
+        #           [    3, 41455],
+        #           [    4, 45470],
+        #           [    5, 45325],
+        #           [    5, 46390]])
+        # Each row is a (gt index, prediction index)
+        # Note how gt items 1, 2, 3, and 5 each have two ties
+
+        pred_inds_to_update = gt_pred_pairs_of_highest_quality[1]
+        matches[pred_inds_to_update] = all_matches[pred_inds_to_update]

tinygrad_repo/test/external/mlperf_retinanet/openimages.py

@@ -0,0 +1,25 @@
+from test.external.mlperf_retinanet.model.boxes import box_iou
+from test.external.mlperf_retinanet.model.utils import Matcher
+
+import torch
+
+# This applies the filtering in https://github.com/mlcommons/training/blob/cdd928d4596c142c15a7d86b2eeadbac718c8da2/single_stage_detector/ssd/model/retinanet.py#L117
+# and https://github.com/mlcommons/training/blob/cdd928d4596c142c15a7d86b2eeadbac718c8da2/single_stage_detector/ssd/model/retinanet.py#L203
+# to match with tinygrad's dataloader implementation.
+def postprocess_targets(targets, anchors):
+    proposal_matcher, matched_idxs = Matcher(0.5, 0.4, allow_low_quality_matches=True), []
+    for anchors_per_image, targets_per_image in zip(anchors, targets):
+        if targets_per_image['boxes'].numel() == 0:
+            matched_idxs.append(torch.full((anchors_per_image.size(0),), -1, dtype=torch.int64,
+                                            device=anchors_per_image.device))
+            continue
+
+        match_quality_matrix = box_iou(targets_per_image['boxes'], anchors_per_image)
+        matched_idxs.append(proposal_matcher(match_quality_matrix))
+
+    for targets_per_image, matched_idxs_per_image in zip(targets, matched_idxs):
+        foreground_idxs_per_image = matched_idxs_per_image >= 0
+        targets_per_image["boxes"] = targets_per_image["boxes"][matched_idxs_per_image[foreground_idxs_per_image]]
+        targets_per_image["labels"] = targets_per_image["labels"][matched_idxs_per_image[foreground_idxs_per_image]]
+
+    return targets

tinygrad_repo/test/external/mlperf_retinanet/presets.py

@@ -0,0 +1,24 @@
+# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/presets.py
+
+from test.external.mlperf_retinanet import transforms as T
+
+class DetectionPresetTrain:
+    def __init__(self, data_augmentation, hflip_prob=0.5, mean=(123., 117., 104.)):
+        if data_augmentation == 'hflip':
+            self.transforms = T.Compose([
+                T.RandomHorizontalFlip(p=hflip_prob),
+                T.ToTensor(),
+            ])
+        else:
+            raise ValueError(f'Unknown data augmentation policy "{data_augmentation}"')
+
+    def __call__(self, img, target):
+        return self.transforms(img, target)
+
+
+class DetectionPresetEval:
+    def __init__(self):
+        self.transforms = T.ToTensor()
+
+    def __call__(self, img, target):
+        return self.transforms(img, target)

tinygrad_repo/test/external/mlperf_retinanet/transforms.py

@@ -0,0 +1,77 @@
+# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/transforms.py
+
+import torch
+import torchvision
+
+from torch import nn, Tensor
+from torchvision.transforms import functional as F
+from torchvision.transforms import transforms as T
+from typing import List, Tuple, Dict, Optional
+
+from PIL import Image
+Image.MAX_IMAGE_PIXELS = None
+from typing import Any
+
+try:
+    import accimage
+except ImportError:
+    accimage = None
+
+@torch.jit.unused
+def _is_pil_image(img: Any) -> bool:
+    if accimage is not None:
+        return isinstance(img, (Image.Image, accimage.Image))
+    else:
+        return isinstance(img, Image.Image)
+
+def get_image_size_tensor(img: Tensor) -> List[int]:
+    # Returns (w, h) of tensor image
+    torchvision.transforms._functional_tensor._assert_image_tensor(img)
+    return [img.shape[-1], img.shape[-2]]
+
+@torch.jit.unused
+def get_image_size_pil(img: Any) -> List[int]:
+    if _is_pil_image(img):
+        return list(img.size)
+    raise TypeError("Unexpected type {}".format(type(img)))
+
+def get_image_size(img: Tensor) -> List[int]:
+    """Returns the size of an image as [width, height].
+    Args:
+        img (PIL Image or Tensor): The image to be checked.
+    Returns:
+        List[int]: The image size.
+    """
+    if isinstance(img, torch.Tensor):
+        return get_image_size_tensor(img)
+
+    return get_image_size_pil(img)
+
+class Compose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, image, target):
+        for t in self.transforms:
+            image, target = t(image, target)
+        return image, target
+
+
+class RandomHorizontalFlip(T.RandomHorizontalFlip):
+    def forward(self, image: Tensor,
+                target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if torch.rand(1) < self.p:
+            image = F.hflip(image)
+            if target is not None:
+                width, _ = get_image_size(image)
+                target["boxes"][:, [0, 2]] = width - target["boxes"][:, [2, 0]]
+                if "masks" in target:
+                    target["masks"] = target["masks"].flip(-1)
+        return image, target
+
+
+class ToTensor(nn.Module):
+    def forward(self, image: Tensor,
+                target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        image = F.to_tensor(image)
+        return image, target

tinygrad_repo/test/external/mlperf_stable_diffusion/external_test_eval.py

@@ -0,0 +1,58 @@
+import unittest, os
+import numpy as np
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from tinygrad import Device, Tensor
+from tinygrad.helpers import getenv, Context
+from tinygrad.nn.state import safe_save, torch_load, get_parameters
+from examples.mlperf.model_eval import eval_stable_diffusion, vae_decode
+from examples.stable_diffusion import AutoencoderKL
+
+def set_eval_params():
+  # override these as needed from cli
+  for k,v in {"MODEL": "stable_diffusion", "GPUS": "8", "EVAL_SAMPLES": "600", "CONTEXT_BS": "816", "DENOISE_BS": "600", "DECODE_BS": "384",
+   "INCEPTION_BS": "560", "CLIP_BS": "240", "DATADIR": "/raid/datasets/stable_diffusion", "CKPTDIR": "/raid/weights/stable_diffusion"}.items():
+    os.environ[k] = getenv(k, v)
+
+class TestEval(unittest.TestCase):
+  def test_eval_ckpt(self):
+    set_eval_params()
+    with TemporaryDirectory(prefix="test-eval") as tmp:
+      os.environ["EVAL_CKPT_DIR"] = tmp
+      # NOTE Although this checkpoint has the original fully trained model from StabilityAI, we are using mlperf code that uses different
+      #   GroupNorm num_groups. Therefore, eval results may not reflect eval results on the original model.
+      # The purpose of using this checkpoint is to have reproducible eval outputs.
+      # Eval code expects file and weight names in a specific format, as .safetensors (not .ckpt), which is why we resave the checkpoint
+      sd_v2 = torch_load(Path(getenv("CKPTDIR", "")) / "sd" / "512-base-ema.ckpt")["state_dict"]
+      sd_v2 = {k.replace("model.diffusion_model.", "", 1): v for k,v in sd_v2.items() if k.startswith("model.diffusion_model.")}
+      safe_save(sd_v2, f"{tmp}/0.safetensors")
+      clip, fid, ckpt = eval_stable_diffusion()
+    assert ckpt == 0
+    if Device.DEFAULT == "NULL":
+      assert clip == 0
+      assert fid > 0 and fid < 1000
+    else:
+      # observed:
+      # clip=0.08369670808315277, fid=301.05236173709545 (if SEED=12345, commit=c01b2c93076e80ae6d1ebca64bb8e83a54dadba6)
+      # clip=0.08415728807449341, fid=300.3710877072948 (if SEED=12345, commit=179c7fcfe132f1a6344b57c9d8cef4eded586867)
+      # clip=0.0828116238117218, fid=301.241909555543 (if SEED=98765, commit=c01b2c93076e80ae6d1ebca64bb8e83a54dadba6)
+      np.testing.assert_allclose(fid, 301.147, rtol=0.1, atol=0)
+      np.testing.assert_allclose(clip, 0.08325, rtol=0.1, atol=0)
+
+  # only tested on 8xMI300x system
+  @unittest.skipUnless(getenv("HANG_OK"), "expected to hang")
+  def test_decoder_beam_hang(self):
+    set_eval_params()
+    for k,v in {"BEAM": "2", "HCQDEV_WAIT_TIMEOUT_MS": "300000", "BEAM_UOPS_MAX": "8000", "BEAM_UPCAST_MAX": "256", "BEAM_LOCAL_MAX": "1024",
+                "BEAM_MIN_PROGRESS": "5", "IGNORE_JIT_FIRST_BEAM": "1"}.items():
+      os.environ[k] = getenv(k, v)
+    with Context(BEAM=int(os.environ["BEAM"])): # necessary because helpers.py has already set BEAM=0 and cached getenv for "BEAM"
+      GPUS = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 8))]
+      vae = AutoencoderKL()
+      for p in get_parameters(vae): p.to_(GPUS).realize()
+      x = Tensor.zeros(48,4,64,64).contiguous().to(GPUS).realize()
+      x.uop = x.uop.multi(0)
+      for _ in range(2): vae_decode(x, vae)
+
+if __name__=="__main__":
+  unittest.main()

tinygrad_repo/test/external/mlperf_stable_diffusion/external_test_models.py

@@ -0,0 +1,114 @@
+import unittest
+import numpy as np
+from pathlib import Path
+from tinygrad import Tensor, dtypes, Device
+from tinygrad.helpers import getenv
+from tinygrad.nn.state import get_parameters
+from extra.models import clip
+from examples.mlperf.initializers import gelu_erf, init_stable_diffusion, attn_f32_softmax
+from typing import Literal
+
+clip_params = {"dims": 1024, "n_heads": 16, "layers": 24, "return_pooled": False, "ln_penultimate": True, "clip_tokenizer_version": "sd_mlperf_v5_0"}
+def get_cond_stage_model(GPUS:list[str]|None=None) -> clip.FrozenOpenClipEmbedder:
+  clip.gelu = gelu_erf
+  model = clip.FrozenOpenClipEmbedder(**clip_params)
+  if GPUS and len(GPUS) > 1:
+    for p in get_parameters(model): p.to_(GPUS)
+  return model
+def get_tokens(BS:int) -> Tensor: return Tensor([0] * 77 * BS, dtype=dtypes.int32).reshape(-1, 77)
+
+class TestOpenClip(unittest.TestCase):
+  def test_tokenizer(self):
+    prompt = "Beautiful is better than ugly.\nExplicit is better than implicit.\nSimple is better than complex.\nComplex is better than complicated."
+    model = get_cond_stage_model()
+    tokens = model.tokenizer.encode(prompt, pad_with_zeros=True)
+    expected = [49406, 1215, 533, 1539, 1126, 8159, 269, 33228, 533, 1539, 1126, 15269, 585, 269, 4129, 533, 1539, 1126, 6324, 269, 6324, 533,
+                1539, 1126, 16621, 269, 49407] + [0]*50
+    self.assertEqual(tokens, expected)
+
+  def test_clip_gelu_init(self):
+    for resblock in get_cond_stage_model().model.transformer.resblocks:
+      self.assertEqual(resblock.mlp.gelu, gelu_erf)
+
+  def test_multigpu_clip_embed(self):
+    BS = 304
+    GPUS = [f"{Device.DEFAULT}:{i}" for i in range(8)]
+    model = get_cond_stage_model(GPUS)
+    tokens = get_tokens(BS)
+    embeds = model.embed_tokens(tokens.shard(GPUS, axis=0)).realize()
+    self.assertEqual(embeds.shape, (BS, 77, 1024))
+    self.assertEqual(embeds.dtype, dtypes.float32)
+
+  def test_multigpu_clip_score(self):
+    BS = 240
+    GPUS = [f"{Device.DEFAULT}:{i}" for i in range(8)]
+    vision_cfg = {'width': 1280, 'layers': 32, 'd_head': 80, 'image_size': 224, 'patch_size': 14}
+    text_cfg = {'width': 1024, 'n_heads': 16, 'layers': 24, 'vocab_size': 49408, 'ctx_length': 77}
+    clip.gelu = gelu_erf
+    clip_encoder = clip.OpenClipEncoder(1024, text_cfg, vision_cfg)
+    for p in get_parameters(clip_encoder): p.to_(GPUS)
+    tokens = get_tokens(BS)
+    imgs = Tensor.zeros(BS,3,224,224).contiguous()
+    scores = clip_encoder.get_clip_score(tokens.shard(GPUS, axis=0), imgs.shard(GPUS, axis=0)).realize()
+    self.assertEqual(scores.shape, (BS,))
+    self.assertEqual(scores.dtype, dtypes.float32)
+
+class TestInitStableDiffusion(unittest.TestCase):
+  def setUp(self):
+    # NOTE: set env variable based on where checkpoints are on the system
+    self.CKPTDIR = Path(getenv("CKPTDIR", "/raid/weights/stable_diffusion"))
+
+  def helper_test_init(self, version:Literal["v2-mlperf-train", "v2-mlperf-eval"]):
+    model, unet, sqrt_acp, sqrt_omacp = init_stable_diffusion(version, self.CKPTDIR / "sd" / "512-base-ema.ckpt", ["CPU"])
+
+    with self.subTest("test that StableDiffusion has correct models"):
+      self.assertEqual(model.model.diffusion_model, unet)
+      has_encoder = True if version=="v2-mlperf-eval" else False
+      self.assertEqual(hasattr(model, "first_stage_model"), has_encoder, "only the eval model uses the encoder")
+      self.assertTrue(isinstance(model.cond_stage_model, clip.FrozenOpenClipEmbedder))
+
+    with self.subTest("test for mlperf unique attributes"):
+      self.assertEqual(model.cond_stage_model.tokenizer.version, 'sd_mlperf_v5_0')
+      self.assertEqual(unet.out[0].num_groups, 16)
+      self.assertEqual(unet.input_blocks[1][1].norm.eps, 1e-6)
+      self.assertEqual(unet.input_blocks[1][1].transformer_blocks[0].attn1.attn, attn_f32_softmax)
+
+    with self.subTest("test loaded clip parameters"):
+      sample = model.cond_stage_model.model.transformer.resblocks[8].mlp.c_fc.bias.flatten()[42:46].numpy()
+      expected = np.array([-0.49812260270118713, -0.3039605915546417, -0.40284937620162964, -0.45069342851638794], dtype=np.float32)
+      np.testing.assert_allclose(sample, expected, rtol=1e-7, atol=0, err_msg="loaded clip parameters are incorrect")
+
+    if version=="v2-mlperf-train":
+      with self.subTest("test that zero_module worked"):
+        self.assertTrue((unet.out[2].weight == 0).all().item(), "expected all zeroes")
+        self.assertTrue((unet.out[2].bias == 0).all().item(), "expected all zeroes")
+    elif version=="v2-mlperf-eval":
+      with self.subTest("test loaded vae parameters"):
+        sample = model.first_stage_model.decoder.up[0]['block'][1].conv2.weight.flatten()[42:46].numpy()
+        expected = np.array([0.08192943036556244, 0.040095631033182144, 0.07541035860776901, 0.1475081741809845], dtype=np.float32)
+        np.testing.assert_allclose(sample, expected, rtol=1e-7, atol=0, err_msg="loaded vae parameters are incorrect")
+
+    with self.subTest("check schedules"):
+      expected = np.array([0.9995748996734619, 0.06826484948396683], dtype=np.float32)
+      np.testing.assert_allclose(sqrt_acp[[0,-1]].numpy(), expected, rtol=1e-7, atol=0, err_msg="sqrt_acp is incorrect")
+      expected = np.array([0.029155133292078972, 0.9976672530174255], dtype=np.float32)
+      np.testing.assert_allclose(sqrt_omacp[[0,-1]].numpy(), expected, rtol=1e-7, atol=0, err_msg="sqrt_omacp is incorrect")
+
+    with self.subTest("check mixed precision"):
+      out = unet.input_blocks[2][1].proj_in(Tensor.randn(320, dtype=dtypes.float32))
+      self.assertEqual(out.dtype, dtypes.bfloat16, "expected float32 to be downcast to bfloat16 by Linear")
+      out = unet.out[2](Tensor.randn(304,320,64,64, dtype=dtypes.float32))
+      self.assertEqual(out.dtype, dtypes.bfloat16, "expected float32 to be downcast to bfloat16 by Conv2d")
+      out = unet.input_blocks[1][1].transformer_blocks[0].norm1(Tensor.randn(320, dtype=dtypes.bfloat16))
+      self.assertEqual(out.dtype, dtypes.float32, "expected bfloat16 to be upcast to float32 by LayerNorm")
+      out = unet.input_blocks[5][0].in_layers[0](Tensor.randn(304, 640, dtype=dtypes.bfloat16))
+      self.assertEqual(out.dtype, dtypes.float32, "expected bfloat16 to be upcast to float32 by GroupNorm")
+
+  def test_train_model(self):
+    self.helper_test_init("v2-mlperf-train")
+
+  def test_eval_model(self):
+    self.helper_test_init("v2-mlperf-eval")
+
+if __name__=="__main__":
+  unittest.main()

tinygrad_repo/test/external/mlperf_stable_diffusion/external_test_train.py

@@ -0,0 +1,23 @@
+import unittest, os
+from tempfile import TemporaryDirectory
+from tinygrad import Tensor
+from tinygrad.helpers import getenv
+from examples.mlperf.model_train import train_stable_diffusion
+
+class TestTrain(unittest.TestCase):
+  def test_train_to_ckpt(self):
+    # train for num_steps, save checkpoint, and stop training
+    num_steps = 42
+    os.environ.update({"MODEL": "stable_diffusion", "TOTAL_CKPTS": "1", "CKPT_STEP_INTERVAL": str(num_steps), "GPUS": "8", "BS": "304"})
+    # NOTE: update these based on where data/checkpoints are on your system
+    if not getenv("DATADIR", ""): os.environ["DATADIR"] = "/raid/datasets/stable_diffusion"
+    if not getenv("CKPTDIR", ""): os.environ["CKPTDIR"] = "/raid/weights/stable_diffusion"
+    with TemporaryDirectory(prefix="test-train") as tmp:
+      os.environ["UNET_CKPTDIR"] = tmp
+      with Tensor.train():
+        saved_ckpts = train_stable_diffusion()
+      expected_ckpt = f"{tmp}/{num_steps}.safetensors"
+      assert len(saved_ckpts) == 1 and saved_ckpts[0] == expected_ckpt
+
+if __name__=="__main__":
+  unittest.main()

tinygrad_repo/test/external/mlperf_unet3d/dice.py

@@ -0,0 +1,94 @@
+# https://github.com/mlcommons/training/blob/master/image_segmentation/pytorch/model/losses.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Dice:
+  def __init__(self,
+               to_onehot_y: bool = True,
+               to_onehot_x: bool = False,
+               use_softmax: bool = True,
+               use_argmax: bool = False,
+               include_background: bool = False,
+               layout: str = "NCDHW"):
+    self.include_background = include_background
+    self.to_onehot_y = to_onehot_y
+    self.to_onehot_x = to_onehot_x
+    self.use_softmax = use_softmax
+    self.use_argmax = use_argmax
+    self.smooth_nr = 1e-6
+    self.smooth_dr = 1e-6
+    self.layout = layout
+
+  def __call__(self, prediction, target):
+    if self.layout == "NCDHW":
+      channel_axis = 1
+      reduce_axis = list(range(2, len(prediction.shape)))
+    else:
+      channel_axis = -1
+      reduce_axis = list(range(1, len(prediction.shape) - 1))
+    num_pred_ch = prediction.shape[channel_axis]
+
+    if self.use_softmax:
+      prediction = torch.softmax(prediction, dim=channel_axis)
+    elif self.use_argmax:
+      prediction = torch.argmax(prediction, dim=channel_axis)
+
+    if self.to_onehot_y:
+      target = to_one_hot(target, self.layout, channel_axis)
+
+    if self.to_onehot_x:
+      prediction = to_one_hot(prediction, self.layout, channel_axis)
+
+    if not self.include_background:
+      assert num_pred_ch > 1, \
+          f"To exclude background the prediction needs more than one channel. Got {num_pred_ch}."
+      if self.layout == "NCDHW":
+          target = target[:, 1:]
+          prediction = prediction[:, 1:]
+      else:
+          target = target[..., 1:]
+          prediction = prediction[..., 1:]
+
+    assert (target.shape == prediction.shape), \
+        f"Target and prediction shape do not match. Target: ({target.shape}), prediction: ({prediction.shape})."
+
+    intersection = torch.sum(target * prediction, dim=reduce_axis)
+    target_sum = torch.sum(target, dim=reduce_axis)
+    prediction_sum = torch.sum(prediction, dim=reduce_axis)
+
+    return (2.0 * intersection + self.smooth_nr) / (target_sum + prediction_sum + self.smooth_dr)
+
+
+def to_one_hot(array, layout, channel_axis):
+  if len(array.shape) >= 5:
+    array = torch.squeeze(array, dim=channel_axis)
+  array = F.one_hot(array.long(), num_classes=3)
+  if layout == "NCDHW":
+    array = array.permute(0, 4, 1, 2, 3).float()
+  return array
+
+
+class DiceCELoss(nn.Module):
+  def __init__(self, to_onehot_y, use_softmax, layout, include_background):
+    super(DiceCELoss, self).__init__()
+    self.dice = Dice(to_onehot_y=to_onehot_y, use_softmax=use_softmax, layout=layout,
+                      include_background=include_background)
+    self.cross_entropy = nn.CrossEntropyLoss()
+
+  def forward(self, y_pred, y_true):
+    cross_entropy = self.cross_entropy(y_pred, torch.squeeze(y_true, dim=1).long())
+    dice = torch.mean(1.0 - self.dice(y_pred, y_true))
+    return (dice + cross_entropy) / 2
+
+
+class DiceScore:
+  def __init__(self, to_onehot_y: bool = True, use_argmax: bool = True, layout: str = "NCDHW",
+               include_background: bool = False):
+    self.dice = Dice(to_onehot_y=to_onehot_y, to_onehot_x=True, use_softmax=False,
+                     use_argmax=use_argmax, layout=layout, include_background=include_background)
+
+  def __call__(self, y_pred, y_true):
+    return torch.mean(self.dice(y_pred, y_true), dim=0)

tinygrad_repo/test/external/mlperf_unet3d/kits19.py

@@ -0,0 +1,165 @@
+# copied from https://github.com/mlcommons/training/blob/5c08ce57e7f582cc4558035d8324a2bf4c8ca225/image_segmentation/pytorch/data_loading/pytorch_loader.py
+
+import random
+import numpy as np
+import scipy.ndimage
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+
+def get_train_transforms():
+    rand_flip = RandFlip()
+    cast = Cast(types=(np.float32, np.uint8))
+    rand_scale = RandomBrightnessAugmentation(factor=0.3, prob=0.1)
+    rand_noise = GaussianNoise(mean=0.0, std=0.1, prob=0.1)
+    train_transforms = transforms.Compose([rand_flip, cast, rand_scale, rand_noise])
+    return train_transforms
+
+
+class RandBalancedCrop:
+    def __init__(self, patch_size, oversampling):
+        self.patch_size = patch_size
+        self.oversampling = oversampling
+
+    def __call__(self, data):
+        image, label = data["image"], data["label"]
+        if random.random() < self.oversampling:
+            image, label, cords = self.rand_foreg_cropd(image, label)
+        else:
+            image, label, cords = self._rand_crop(image, label)
+        data.update({"image": image, "label": label})
+        return data
+
+    @staticmethod
+    def randrange(max_range):
+        return 0 if max_range == 0 else random.randrange(max_range)
+
+    def get_cords(self, cord, idx):
+        return cord[idx], cord[idx] + self.patch_size[idx]
+
+    def _rand_crop(self, image, label):
+        ranges = [s - p for s, p in zip(image.shape[1:], self.patch_size)]
+        cord = [self.randrange(x) for x in ranges]
+        low_x, high_x = self.get_cords(cord, 0)
+        low_y, high_y = self.get_cords(cord, 1)
+        low_z, high_z = self.get_cords(cord, 2)
+        image = image[:, low_x:high_x, low_y:high_y, low_z:high_z]
+        label = label[:, low_x:high_x, low_y:high_y, low_z:high_z]
+        return image, label, [low_x, high_x, low_y, high_y, low_z, high_z]
+
+    def rand_foreg_cropd(self, image, label):
+        def adjust(foreg_slice, patch_size, label, idx):
+            diff = patch_size[idx - 1] - (foreg_slice[idx].stop - foreg_slice[idx].start)
+            sign = -1 if diff < 0 else 1
+            diff = abs(diff)
+            ladj = self.randrange(diff)
+            hadj = diff - ladj
+            low = max(0, foreg_slice[idx].start - sign * ladj)
+            high = min(label.shape[idx], foreg_slice[idx].stop + sign * hadj)
+            diff = patch_size[idx - 1] - (high - low)
+            if diff > 0 and low == 0:
+                high += diff
+            elif diff > 0:
+                low -= diff
+            return low, high
+
+        cl = np.random.choice(np.unique(label[label > 0]))
+        foreg_slices = scipy.ndimage.find_objects(scipy.ndimage.measurements.label(label==cl)[0])
+        foreg_slices = [x for x in foreg_slices if x is not None]
+        slice_volumes = [np.prod([s.stop - s.start for s in sl]) for sl in foreg_slices]
+        slice_idx = np.argsort(slice_volumes)[-2:]
+        foreg_slices = [foreg_slices[i] for i in slice_idx]
+        if not foreg_slices:
+            return self._rand_crop(image, label)
+        foreg_slice = foreg_slices[random.randrange(len(foreg_slices))]
+        low_x, high_x = adjust(foreg_slice, self.patch_size, label, 1)
+        low_y, high_y = adjust(foreg_slice, self.patch_size, label, 2)
+        low_z, high_z = adjust(foreg_slice, self.patch_size, label, 3)
+        image = image[:, low_x:high_x, low_y:high_y, low_z:high_z]
+        label = label[:, low_x:high_x, low_y:high_y, low_z:high_z]
+        return image, label, [low_x, high_x, low_y, high_y, low_z, high_z]
+
+
+class RandFlip:
+    def __init__(self):
+        self.axis = [1, 2, 3]
+        self.prob = 1 / len(self.axis)
+
+    def flip(self, data, axis):
+        data["image"] = np.flip(data["image"], axis=axis).copy()
+        data["label"] = np.flip(data["label"], axis=axis).copy()
+        return data
+
+    def __call__(self, data):
+        for axis in self.axis:
+            if random.random() < self.prob:
+                data = self.flip(data, axis)
+        return data
+
+
+class Cast:
+    def __init__(self, types):
+        self.types = types
+
+    def __call__(self, data):
+        data["image"] = data["image"].astype(self.types[0])
+        data["label"] = data["label"].astype(self.types[1])
+        return data
+
+
+class RandomBrightnessAugmentation:
+    def __init__(self, factor, prob):
+        self.prob = prob
+        self.factor = factor
+
+    def __call__(self, data):
+        image = data["image"]
+        if random.random() < self.prob:
+            factor = np.random.uniform(low=1.0-self.factor, high=1.0+self.factor, size=1)
+            image = (image * (1 + factor)).astype(image.dtype)
+            data.update({"image": image})
+        return data
+
+
+class GaussianNoise:
+    def __init__(self, mean, std, prob):
+        self.mean = mean
+        self.std = std
+        self.prob = prob
+
+    def __call__(self, data):
+        image = data["image"]
+        if random.random() < self.prob:
+            scale = np.random.uniform(low=0.0, high=self.std)
+            noise = np.random.normal(loc=self.mean, scale=scale, size=image.shape).astype(image.dtype)
+            data.update({"image": image + noise})
+        return data
+
+
+class PytTrain(Dataset):
+    def __init__(self, images, labels, **kwargs):
+        self.images, self.labels = images, labels
+        self.train_transforms = get_train_transforms()
+        patch_size, oversampling = kwargs["patch_size"], kwargs["oversampling"]
+        self.patch_size = patch_size
+        self.rand_crop = RandBalancedCrop(patch_size=patch_size, oversampling=oversampling)
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        data = {"image": np.load(self.images[idx]), "label": np.load(self.labels[idx])}
+        data = self.rand_crop(data)
+        data = self.train_transforms(data)
+        return data["image"], data["label"]
+
+
+class PytVal(Dataset):
+    def __init__(self, images, labels):
+        self.images, self.labels = images, labels
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        return np.load(self.images[idx]), np.load(self.labels[idx])

tinygrad_repo/test/external/process_replay/README.md

@@ -0,0 +1,19 @@
+# Process replay tests
+
+Process replay is a tool for creating a diff of generated kernels between two commits. By default, process replay doesn't assert kernel diffs.
+
+Refactor and speedup PRs must enable the assert by including `[pr]` in the pull request title.
+
+Note that process replay [early stops when over 20% of kernels change, for speed.](https://github.com/tinygrad/tinygrad/pull/5480).
+
+## Running locally
+
+To run process replay locally:
+
+(optional: clear previous process replay runs with `test/external/process_replay/reset.py`)
+
+1. Run tests with `CAPTURE_PROCESS_REPLAY=1` in your branch. This will pickle process inputs to CACHEDB.
+2. Checkout master
+3. Run `test/external/process_replay/process_replay.py`
+
+For reference, see `test/external/process_replay/local.sh`.

tinygrad_repo/test/external/process_replay/local.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+set -e
+HEAD=$(git rev-parse --abbrev-ref HEAD)
+python test/external/process_replay/reset.py
+CAPTURE_PROCESS_REPLAY=1 python test/backend/test_ops.py TestOps.test_add
+git checkout master
+git checkout $HEAD -- test/external/process_replay/process_replay.py
+ASSERT_PROCESS_REPLAY=${ASSERT_PROCESS_REPLAY:-1} python test/external/process_replay/process_replay.py
+git checkout $HEAD

tinygrad_repo/test/external/process_replay/process_replay.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python3
+# compare kernels created by HEAD against master
+import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, functools, base64, codecs
+from dataclasses import replace
+from typing import Callable, Any
+
+ASSERT_DIFF = int((flag:="[PR]") in os.getenv("COMMIT_MESSAGE", flag) or flag in os.getenv("PR_TITLE", flag))
+if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0
+
+try:
+  from tinygrad.renderer import Renderer
+  from tinygrad.codegen import to_program
+  from tinygrad.uop.ops import UOp, Ops
+  from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
+except ImportError as e:
+  print(repr(e))
+  exit(int(ASSERT_DIFF))
+
+# *** process replay settings
+
+# internal
+PAGE_SIZE = getenv("PAGE_SIZE", 100)
+REF = os.getenv("GITHUB_REF_NAME", "")
+MAX_DIFF_PCT = getenv("PROCESS_REPLAY_MAX_DIFF_PCT", 20)
+TABLE_NAME = f"process_replay_{VERSION}"
+os.environ["CAPTURE_PROCESS_REPLAY"] = "0"
+early_stop = multiprocessing.Event()
+logging.basicConfig(level=logging.INFO, format="%(message)s")
+MAX_LINES = 500
+def trunc_log(x):
+  if len(lines:=(x if isinstance(x, str) else repr(x)).splitlines()) > MAX_LINES:
+    lines = lines[:MAX_LINES]+[f"WARN: truncated string with {len(lines)} lines"]
+  logging.info("\n".join(lines))
+
+# user config
+SKIP_PROCESS_REPLAY = (k:="[skip_process_replay]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", "")
+# uncomment this to disable by default
+#SKIP_PROCESS_REPLAY = not ASSERT_DIFF and not ((k:="[p]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", ""))
+if REF == "master": SKIP_PROCESS_REPLAY = True
+class ProcessReplayWarning(Warning): pass
+
+# *** replay the function and convert return values to string
+
+def replay_to_program(p:UOp, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
+  if ast.op is Ops.PROGRAM: input_ast = ast
+  else:
+    sink_arg = ast.arg
+    if sink_arg.beam: sink_arg = replace(sink_arg, opts_to_apply=p.src[0].arg.applied_opts)
+    input_ast = ast.replace(arg=replace(sink_arg, name=p.src[0].arg.name))
+  p2 = to_program(input_ast, renderer=renderer)
+  device = p.src[1].arg
+  def to_str(ret:UOp) -> str:
+    src = ret.src[3].arg
+    # PYTHON renderer pickles UOps, first unpickle and decode here
+    if device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(src))])
+    return src
+  # properly color the name arg
+  ast_repr = codecs.decode(str(input_ast), "unicode_escape")
+  return to_str(p2), to_str(p), (ast_repr, renderer)
+
+replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {}
+replayers["do_to_program"] = replay_to_program
+
+# *** run replayers on captured rows and print diffs
+
+def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
+  if ASSERT_DIFF: warnings.filterwarnings("error", category=ProcessReplayWarning)
+  if early_stop.is_set(): return None
+  conn = db_connection()
+  cur = conn.cursor()
+  cur.execute(f"SELECT val FROM '{TABLE_NAME}' LIMIT ? OFFSET ?", (PAGE_SIZE, offset))
+  changed = 0
+  for row in cur.fetchall():
+    if changed > MAX_DIFF_PCT:
+      warnings.warn(f"detected changes in over {MAX_DIFF_PCT}%. skipping further diff generation.", ProcessReplayWarning)
+      early_stop.set()
+      break
+    name, loc = "", ""
+    try:
+      name, args, kwargs, ctx_vals, loc, ret = pickle.loads(row[0])
+      ctx_vars = {k:v.value for k,v in ctx_vals.items() if k not in ("DEBUG", "CAPTURE_PROCESS_REPLAY")
+                  and (var:=ContextVar._cache.get(k)) is not None and var.value != v.value}
+      if (replayer:=fxns.get(name)) is None: continue
+      with Context(**ctx_vars):
+        if (ret:=replayer(ret, *args, **kwargs)) is None: continue
+        good, compare, metadata = ret
+      if good != compare:
+        for m in metadata: trunc_log(m)
+        logging.info(loc)
+        for line in difflib.unified_diff(good.splitlines(), compare.splitlines()):
+          logging.info(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None))
+        if ctx_vars: logging.info(ctx_vars)
+        warnings.warn("PROCESS REPLAY DETECTED CHANGE", ProcessReplayWarning)
+    except Exception as e:
+      changed += 1
+      warnings.warn(f"{name=} {loc=} {e=}", ProcessReplayWarning)
+  cur.close()
+
+# *** generic runner to map rows of a table to a function in parallel
+
+def _pmap(fxns:dict[str, Callable]) -> None:
+  conn = db_connection()
+  cur = conn.cursor()
+  try: row_count = cur.execute(f"select count(*) from '{TABLE_NAME}'").fetchone()[0]
+  except sqlite3.OperationalError:
+    raise RuntimeError(f"{TABLE_NAME} isn't accessible in master, did DB_VERSION change?")
+  finally:
+    cur.close()
+
+  with multiprocessing.get_context("spawn").Pool(multiprocessing.cpu_count()) as pool:
+    bar = tqdm(total=row_count)
+    for _ in pool.imap_unordered(functools.partial(diff, fxns=fxns), range(0, row_count, s:=min(PAGE_SIZE, row_count))): bar.update(s)
+    pool.close()
+    pool.join()
+    pool.terminate()
+
+# *** main loop
+
+if __name__ == "__main__":
+  if SKIP_PROCESS_REPLAY:
+    logging.info("skipping process replay.")
+    exit(0)
+
+  logging.info(f"running process replay with {ASSERT_DIFF=}")
+  try: _pmap(replayers)
+  except Exception as e:
+    logging.info(f"process replay err: {e}")
+    exit(int(ASSERT_DIFF))

tinygrad_repo/test/external/process_replay/reset.py

@@ -0,0 +1,4 @@
+#!/usr/bin/env python3
+from tinygrad.helpers import db_connection, VERSION
+cur = db_connection()
+cur.execute(f"drop table if exists process_replay_{VERSION}")

tinygrad_repo/test/external/sglang_llama/external_llama_eval.py

@@ -0,0 +1,100 @@
+from lm_eval import simple_evaluate
+from lm_eval.api.instance import Instance
+from lm_eval.api.model import LM
+from lm_eval.tasks import TaskManager
+from pathlib import Path
+import json, argparse
+
+from examples.llama3 import build_transformer, Tokenizer, MODEL_PARAMS
+from tinygrad import Tensor, Device
+from tinygrad.helpers import tqdm
+
+class LLaMaAdaptor(LM):
+  def __init__(
+    self,
+    model_size: str,
+    checkpoint_path: Path,
+    max_length: int,
+    quantize: str | None,
+  ):
+    super().__init__()
+    self.max_length = max_length
+    self.tokenizer = Tokenizer(str((checkpoint_path if checkpoint_path.is_dir() else checkpoint_path.parent) / "tokenizer.model"))
+    self.model = build_transformer(checkpoint_path, model_size=model_size, quantize=quantize, max_context=self.max_length)
+    self.last_seen_toks = []
+  def _prefill(self, toks, temperature) -> int:
+    start_pos = 0
+    # we can skip part of the prompt if it is the same as last
+    for i, (a, b) in enumerate(zip(toks, self.last_seen_toks)):
+      if a != b: break
+    else: i = min(len(toks), len(self.last_seen_toks))
+    start_pos += i
+    self.last_seen_toks = toks
+    toks = toks[i:]
+
+    # prefill the model
+    for tok in toks:
+      self.model(Tensor([[tok]]), start_pos, temperature).realize()
+      start_pos += 1
+    return start_pos
+
+  @property
+  def tokenizer_name(self) -> str: pass
+  def chat_template(self, chat_template: bool | str = False) -> str: pass
+  def apply_chat_template(self, chat_history: list[dict[str, str]], add_generation_prompt: bool = True) -> str:
+    ret = ""
+    for message in chat_history:
+      ret += f"<|start_header_id|>{message['role']}<|end_header_id|>\n\n{message['content'].strip()}<|eot_id|>"
+    if add_generation_prompt: ret += "<|start_header_id|>assistant<|end_header_id|>\n\n"
+    return ret
+
+  def generate_until(self, requests: list[Instance]) -> list[str]:
+    continuations = []
+    for request in tqdm(requests):
+      prompt, args = request.args
+      until = [self.tokenizer.encode(tok) for tok in args.get("until", [])]
+      toks = [self.tokenizer.bos_id] + self.tokenizer.encode(prompt,allow_special=True)
+      prompt_len = len(toks)
+      max_gen_toks = args.get("max_gen_toks") or args.get("max_length") or self.max_length-prompt_len
+      assert self.max_length >= max_gen_toks, "This eval needs a longer context length"
+      temperature = args.get("temperature", 0.0)
+      start_pos = self._prefill(toks[:-1], temperature)
+
+      for _ in range(max_gen_toks):
+        next_tok = self.model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
+        if next_tok in self.tokenizer.stop_tokens or next_tok in until: break
+        toks.append(next_tok)
+        start_pos += 1
+
+      continuations.append(self.tokenizer.decode(toks[prompt_len:]))
+    return continuations
+
+  def loglikelihood(self, requests: list[Instance]) -> list[tuple[float, bool]]: raise NotImplementedError() # needs changes to extra/models/llama.py
+  def loglikelihood_rolling(self, requests: list[Instance]) -> list[tuple[float, bool]]: raise NotImplementedError()
+
+if __name__ == '__main__':
+  print(f"using {Device.DEFAULT} backend")
+
+  parser = argparse.ArgumentParser(description='Run LLaMA evals in tinygrad', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  parser.add_argument('--size', type=str, default="8B", help=f"Size of model to use [{', '.join(list(MODEL_PARAMS.keys()))}]")
+  parser.add_argument('--chat', action='store_true', help="Use chat model")
+  parser.add_argument('--ctx', type=int, default=8192, help="Max context length")
+  parser.add_argument('--quantize', type=str, default=None, help="Quantize the weights to int8 or int4 in memory")
+  parser.add_argument('--eval', type=str, default="mgsm_en_cot_sglang", help="Run in evaluation mode")
+  parser.add_argument('--limit', type=int, default=None, help="Limit tests in eval")
+  parser.add_argument('--num_fewshot', type=int, default=None, help="Number of examples to add to context")
+  parser.add_argument('--model', type=Path, default="./weights/LLaMa/", help="Location of the weights")
+  parser.add_argument('--output_path', type=Path, default=None, help="Location of the log file")
+  args = parser.parse_args()
+
+  # run eval and exit
+  adaptor = LLaMaAdaptor(model_size=args.size, quantize=args.quantize,
+                         checkpoint_path=args.model, max_length=args.ctx)
+  task_manager = TaskManager(include_path="./")
+  results = simple_evaluate(model=adaptor, tasks=args.eval.split(","), task_manager=task_manager, apply_chat_template=args.chat,
+                            num_fewshot=args.num_fewshot, limit=args.limit)
+
+  if args.output_path: args.output_path.write_text(json.dumps(results, indent=2))
+  for task_name, val in results["results"].items():
+    print(f"{task_name}:")
+    print("\n".join(f"\t{k}: {v}" for k, v in val.items() if k != "alias"))