IQ.Pilot Release Commit @ 0798119

tinygrad_repo/test/backend/test_arange.py

@@ -0,0 +1,246 @@
+import unittest
+import numpy as np
+from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
+from tinygrad.helpers import Context, getenv, DEV
+from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
+from tinygrad.renderer.ptx import PTXRenderer
+from test.helpers import needs_second_gpu
+
+class TestArange(unittest.TestCase):
+  def _get_flops(self, tensor, desired):
+    GlobalCounters.reset()
+    linear = compile_linear(tensor.schedule_linear())
+    self.assertEqual(len(linear.src), 1)
+    run_linear(linear)
+    np.testing.assert_equal(tensor.numpy(), desired)
+    return estimate_uop(linear.src[-1]).ops
+
+  def test_arange_complexity(self):
+    self.assertEqual(self._get_flops(Tensor.arange(256), np.arange(256)), 0)
+    self.assertEqual(self._get_flops(Tensor.arange(2560), np.arange(2560)), 0)
+
+  @unittest.skipIf(Device.DEFAULT == "CL", "flaky in CI")
+  def test_arange_cumsum(self):
+    np.testing.assert_equal(Tensor.arange(513).cumsum(0).numpy(), np.arange(513).cumsum())
+
+  def test_arange_cat(self):
+    t = Tensor.arange(2, dtype=dtypes.int)+Tensor([3])
+    self.assertEqual(t.cat(t).tolist(), [3, 4, 3, 4])
+
+  def test_eye_complexity(self):
+    with Context(NOOPT=1):
+      # NOTE: not every backend supports CMPEQ
+      self.assertLessEqual(self._get_flops(Tensor.eye(2560).contiguous(), np.eye(2560)), 2*2560*2560)
+
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexing is weird")
+  def test_tri_complexity(self):
+    with Context(NOOPT=1):
+      t = Tensor.ones(256, 256).contiguous().realize()
+      linear = compile_linear(t.triu().schedule_linear())
+      self.assertLessEqual(estimate_uop(linear.src[-1]).ops, 4 * 256 * 256)
+
+DSET, DDIM = 2048, 32
+
+class TestIndexing(unittest.TestCase):
+  def test_arange_2_reduce(self):
+    needle = Tensor.zeros(16384, dtype=dtypes.int).contiguous()
+    needle[1337] = 1
+    needle.realize()
+    with Context(NOOPT=1):
+      GlobalCounters.reset()
+      out = ((Tensor.arange(1,16385)-1)*needle).sum()
+      linear, var_vals = out.linear_with_vars()
+      self.assertEqual(len(linear.src), 1)
+      run_linear(linear, var_vals)
+    self.assertEqual(out.item(), 1337)
+
+  def test_manual_index(self):
+    dataset = Tensor.rand(DSET, DDIM).realize()
+    idxs = Tensor([0,3,5,6]).realize()
+    real_index = dataset.numpy()[idxs.numpy()]
+    print("*** indexing ***")
+    with Context(NOOPT=1):
+      GlobalCounters.reset()
+      rng = Tensor.arange(DSET, dtype=dtypes.int).reshape(1, 1, DSET, 1).expand(4, DDIM, DSET, 1)
+      idxs = idxs.reshape(4,1,1,1).expand(4, DDIM, DSET, 1)
+      reshape_dataset = dataset.T.reshape(1, DDIM, DSET, 1).expand(4, DDIM, DSET, 1)
+      full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1, buffer=False))
+      X = full.sum(axis=(2,3))
+      linear, var_vals = X.linear_with_vars()
+      self.assertEqual(len(linear.src), 1)
+      run_linear(linear, var_vals)
+      assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
+    np.testing.assert_allclose(real_index, X.numpy())
+
+  def test_index_variable(self):
+    dataset = Tensor.rand(DSET, DDIM).realize()
+    v = Variable("v", 0, DDIM-1)
+    with Context(NOOPT=1):
+      GlobalCounters.reset()
+      vb = Tensor(v.bind(12))
+      comp = dataset[vb].numpy()
+      # no global ops because they are all indexing
+      self.assertEqual(GlobalCounters.global_ops, 0)
+    np.testing.assert_allclose(comp, dataset.numpy()[12])
+
+  def test_index(self):
+    dataset = Tensor.rand(DSET, DDIM).realize()
+    idxs = Tensor([0,3,5,6]).realize()
+    real_index = dataset.numpy()[idxs.numpy()]
+    print("*** indexing ***")
+    with Context(NOOPT=1):
+      GlobalCounters.reset()
+      X = dataset[idxs]
+      assert X.shape == (4,DDIM)
+      linear, var_vals = X.linear_with_vars()
+      self.assertEqual(len(linear.src), 1)
+      run_linear(linear, var_vals)
+      assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
+    np.testing.assert_allclose(real_index, X.numpy())
+
+  def test_index_fused(self, noopt=1):
+    dataset = Tensor.rand(DSET, DDIM).realize()
+    idxs = Tensor([0,3,5,6]).realize()
+    real_index = dataset.numpy()[idxs.numpy()]
+    print("*** indexing ***")
+    with Context(NOOPT=noopt):
+      GlobalCounters.reset()
+      X = dataset[idxs]
+      assert X.shape == (4,DDIM)
+      linear, var_vals = X.linear_with_vars()
+      self.assertEqual(len(linear.src), 1)
+      run_linear(linear, var_vals)
+      assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops} != {4*DSET}"
+    np.testing.assert_allclose(real_index, X.numpy())
+  @unittest.skip("not ready")
+  def test_index_fused_opt(self): self.test_index_fused(0)
+
+  def test_index_fused_out_of_bounds(self):
+    dataset = Tensor.rand(256, 256).realize()
+    idxs = Tensor([-19238, -257, 256, 495, 10982377]).realize()
+    with Context(NOOPT=1):
+      X = dataset[idxs]
+      np.testing.assert_equal(X.numpy(), 0)
+
+  def test_index_mnist(self, noopt=1, op_limit=512*784*13, split_reduceop=0):
+    # WEBGPU generates more ops due to bitpacking of < 4-byte dtypes
+    if Device.DEFAULT == "WEBGPU": op_limit *= 15
+    # from tinygrad.nn.datasets import mnist
+    X_train, Y_train = Tensor.randint(DSET, 1, 28, 28, dtype='uchar').realize(), Tensor.randint(DSET, dtype='uchar').realize()
+    with Context(NOOPT=noopt, SPLIT_REDUCEOP=split_reduceop):
+      samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0]).realize()
+      GlobalCounters.reset()
+      x = X_train[samples].numpy()
+      y = Y_train[samples].numpy()
+      assert GlobalCounters.global_ops < op_limit, f"too many ops {GlobalCounters.global_ops} != {op_limit}"
+    np.testing.assert_allclose(X_train.numpy()[samples.numpy()], x)
+    np.testing.assert_allclose(Y_train.numpy()[samples.numpy()], y)
+
+  def test_index_mnist_opt(self): self.test_index_mnist(0)
+  def test_index_mnist_split(self): self.test_index_mnist(1, split_reduceop=1)
+  def test_index_mnist_opt_split(self): self.test_index_mnist(0, split_reduceop=1)
+
+  def test_llama_embedding(self, noopt=1, op_limit=65536):
+    # llama3 is 128256
+    vocab_size, embed_size = (10, 3)
+    emb = nn.Embedding(vocab_size, embed_size)
+    emb_w = emb.weight.numpy()
+    x = Tensor([1,2,3,4])
+    with Context(NOOPT=noopt):
+      GlobalCounters.reset()
+      z = emb(x).realize()
+      self.assertLessEqual(GlobalCounters.global_ops, op_limit)
+      self.assertEqual(GlobalCounters.kernel_count, 2)
+    if getenv("CHECK", 1):
+      import torch
+      with torch.no_grad():
+        torch_emb = torch.nn.Embedding(vocab_size, embed_size).eval()
+        torch_emb.weight[:] = torch.tensor(emb_w, dtype=torch.float32)
+      torch_z = torch_emb(torch.tensor(x.numpy()))
+      # TODO: reshape to match torch, should we do this in nn?
+      np.testing.assert_allclose(z.numpy().reshape(4, embed_size), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
+  # at least the arange is being fused
+  def test_llama_embedding_opt(self): self.test_llama_embedding(0, 1_736_704_000)
+
+  # NOTE: call doesn't work with SPEC=2
+  @unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
+  @Context(USE_ATOMICS=1, SPEC=1)
+  def test_llama_8b_embedding_backward(self):
+    from tinygrad.renderer.cstyle import CStyleLanguage
+    if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
+    vocab_size, embed_size = 1000, 128
+    bs, seqlen = 4, 256
+    idx = Tensor.randint(bs, seqlen, high=vocab_size)
+    emb = nn.Embedding(vocab_size, embed_size)
+    emb.weight = Tensor.ones(vocab_size, embed_size)
+    gt = Tensor.zeros(bs, seqlen, embed_size)
+    Tensor.realize(idx, emb.weight, gt)
+    GlobalCounters.reset()
+    loss = (emb(idx)-gt).square().sum()
+    loss.backward()
+    emb.weight.grad.realize()
+    bwd_ops = GlobalCounters.global_ops
+    print(f"embedding bwd: {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops")
+    self.assertLess(bwd_ops, bs*seqlen*embed_size*20, f"backward ops {bwd_ops:,} should be less than 20 per with atomic scatter-add")
+    # correctness check
+    expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
+    for i in idx.flatten().numpy(): expected_grad[i] += 2
+    np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
+
+  @needs_second_gpu
+  @unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
+  @Context(USE_ATOMICS=1, SPEC=1)
+  def test_embedding_backward_vocab_sharded(self):
+    from tinygrad.renderer.cstyle import CStyleLanguage
+    if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
+    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
+    vocab_size, embed_size = 1000, 128
+    bs, seqlen = 4, 256
+    idx = Tensor.randint(bs, seqlen, high=vocab_size)
+    emb = nn.Embedding(vocab_size, embed_size)
+    emb.weight = Tensor.ones(vocab_size, embed_size)
+    gt = Tensor.zeros(bs, seqlen, embed_size)
+    Tensor.realize(idx, emb.weight, gt)
+    # compute expected grad on single device
+    expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
+    for i in idx.flatten().numpy(): expected_grad[i] += 2
+    # now shard the embedding weight on vocab axis and recompute
+    emb.weight = Tensor.ones(vocab_size, embed_size)
+    emb.weight.shard_(devices, axis=0)
+    idx = idx.shard(devices, axis=None)
+    gt = gt.shard(devices, axis=None)
+    Tensor.realize(idx, emb.weight, gt)
+    loss = (emb(idx)-gt).square().sum()
+    loss.backward()
+    np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
+
+  @unittest.skipUnless(Device.DEFAULT == "AMD" or (Device.DEFAULT == "NULL" and DEV.arch.startswith("gfx")), "tests AMD bf16 cast overhead")
+  def base_test_llama_8b_rope_backward(self, dtype, ops_scale=1):
+    from extra.models.llama import precompute_freqs_cis, apply_rotary_emb
+    bs, seqlen, dim, n_heads = 1, 512, 256, 4
+    head_dim = dim // n_heads
+    x = Tensor.randn(bs, seqlen, dim, dtype=dtype)
+    wq = Tensor.randn(dim, dim, dtype=dtype)
+    freqs_cis = precompute_freqs_cis(head_dim, seqlen).cast(dtype)
+    Tensor.realize(x, wq, freqs_cis)
+    xq = (x @ wq.T)
+    # main llama does not fuse it
+    #xq = xq.contiguous_backward()
+    xq = xq.reshape(bs, seqlen, n_heads, head_dim)
+    xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis)
+    xq_rope.sum().backward()
+    linear = compile_linear(wq.grad.schedule_linear())
+    assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}"
+    bwd_ops = estimate_uop(linear.src[0]).ops
+    expected_ops = bs*seqlen*dim*dim*ops_scale
+    print(f"rope matmul bwd ({dtype}): {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops")
+    self.assertLess(bwd_ops, expected_ops, f"rope bwd ops {bwd_ops:,} should be < {ops_scale} per (got {bwd_ops/(bs*seqlen*dim*dim):.1f})")
+
+  def test_llama_8b_rope_backward_f16(self):
+    self.base_test_llama_8b_rope_backward(dtypes.float16, ops_scale=2)
+  # bfloat16 on non CDNA4 has ~10x ops overhead because of the software emulation
+  def test_llama_8b_rope_backward_bf16(self):
+    self.base_test_llama_8b_rope_backward(dtypes.bfloat16, ops_scale=2 if Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950") else 25)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/backend/test_asm_gemm.py

@@ -0,0 +1,252 @@
+import unittest
+from tinygrad import Tensor, Device, dtypes, Context
+from tinygrad.helpers import getenv, system, DEV
+from extra.gemm.cdna_asm_gemm import asm_gemm
+from test.helpers import needs_second_gpu
+from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX
+
+# On non CDNA4 it will only validate the Tensor.custom_kernel integration
+# Use DEV=NULL:HIP:gfx950 to also test the assembly
+def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950")
+
+def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=None, gpus:int=1) -> None:
+  Tensor.manual_seed(0)
+  input_dtype = dtypes.bfloat16 if dtype == FP8_DTYPE else dtype
+  a_rand = Tensor.randn(a_shape, dtype=dtypes.float).sub(0.5).cast(input_dtype)
+  b_rand = Tensor.randn(b_shape, dtype=dtypes.float).sub(0.5).cast(input_dtype)
+  with Context(DEBUG=0):
+    Tensor.realize(a_rand, b_rand)
+
+  devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(gpus)) if (multi:=gpus>1) else None
+
+  if dtype == FP8_DTYPE:
+    a_rand, x_scale, _ = quantize_fp8(a_rand)
+    b_rand, w_scale, _ = quantize_fp8(b_rand)
+    grad_amax_state = Tensor.full((), FP8_MAX, dtype=dtypes.float32, device=devs).contiguous()
+    with Context(DEBUG=0):
+      Tensor.realize(a_rand, x_scale, b_rand, w_scale, grad_amax_state)
+
+  # clone all inputs before any backward: a clone copies the source's current .grad
+  a, b = a_rand.clone(), b_rand.clone()
+  if dtype == FP8_DTYPE:
+    a_ref, b_ref = a_rand.detach().cast(dtypes.bfloat16), b_rand.detach().cast(dtypes.bfloat16)
+  else:
+    a_ref, b_ref = a_rand.clone(), b_rand.clone()
+  if multi: a, b = a.shard(devs, axis=a_shard), b.shard(devs, axis=b_shard)
+  if dtype == FP8_DTYPE:
+    tst = asm_gemm(a, b, x_scale=x_scale, w_scale=w_scale, grad_amax_state=grad_amax_state)
+  else:
+    tst = asm_gemm(a, b)
+  tst.sum().backward()
+  Tensor.realize(tst, a.grad, b.grad)
+
+  if multi: a_ref, b_ref = a_ref.shard(devs, axis=a_shard), b_ref.shard(devs, axis=b_shard)
+  if dtype == FP8_DTYPE:
+    ref = ((a_ref @ b_ref) * x_scale * w_scale).cast(dtypes.bfloat16)
+  else:
+    ref = a_ref @ b_ref
+  ref.sum().backward()
+  Tensor.realize(ref, a_ref.grad, b_ref.grad)
+
+  # no validation on the NULL device
+  if a_rand.device.startswith("NULL"): return None
+  atol, rtol = (2e-1, 1e-2) if dtype == dtypes.bfloat16 else (256, 1e-2) if dtype == FP8_DTYPE else (1e-2, 1e-3)
+  # allow more rtol for multi because of ALLREDUCE_CAST
+  grad_atol, grad_rtol = (16895, 0.125) if dtype == FP8_DTYPE else (atol, 2e-2 if multi else rtol)
+  with Context(DEBUG=0):
+    # enable for debugging, slow for larger gemms
+    if getenv("USE_NPY"):
+      import numpy as np
+      np.testing.assert_allclose(tst.numpy(), ref.numpy(), atol=atol, rtol=rtol)
+      np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), atol=grad_atol, rtol=grad_rtol)
+      np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), atol=grad_atol, rtol=grad_rtol)
+    assert tst.allclose(ref, atol=atol, rtol=rtol).item(), "forward mismatch"
+    assert a.grad.allclose(a_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_a mismatch"
+    assert b.grad.allclose(b_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_b mismatch"
+
+def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
+  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
+
+def verify_asm_gemm_k_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=8) -> None:
+  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=1, b_shard=0, gpus=gpus)
+
+def verify_asm_gemm_n_sharded(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
+
+def verify_asm_gemm_m_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
+
+def verify_asm_gemm_n_sharded_2d(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
+
+def verify_asm_gemm_k_sharded_3d(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=2, b_shard=0, gpus=gpus)
+
+# 128x smaller than usual
+# uses the UOp GEMM, runs on non CDNA4 and CI
+@unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
+class TestGemm(unittest.TestCase):
+  def setUp(self):
+    if is_cdna4(): self.skipTest("shapes are too small for the assembly GEMM")
+  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 32), N, N, dtype=dtypes.half)
+  def test_gemm(self): verify_asm_gemm(1, 64, 32, 112)
+  def test_gemm_batched(self): verify_asm_gemm(2, 64, 32, 32)
+  @needs_second_gpu
+  def test_gemm_multi(self): verify_asm_gemm(2, 64, 32, 32, gpus=2)
+  @needs_second_gpu
+  def test_gemm_k_sharded(self): verify_asm_gemm_k_sharded(64, 64, 2*64, gpus=2)
+  @needs_second_gpu
+  def test_gemm_m_sharded(self): verify_asm_gemm_m_sharded(2*64, 64, 32, gpus=2)
+  @needs_second_gpu
+  def test_gemm_n_sharded(self): verify_asm_gemm_n_sharded(1, 64, 64, 32, gpus=2)
+  @needs_second_gpu
+  def test_gemm_n_sharded_2d(self): verify_asm_gemm_n_sharded_2d(64, 2*64, 32, gpus=2)
+  @needs_second_gpu
+  def test_gemm_k_sharded_3d(self): verify_asm_gemm_k_sharded_3d(1, 64, 32, 2*64, gpus=2)
+
+# uses the smallest size for the cdna assembly gemm
+class TestAsmGEMM(unittest.TestCase):
+  def setUp(self):
+    if not is_cdna4():
+      self.skipTest("assembly gemm is only for cdna4")
+
+  def test_tiny(self): verify_asm_gemm(1, 256, 256, 64)
+
+  def test_verify_with_numpy(self):
+    import numpy as np
+    M, N, K = 256, 256, 64
+    rng = np.random.default_rng(0)
+    a_np = (rng.random((M, K), dtype=np.float32) - 0.5).astype(np.half)
+    b_np = (rng.random((K, N), dtype=np.float32) - 0.5).astype(np.half)
+    c_np = a_np @ b_np
+    a, b = Tensor(a_np), Tensor(b_np)
+    c = asm_gemm(a, b)
+    c.realize()
+    # no validation on the NULL device
+    if a.device.startswith("NULL"): return None
+    np.testing.assert_allclose(c.numpy(), c_np, atol=2e-3, rtol=5e-2)
+
+  def test_unsupported_batch(self):
+    with self.assertRaisesRegex(AssertionError, "batch size"):
+      verify_asm_gemm(3, 256, 256, 256)
+
+  def test_unsupported_k(self):
+    with self.assertRaisesRegex(AssertionError, "not a multiple"):
+      verify_asm_gemm(1, 1024, 1024, 100)
+  def test_unsupported_m(self):
+    with self.assertRaisesRegex(AssertionError, "not a multiple"):
+      verify_asm_gemm(1, 1000, 256, 256)
+  def test_unsupported_n(self):
+    with self.assertRaisesRegex(AssertionError, "not a multiple"):
+      verify_asm_gemm(1, 256, 1000, 256)
+
+# test the Asm GEMM with Llama shapes, only run on the real machine for speed
+class TestGemmLlama(unittest.TestCase):
+  dtype = dtypes.bfloat16
+
+  def setUp(self):
+    if not is_cdna4() or DEV.interface.startswith("MOCK"):
+      self.skipTest("very slow on non mi350x")
+
+  def test_empty(self): asm_gemm(Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype), Tensor.empty(N, N, dtype=self.dtype)).realize()
+
+  def test_empty_bw(self):
+    x = Tensor.empty(1, N:=getenv("N", 4096), N, dtype=self.dtype)
+    y = Tensor.empty((N, N), dtype=self.dtype)
+    if self.dtype == FP8_DTYPE:
+      x_scale = Tensor.empty((), dtype=dtypes.float32)
+      w_scale = Tensor.empty((), dtype=dtypes.float32)
+      grad_amax_state = Tensor.empty((), dtype=dtypes.float32).contiguous()
+      z = asm_gemm(x, y, x_scale=x_scale, w_scale=w_scale, grad_amax_state=grad_amax_state)
+    else:
+      z = asm_gemm(x, y)
+    z.sum().backward()
+    Tensor.realize(z, x.grad, y.grad)
+    # FP8 GEMM stores bf16 output and its backward produces bf16 gradients.
+    grad_dtype = dtypes.bfloat16 if self.dtype == FP8_DTYPE else self.dtype
+    assert z.dtype == dtypes.bfloat16
+    assert x.grad.dtype == y.grad.dtype == grad_dtype
+
+  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=self.dtype)
+  def test_gemm(self): verify_asm_gemm(1, 8192, 4096, 14336, dtype=self.dtype)
+  def test_gemm_batched(self): verify_asm_gemm(2, 8192, 4096, 4096, dtype=self.dtype)
+
+  def test_gemm1(self): verify_asm_gemm(8, 8192, 4096, 14336, dtype=self.dtype, gpus=8)
+  @unittest.skip("disabled, asm in this shape is slower than tinygrad")
+  def test_gemm2(self): verify_asm_gemm(8, 8192, 128256, 4096, dtype=self.dtype, gpus=8)
+  def test_gemm3(self): verify_asm_gemm(8, 8192, 14336, 4096, dtype=self.dtype, gpus=8)
+  def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, dtype=self.dtype, gpus=8)
+  def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, dtype=self.dtype, gpus=8)
+  def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, dtype=self.dtype, gpus=8)
+  @unittest.skip("disabled, asm in this shape is slower than tinygrad")
+  def test_gemm7(self): verify_asm_gemm(1, 8192, 128256, 4096, dtype=self.dtype)
+  def test_gemm8(self): verify_asm_gemm(1, 4096, 14336, 8192, dtype=self.dtype)
+  def test_gemm9(self): verify_asm_gemm(8, 4096, 14336, 8192, dtype=self.dtype, gpus=8)
+  def test_gemm10(self): verify_asm_gemm(1, 4096, 8192, 4096, dtype=self.dtype)
+  def test_gemm_previously_unsupported(self): verify_asm_gemm(8, 1024, 1024, 4096, gpus=8)
+  def test_k_sharded_1(self): verify_asm_gemm_k_sharded(14336, 4096, 8*8192, dtype=self.dtype, gpus=8)
+  def test_k_sharded_2(self): verify_asm_gemm_k_sharded(4096, 14336, 8*8192, dtype=self.dtype, gpus=8)
+  def test_k_sharded_3(self): verify_asm_gemm_k_sharded(4096, 4096, 8*8192, dtype=self.dtype, gpus=8)
+
+  # M-sharded 2D
+  def test_m_sharded_1(self): verify_asm_gemm_m_sharded(8*8192, 4096, 4096, dtype=self.dtype, gpus=8)
+  def test_m_sharded_2(self): verify_asm_gemm_m_sharded(8*4096, 14336, 4096, dtype=self.dtype, gpus=8)
+
+  # N-sharded 2D
+  def test_n_sharded_2d_1(self): verify_asm_gemm_n_sharded_2d(8192, 8*4096, 4096, dtype=self.dtype, gpus=8)
+  def test_n_sharded_2d_2(self): verify_asm_gemm_n_sharded_2d(4096, 8*14336, 4096, dtype=self.dtype, gpus=8)
+
+  # tensor parallel shapes (Llama 8B, MP=8)
+  def test_tp_n_sharded_wq(self): verify_asm_gemm_n_sharded(1, 8192, 4096, 4096, dtype=self.dtype, gpus=8)
+  def test_tp_n_sharded_w1(self): verify_asm_gemm_n_sharded(1, 8192, 14336, 4096, dtype=self.dtype, gpus=8)
+  def test_tp_k_sharded_wo(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 4096, dtype=self.dtype, gpus=8)
+  def test_tp_k_sharded_w2(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 14336, dtype=self.dtype, gpus=8)
+
+  # more shapes: vary M, N, K independently
+  def test_shape_small_square(self): verify_asm_gemm(1, 256, 256, 256)
+  def test_shape_small_rect_m(self): verify_asm_gemm(1, 512, 256, 256)
+  def test_shape_small_rect_n(self): verify_asm_gemm(1, 256, 512, 256)
+  def test_shape_small_rect_k(self): verify_asm_gemm(1, 256, 256, 512)
+  def test_shape_tall(self): verify_asm_gemm(1, 2048, 256, 256)
+  def test_shape_wide(self): verify_asm_gemm(1, 256, 2048, 256)
+  def test_shape_deep(self): verify_asm_gemm(1, 256, 256, 4096)
+  def test_shape_non_square(self): verify_asm_gemm(1, 1024, 2048, 512)
+  def test_shape_batched_small(self): verify_asm_gemm(2, 256, 256, 256)
+  def test_shape_batched_rect(self): verify_asm_gemm(2, 512, 1024, 256)
+  # K edge cases: iters=1,2,3 exercise different loop paths
+  def test_shape_k64(self): verify_asm_gemm(1, 256, 256, 64)
+  def test_shape_k128(self): verify_asm_gemm(1, 256, 256, 128)
+  def test_shape_k192(self): verify_asm_gemm(1, 256, 256, 192)
+
+  def test_llama3_out1(self): verify_asm_gemm(1, 8192, 128256, 4096, dtype=self.dtype)
+  def test_llama3_out2(self): verify_asm_gemm(1, 8192, 4096, 128256, dtype=self.dtype)
+  def test_llama3_out3(self): verify_asm_gemm(1, 4096, 128256, 8192, dtype=self.dtype)
+
+def has_hipcc():
+  try: system("hipcc --version")
+  except Exception: return False
+  return True
+
+@unittest.skipUnless(has_hipcc(), "FP8 gemm requires hipcc to compile")
+class TestGemmLlamaFP8(TestGemmLlama): dtype = FP8_DTYPE
+
+class TestMagicGu(unittest.TestCase):
+  def test_magicgu_matches_old(self):
+    from extra.gemm.cdna_asm_gemm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
+    old_iters_args = {64: (67108864, 0), 128: (33554432, 0), 224: (613566757, 2147483656)}
+    old_gemm_shapes = [
+      (8192, 4096, 4096), (8192, 14336, 4096), (8192, 4096, 14336),
+      (8192, 8192, 8192), (4096, 4096, 4096), (4096, 14336, 4096),
+      (4096, 14336, 8192), (4096, 4096, 14336), (14336, 4096, 8192),
+      (4096, 8192, 14336), (4096, 4096, 8192), (4096, 8192, 4096),
+    ]
+    for M, N, K in old_gemm_shapes:
+      iters = K // TILE_K
+      total = (M // TILE_M) * (N // TILE_N) * iters
+      for batch in [1, 2]:
+        magic, shift = _magicgu_mulhi(iters, total * batch)
+        old_magic, old_shift = old_iters_args[iters]
+        self.assertEqual((magic, shift), (old_magic, old_shift), f"mismatch for ({M},{N},{K}) batch={batch} iters={iters}")
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/backend/test_const_folding.py

@@ -0,0 +1,209 @@
+import unittest, math
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.dtype import DTYPES_DICT
+from tinygrad.uop.ops import Ops, UOp
+import numpy as np
+from test.helpers import not_support_multi_device
+
+def _check_ast_count(desired_count:int, t:Tensor):
+  # NOTE: this has side effect because everything can be scheduled only once
+  schedule = t.schedule_linear()
+  asts = [s for s in schedule.src if s.src[0].op is Ops.SINK]
+  len(asts)
+  # NOT SUPPORTED ANYMORE
+  #assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
+
+class TestMovedConstFolding(unittest.TestCase):
+  def test_contiguous_deviceless_const(self):
+    t = Tensor(UOp.const(dtypes.float, 2.0)).contiguous()
+    self.assertIs(t.uop.op, Ops.CONST)
+    self.assertIsNone(t.uop.device)
+
+  def test_add_shrunk_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(6).shrink(((1, 5),)))
+
+  def test_add_padded_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(2).pad(((1, 1),)))
+
+  def test_mul_shrunk_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(6).shrink(((1, 5),)))
+
+  def test_add_padded_one(self):
+    _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
+
+  def test_copy_padded_const(self):
+    schedule = Tensor.ones(4, device="CPU:0", buffer=False).pad(((1, 1),)).to("CPU:1").schedule_linear()
+    assert not any(si.src[0].op is Ops.COPY for si in schedule.src), "const copy should be folded"
+    np.testing.assert_equal(Tensor.ones(4, device="CPU:0", buffer=False).pad(((1, 1),)).to("CPU:1").numpy(), [0, 1, 1, 1, 1, 0])
+
+  def test_cast_padded(self):
+    # NOTE: it's always 1 kernel when calling .numpy, limitation of _check_ast_count
+    _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
+    np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])
+    _check_ast_count(1, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16))
+    np.testing.assert_equal(Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16).numpy(), [0, 65535, 65535, 65535, 65535, 0])
+    # folded
+    _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64))
+    np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64).numpy(), [0, 1, 1, 1, 1, 0])
+
+class TestReduceOpsConstFolding(unittest.TestCase):
+  def test_const_sum(self):
+    _check_ast_count(0, Tensor.ones(4, 5, 6).sum())
+    np.testing.assert_equal(Tensor.ones(4, 5, 6).sum().numpy(), 4 * 5 * 6)
+    _check_ast_count(0, Tensor.ones(4, 5, 6).sum(axis=0))
+    np.testing.assert_equal(Tensor.ones(4, 5, 6).sum(axis=0).numpy(), np.full((5, 6), 4))
+    _check_ast_count(0, Tensor(4).sum())
+    np.testing.assert_equal(Tensor(4).sum().numpy(), 4)
+
+  def test_padded_const_sum(self):
+    _check_ast_count(0, Tensor.ones(4).pad(((1, 1),)).sum())
+    np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).sum().numpy(), 4)
+
+    # NOTE: cannot just count the non-padded area because some Ops f do not have f(0) = 0.
+    _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).exp().sum())
+    np.testing.assert_allclose(Tensor.ones(4).pad(((1, 1),)).exp().sum().numpy(), 4 * math.e + 2)
+
+  def test_bool_zero_max(self):
+    _check_ast_count(0, Tensor.full((1, 2), True).shrink(((0, 1), (0, 0))).max((1, 0)))
+    np.testing.assert_equal(Tensor.full((1, 2), True).shrink(((0, 1), (0, 0))).max((1, 0)).numpy(), False)
+
+  def test_zero_size_ops(self):
+    for reduceop in [lambda x:x.prod(), lambda x:x.sum()]: # lambda x:x.max() NOTE: numpy gives "reduction operation maximum which has no identity"
+      _check_ast_count(0, reduceop(Tensor.empty(1, 0)))
+      np.testing.assert_equal(reduceop(Tensor.empty(shape:=(1, 0))).numpy(), reduceop(np.empty(shape)))
+
+  def test_zero_size_ops_view(self):
+    for reduceop in [lambda x:x.prod(), lambda x:x.sum()]:
+      _check_ast_count(0, reduceop(Tensor.empty(1, 0, 4).permute((1, 2, 0)).contiguous()))
+      np.testing.assert_equal(reduceop(Tensor.empty(shape:=(1, 0))).numpy(), reduceop(np.empty((shape))))
+
+  def test_zero_size_ops_realized(self):
+    for reduceop in [lambda x:x.prod(), lambda x:x.sum()]:
+      _check_ast_count(0, reduceop((Tensor.randn(0, 1)+1).realize()))
+      np.testing.assert_equal(reduceop((Tensor.randn(shape:=(0, 1))+1).realize()).numpy(), reduceop(np.empty(shape)))
+
+  def test_zero_size_realize_folded(self):
+    # non contiguous folded output doesn't realize
+    _check_ast_count(0, Tensor.empty(1, 0).sum())
+    # contiguous folded const can still schedule
+    a = Tensor.empty(1, 0).sum().contiguous()
+    _check_ast_count(2, a+2)
+    self.assertIs(a.uop.base.op, Ops.BUFFER)
+    np.testing.assert_equal((Tensor.empty(1, 0).sum().contiguous()+2).numpy(), 2)
+    # otherwise we just fuse it
+    _check_ast_count(1, (Tensor.empty(1, 0).sum()+2).contiguous())
+    np.testing.assert_equal((Tensor.empty(1, 0).sum()+2).numpy(), 2)
+
+  def test_const_prod(self):
+    _check_ast_count(0, Tensor.full((2, 3), fill_value=2).prod())
+    np.testing.assert_equal(Tensor.full((2, 3), fill_value=2).prod().numpy(), 2**(2*3))
+    _check_ast_count(0, Tensor.full((4, 5, 6), fill_value=2).prod(axis=0))
+    np.testing.assert_equal(Tensor.full((4, 5, 6), fill_value=2).prod(axis=0).numpy(), np.full((5, 6), 2**4))
+    _check_ast_count(0, Tensor(4).prod())
+    np.testing.assert_equal(Tensor(4).prod().numpy(), 4)
+
+  def test_const_max(self):
+    _check_ast_count(0, Tensor.ones(4, 5, 6).max())
+    np.testing.assert_equal(Tensor.ones(4, 5, 6).max().numpy(), 1)
+    _check_ast_count(0, Tensor(4).max())
+    np.testing.assert_equal(Tensor(4).max().numpy(), 4)
+
+  def test_sum_output_dtype(self):
+    # sum output dtype can be different from input
+    for dt in DTYPES_DICT.values():
+      if dt in Device[Device.DEFAULT].renderer.supported_dtypes():
+        t = Tensor.ones(16, dtype=dt).reshape(4, 4)
+        assert t.sum().dtype == t.contiguous().sum().dtype
+
+@unittest.skipIf(not_support_multi_device() or True, "no multi, RANGEIFY doesn't support multi const folding")
+class TestMultiConstFolding(unittest.TestCase):
+  def test_multi_const_folding_literal(self):
+    ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
+    t = Tensor.arange(16).float().clone().to(ds).realize()
+
+    # non const folding case creates one ast on each shard
+    _check_ast_count(4, t + 1)
+    _check_ast_count(4, 1 + t)
+    _check_ast_count(4, t * 2)
+    _check_ast_count(4, 2 * t)
+
+    # const folded
+    _check_ast_count(0, t + 0)
+    _check_ast_count(0, 0 + t)
+    _check_ast_count(0, t * 0)
+    _check_ast_count(0, 0 * t)
+    _check_ast_count(0, t * 1)
+    _check_ast_count(0, 1 * t)
+    np.testing.assert_equal((t + 0).numpy(), np.arange(16))
+    np.testing.assert_equal((t * 0).numpy(), [0] * 16)
+    np.testing.assert_equal((t * 1).numpy(), np.arange(16))
+
+    _check_ast_count(0, t ** 0)
+    _check_ast_count(0, t ** 1)
+    _check_ast_count(0, 1 ** t)
+
+  def test_multi_const_folding_tensor(self):
+    ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
+    t = Tensor.arange(16).float().clone().to(ds).realize()
+    zero = Tensor.zeros(16).to(ds).realize()
+    one = Tensor.ones(16).to(ds).realize()
+
+    # const folded
+    _check_ast_count(0, t + zero)
+    _check_ast_count(0, zero + t)
+    _check_ast_count(0, t * zero)
+    _check_ast_count(0, zero * t)
+    _check_ast_count(0, t * one)
+    _check_ast_count(0, one * t)
+    np.testing.assert_equal((t + zero).numpy(), np.arange(16))
+    np.testing.assert_equal((t * zero).numpy(), [0] * 16)
+    np.testing.assert_equal((t * one).numpy(), np.arange(16))
+    _check_ast_count(0, t ** zero)
+    _check_ast_count(0, t ** one)
+    _check_ast_count(0, one ** t)
+    np.testing.assert_equal((t ** zero).numpy(), [1] * 16)
+    np.testing.assert_equal((t ** one).numpy(), np.arange(16))
+    np.testing.assert_equal((one ** t).numpy(), [1] * 16)
+
+class TestThreefryConstFolding(unittest.TestCase):
+  def test_threefry(self):
+    x = UOp.const(dtypes.uint64, 5, Device.DEFAULT, ()).threefry(UOp.const(dtypes.uint64, 10, Device.DEFAULT, ()))
+    self.assertIs(x.simplify().op, Ops.CONST)
+
+class TestTautologicalCompare(unittest.TestCase):
+  # without const folding, these would have triggered -Wtautological-compare in clang
+  def test_lt_false(self):
+    # bool < False is always false
+    np.testing.assert_equal((Tensor([True, False]) < False).numpy(), [False, False])
+
+  def test_true_lt(self):
+    # True < bool is always false
+    np.testing.assert_equal((True < Tensor([True, False])).numpy(), [False, False])
+
+  def test_truth_table(self):
+    np.testing.assert_equal((Tensor(False) < Tensor(False)).numpy(), False)
+    np.testing.assert_equal((Tensor(False) < Tensor(True)).numpy(), True)
+    np.testing.assert_equal((Tensor(True) < Tensor(False)).numpy(), False)
+    np.testing.assert_equal((Tensor(True) < Tensor(True)).numpy(), False)
+
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support NaN comparison correctly")
+  def test_a_eq_a(self):
+    # self eq is always true for int or bool
+    a = Tensor([1, 2, 3])
+    np.testing.assert_equal((a == a).numpy(), [True, True, True])
+
+    # not true for nan
+    a = Tensor([math.nan, 1.0, 2.0])
+    np.testing.assert_equal((a == a).numpy(), [False, True, True])
+
+  def test_a_ne_a(self):
+    # self not eq is always false for int or bool
+    a = Tensor([1, 2, 3])
+    np.testing.assert_equal((a != a).numpy(), [False, False, False])
+
+    # not true for nan
+    a = Tensor([math.nan, 1.0, 2.0])
+    np.testing.assert_equal((a != a).numpy(), [True, False, False])
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_custom_kernel.py

@@ -0,0 +1,446 @@
+import unittest
+from tinygrad import Tensor, UOp, GlobalCounters, Context, Device
+from tinygrad.dtype import AddrSpace, dtypes
+from tinygrad.uop.ops import KernelInfo, AxisType, Ops
+
+# **** kernels ****
+
+def custom_arange_kernel(C:UOp) -> UOp:
+  i = UOp.range(C.shape[0], 0)
+  return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.shape[0]}"))
+
+def custom_eye_kernel(C:UOp) -> UOp:
+  i = UOp.range(C.shape[0], 0)
+  j = UOp.range(C.shape[1], 1)
+  return C[i, j].store((i.eq(j)).cast(C.dtype.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.numel()}"))
+
+def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
+  A,B = A.flatten(), B.flatten()
+  assert B.numel() == A.numel()
+  i = UOp.range(A.numel(), 0)
+  return B[i].store(A[i] + 1).end(i).sink(arg=KernelInfo(name=f"add_one_{A.numel()}"))
+
+def custom_elementwise_add_kernel(C:UOp, A:UOp, B:UOp) -> UOp:
+  C,A,B = C.flatten(), A.flatten(), B.flatten()
+  i = UOp.range(C.numel(), 0)
+  return C[i].store(A[i]+B[i]).end(i).sink(arg=KernelInfo(name=f"custom_add_kernel_{C.numel()}")).simplify()
+
+def custom_elementwise_addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
+  C,D,A,B = C.flatten(), D.flatten(), A.flatten(), B.flatten()
+  assert C.numel() == D.numel()
+  i = UOp.range(C.numel(), 0)
+  store_c = C[i].store(A[i]+B[i])
+  store_d = D[i].store(A[i]*B[i])
+  return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name=f"custom_addmul_kernel_{C.numel()}")).simplify()
+
+def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
+  assert A.shape[1] == B.shape[0]
+  i, j, k = UOp.range(C.shape[0], 0), UOp.range(C.shape[1], 1), UOp.range(A.shape[1], 2, axis_type=AxisType.REDUCE)
+  C = C[i, j].set(0.0)
+  C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
+  prog = C.end(i, j)
+  return prog.sink(arg=KernelInfo(name=f"custom_gemm_{C.shape[0]}_{C.shape[1]}_{A.shape[1]}", opts_to_apply=()))
+
+def custom_sum(B:UOp, A:UOp) -> UOp:
+  i = UOp.range(A.shape[0], 0, axis_type=AxisType.REDUCE)
+  B = B[0].set(0.0)
+  B = B[0].set(B.after(i)[0] + A[i], end=i)
+  return B.sink(arg=KernelInfo(name=f"custom_sum_{A.shape[0]}", opts_to_apply=()))
+
+def flip_contract_kernel(dest:UOp, src:UOp):
+  i = UOp.range(dest.shape[0], 0)
+  j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
+  vec = src[i, j].contract(j)
+  store = UOp.group(*[dest[i, k].store(vec.gep(3-k)) for k in range(4)])
+  return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))
+
+def slice_sum_kernel(dest:UOp, src:UOp):
+  G = UOp.range(src.shape[0], 0)
+  slice_src = src[G, :]
+  reg = UOp.placeholder((1,), dest.dtype.base, 0, addrspace=AddrSpace.REG)
+  reg = reg.after(G)[0].set(0)
+  R = UOp.range(src.shape[1], 1, AxisType.REDUCE)
+  reg = reg[0].set(reg.after(R)[0] + slice_src[R], end=R)
+  ast = dest[G].set(reg[0], end=G)
+  return ast.sink(arg=KernelInfo(name=f"slice_sum_{src.shape[0]}_{src.shape[1]}", opts_to_apply=()))
+
+def simple_qkv_kernel(O:UOp, Q:UOp, K:UOp, V:UOp) -> UOp:
+  # attention without softmax
+  N, d = Q.shape[0], Q.shape[1]
+
+  i = UOp.range(N, 0)  # output row
+  d_out = UOp.range(d, 1)  # output column
+  j = UOp.range(N, 2, axis_type=AxisType.REDUCE)
+
+  k_inner = UOp.range(d, 3, axis_type=AxisType.REDUCE)
+  qk_acc = UOp.placeholder((1,), Q.dtype.base, 0, addrspace=AddrSpace.REG)
+  qk_acc = qk_acc.after(i, j)[0].set(0.0)
+  qk_acc = qk_acc[0].set(qk_acc.after(k_inner)[0] + Q[i, k_inner] * K[j, k_inner], end=k_inner)
+  qk_score = qk_acc[0] / (d ** 0.5)
+
+  out_acc = UOp.placeholder((1,), Q.dtype.base, 1, addrspace=AddrSpace.REG)
+  out_acc = out_acc.after(i, d_out)[0].set(0.0)
+  out_acc = out_acc[0].set(out_acc.after(j)[0] + qk_score * V[j, d_out], end=j)
+
+  store = O[i, d_out].store(out_acc[0])
+  return store.end(d_out).end(i).sink(arg=KernelInfo(name=f"simple_qkv_{N}_{d}", opts_to_apply=()))
+
+# **** backward callbacks ****
+
+def backward_gemm(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
+  out, a, b = kernel.src[1:]
+  grad_a = (Tensor(gradient) @ Tensor(b).T).uop
+  grad_b = (Tensor(a).T @ Tensor(gradient)).uop
+  return (None, grad_a, grad_b)
+
+def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
+  out, a, b = kernel.src[1:]
+  grad_a = Tensor.empty_like(Tensor(a)).custom_kernel(Tensor(gradient), Tensor(b).T, fxn=custom_gemm)[0].uop
+  grad_b = Tensor.empty_like(Tensor(b)).custom_kernel(Tensor(a).T, Tensor(gradient), fxn=custom_gemm)[0].uop
+  return (None, grad_a, grad_b)
+
+# **** tests ****
+
+class TestCustomKernel(unittest.TestCase):
+  def test_empty(self):
+    a = Tensor.empty(1)
+    a = Tensor.custom_kernel(a, fxn=lambda _: UOp.sink(arg=KernelInfo()))[0]
+    a.realize()
+
+  def test_simple(self):
+    a = Tensor.ones(16, 16).contiguous()
+    b = Tensor.ones(16, 16).contiguous()
+    c = Tensor.empty(16, 16)
+
+    c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
+
+    out = c.flatten().tolist()
+    assert all(x == 2 for x in out), "all 2"
+
+  def test_simple_sharded(self):
+    devs = ("CPU:0", "CPU:1")
+
+    a = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
+    b = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
+    # ugly construction to get a sharded empty tensor
+    c = Tensor(Tensor.empty(8, 16, device=devs).uop.multi(0), device=devs)
+    c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
+    out = c.flatten().tolist()
+    assert all(x == 2 for x in out), "all 2"
+
+  def test_sharded_add_one(self):
+    # PYTHON backend explicitly checks for OOB access for wrong multi shape regression
+    devs = ("PYTHON:0", "PYTHON:1")
+    a = Tensor.ones(4, 4).contiguous().shard(devs, axis=0)
+    c = Tensor(Tensor.empty(2, 4, device=devs).uop.multi(0), device=devs)
+    c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
+    assert (c == 2).all().item()
+
+  def test_multioutput(self):
+    a = Tensor.full((16, 16), 3.).contiguous()
+    b = Tensor.full((16, 16), 3.).contiguous()
+    c = Tensor.empty(16, 16)
+    d = Tensor.empty(16, 16)
+
+    c,d = Tensor.custom_kernel(c,d,a,b, fxn=custom_elementwise_addmul_kernel)[:2]
+    Tensor.realize(c,d)
+
+    assert all(x == 6 for x in c.flatten().tolist()), "all 6"
+    assert all(x == 9 for x in d.flatten().tolist()), "all 9"
+
+  def test_arange(self):
+    ref = Tensor.arange(100)
+    tst = Tensor.empty_like(ref)
+    tst = tst.custom_kernel(fxn=custom_arange_kernel)[0]
+    self.assertTrue((ref == tst).all().item())
+
+  def test_eye(self):
+    ref = Tensor.eye(1024).clone().realize()
+    tst = Tensor.empty_like(ref)
+    tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
+    self.assertTrue((ref == tst).all().item())
+
+  @unittest.skip("contract shouldn't be supported here")
+  def test_flip_contract(self):
+    a = Tensor.randn(10,4)
+    b = Tensor.empty_like(a)
+    b = b.custom_kernel(a, fxn=flip_contract_kernel)[0]
+    self.assertTrue((a.flip(1) == b).all().item())
+
+  def test_noncontig(self):
+    a = Tensor.ones(16, 16).contiguous()
+    tst = Tensor.empty_like(a)
+    b = a+1
+    b_p1 = Tensor.custom_kernel(tst, b, fxn=custom_add_one_kernel)[0]
+    self.assertTrue((b_p1 == 3).all().item())
+
+  def test_sum(self):
+    a = Tensor([1.0, 2, 3, 4, 5])
+    tst = Tensor.empty(1)
+    b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
+    self.assertEqual(b.item(), 15)
+
+  def test_sum_int(self):
+    a = Tensor([1, 2, 3, 4, 5])
+    tst = Tensor.empty(1, dtype=a.dtype)
+    b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
+    self.assertEqual(b.item(), 15)
+
+  def test_slice_sum(self):
+    A = Tensor.randn(16, 16).contiguous()
+    B = Tensor.empty(16)
+    B = Tensor.custom_kernel(B, A, fxn=slice_sum_kernel)[0]
+    self.assertTrue(B.allclose(A.sum(1)).item())
+
+  def test_gemm(self):
+    N = 16
+    a = Tensor.randn(N, N)
+    b = Tensor.randn(N, N)
+    c = Tensor.empty(N, N)
+
+    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
+    err = (tst - (a@b)).square().max()
+    self.assertLess(err.item(), 1e-6)
+
+  def test_gemm_multi(self):
+    devs = ("CPU:0", "CPU:1")
+    N = 16
+    a = Tensor.randn(N, N).shard_(devs, axis=0)
+    b = Tensor.randn(N, N).to(devs)
+    c = Tensor(Tensor.empty(N//2, N, device=devs).uop.multi(0), device=devs)
+    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
+    err = (tst - (a@b)).square().max()
+    self.assertLess(err.item(), 1e-6)
+
+  def test_gemm_backward_custom(self): self.test_gemm_backward(True)
+  # NOTE: grad_fxn doesn't work with pyrender
+  def test_gemm_backward(self, custom_backward_gemm=False):
+    N = 4
+    a_rand = Tensor.randn(N, 8)
+    b_rand = Tensor.randn(8, N)
+    Tensor.realize(a_rand, b_rand)
+
+    a, b = Tensor(a_rand.numpy()), Tensor(b_rand.numpy())
+    c = Tensor.empty(N, N)
+    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm, grad_fxn=backward_gemm_custom if custom_backward_gemm else backward_gemm)[0]
+    tst.sum().backward()
+    grad_a, grad_b = a.grad, b.grad
+    Tensor.realize(tst, grad_a, grad_b)
+
+    a, b = Tensor(a_rand.numpy()), Tensor(b_rand.numpy())
+    ref = (a@b)
+    ref.sum().backward()
+    real_grad_a, real_grad_b = a.grad, b.grad
+    Tensor.realize(ref, real_grad_a, real_grad_b)
+
+    err = (tst - ref).square().max()
+    self.assertLess(err.item(), 1e-6)
+
+    err = (grad_a - real_grad_a).square().max()
+    self.assertLess(err.item(), 1e-6)
+
+    err = (grad_b - real_grad_b).square().max()
+    self.assertLess(err.item(), 1e-6)
+
+  def test_simple_qkv(self):
+    N, d = 8, 4
+    Q = Tensor.randn(N, d)
+    K = Tensor.randn(N, d)
+    V = Tensor.randn(N, d)
+    O = Tensor.empty(N, d)
+
+    O_custom = Tensor.custom_kernel(O, Q, K, V, fxn=lambda o,q,k,v: simple_qkv_kernel(o,q,k,v))[0]
+    O_ref = ((Q @ K.T) / (d ** 0.5)) @ V
+
+    Tensor.realize(O_custom, O_ref)
+    err = (O_custom - O_ref).square().max()
+    self.assertLess(err.item(), 1e-6)
+
+  def test_multi_after_schedule_order(self):
+    """Test correct scheduling order when custom_kernel has multiple outputs.
+
+    custom_kernel with 4 arguments creates 4 AFTERs from the same kernel.
+    The custom_kernel depends on both A2 and B2, so it must be scheduled after both.
+    E only depends on A2, so E can run before custom_kernel finishes waiting for B2.
+
+    Expected schedule order: [A2, B2, E, custom_addmul, final_sum]
+    The custom_addmul kernel should be at index 3.
+    """
+
+    A, B = Tensor.empty(4, 4), Tensor.empty(4, 4)
+    A2 = (A + 1).contiguous()                      # kernel 0: depends on A
+    B2 = (B * 2).contiguous()                      # kernel 1: depends on B
+    C, D = Tensor.empty(4, 4), Tensor.empty(4, 4)
+    C, D, _, _ = Tensor.custom_kernel(C, D, A2, B2, fxn=custom_elementwise_addmul_kernel)  # depends on A2 AND B2
+    E = (A2 * 3).contiguous()                      # kernel 2: depends only on A2
+    result = (C + D + E).sum()                     # kernel 3: custom_addmul, then kernel 4: sum
+    schedule = result.schedule_linear().src
+
+    # Find the custom_addmul kernel position
+    custom_idx = next((i for i, item in enumerate(schedule)
+                       if hasattr(item.src[0], "arg") and hasattr(item.src[0].arg, "name")
+                       and "custom_addmul" in item.src[0].arg.name), None)
+
+    self.assertIsNotNone(custom_idx, "custom_addmul kernel not found in schedule")
+    self.assertEqual(custom_idx, 3, f"custom_addmul should be at index 3, got {custom_idx}")
+
+  def test_invalids_into_custom_kernel_no_empty_kernel(self):
+    from tinygrad.engine.realize import compile_linear
+    a = Tensor.full((4, 4), 3.).contiguous()
+    b = Tensor.full((4, 4), 2.).contiguous()
+    Tensor.realize(a, b)
+    out = Tensor.invalids(*a.shape, dtype=a.dtype)
+    out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel)
+    compiled = compile_linear(out.schedule_linear())
+    for call in compiled.src:
+      prg = call.src[0]
+      if prg.op is not Ops.PROGRAM: continue
+      self.assertTrue(len(prg.arg.globals) > 0, f"empty kernel compiled (no globals): name={prg.arg.name}")
+
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "kernel timing not supported")
+  def test_invalids_into_custom_kernel_with_beam(self):
+    a = Tensor.full((4, 4), 3.).contiguous()
+    b = Tensor.full((4, 4), 2.).contiguous()
+    Tensor.realize(a, b)
+    with Context(BEAM=1, IGNORE_BEAM_CACHE=1):
+      out = Tensor.invalids(*a.shape, dtype=a.dtype)
+      out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel)
+      result = out.flatten().tolist()
+    self.assertTrue(all(x == 5 for x in result), f"expected all 5.0, got {result}")
+
+  @unittest.skip("what are anonymous buffers?")
+  def test_anonymous_buffers_in_function(self):
+    """Test that custom kernels with anonymous output buffers work inside @function."""
+    a = Tensor.full((4, 4), 3.).contiguous()
+    b = Tensor.full((4, 4), 2.).contiguous()
+    Tensor.realize(a, b)
+
+    def custom_add_with_tmp(o1:UOp, o2:UOp, A:UOp, B:UOp) -> UOp:
+      o1,o2,A,B = o1.flatten(), o2.flatten(), A.flatten(), B.flatten()
+      i = UOp.range(o1.numel(), 0)
+      store_o1 = o1[i].store(A[i]+B[i])
+      store_o2 = o2[i].store(A[i]+B[i]+2)
+      return UOp.group(store_o1, store_o2).end(i).sink(arg=KernelInfo(name=f"add_with_tmp_{o1.numel()}")).simplify()
+
+    from tinygrad import function
+    @function(precompile=True)
+    def run(x:Tensor, w:Tensor) -> Tensor:
+      out = Tensor.invalids(*x.shape, dtype=x.dtype)
+      tmp = Tensor.invalids(*x.shape, dtype=x.dtype)
+      out, tmp = Tensor.custom_kernel(out, tmp, x, w, fxn=custom_add_with_tmp)[:2]
+      return out+tmp
+
+    result = run(a, b).flatten().tolist()
+    expected = (3+2)*2+2
+    assert all(x == expected for x in result), f"expected all {expected}, got {result}"
+
+  def test_custom_kernel_sched(self, use_custom=False):
+    x = Tensor.arange(32).reshape(8, 4).clone().realize()
+    y = Tensor.empty_like(x)
+    y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
+    if use_custom:
+      z = Tensor.empty_like(x)
+      z = Tensor.custom_kernel(y, y.T.T, fxn=custom_add_one_kernel)[0]
+    else: z = y.T.T+1
+    GlobalCounters.reset()
+    z.realize()
+    self.assertEqual(GlobalCounters.kernel_count, 2)
+    self.assertEqual(z.tolist(), x.add(2).tolist())
+
+  @unittest.expectedFailure
+  def test_custom_kernel_sched_copy(self): self.test_custom_kernel_sched(use_custom=True)
+
+  @unittest.expectedFailure
+  def test_sliced_buffer_function(self):
+    x = Tensor.arange(32).reshape(8, 4).clone().realize()
+    from tinygrad import function
+    @function(precompile=True)
+    def run(x:Tensor) -> Tensor:
+      y = Tensor.invalids(*x.shape, dtype=x.dtype)
+      return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
+    GlobalCounters.reset()
+    y = run(x[0]).realize()
+    # it's copying the input and the output
+    self.assertEqual(GlobalCounters.kernel_count, 1)
+    self.assertEqual(y.tolist(), [1, 2, 3, 4])
+
+  @Context(DEV="CPU")
+  def test_simple_from_source(self):
+    a = Tensor([0., 1., 2.]).realize()
+
+    src = "void test_src(float* restrict a) { a[0] = 1.0; }"
+    # TODO: it currently requires a compiler for Ops.BINARY
+    from tinygrad.device import Device
+    binary = Device[a.device].renderer.compiler.compile(src)
+    def custom_src_kernel(A:UOp) -> UOp:
+      sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="CPU"), UOp(Ops.LINEAR, src=tuple(sink.toposort())),
+                                   UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
+
+    a = Tensor.custom_kernel(a, fxn=custom_src_kernel)[0]
+    self.assertEqual(a.tolist(), [1., 1., 2.])
+
+class TestUOpReduce(unittest.TestCase):
+  def test_uop_sum(self):
+    a = Tensor([1.0, 2, 3, 4, 5])
+    self.assertAlmostEqual(Tensor(a.uop.sum(axis=0)).item(), 15.0)
+
+  def test_uop_sum_2d(self):
+    a = Tensor.arange(6).reshape(2, 3).float()
+    result = Tensor(a.uop.sum(axis=1)).numpy()
+    assert result[0] == 3 and result[1] == 12
+
+  def test_uop_sum_all(self):
+    a = Tensor.arange(6).reshape(2, 3).float()
+    self.assertAlmostEqual(Tensor(a.uop.sum()).item(), 15.0)
+
+  def test_uop_sum_keepdim(self):
+    a = Tensor.arange(6).reshape(2, 3).float()
+    result = Tensor(a.uop.sum(axis=1, keepdim=True))
+    assert result.shape == (2, 1)
+
+  def test_uop_sum_negative_axis(self):
+    a = Tensor.arange(6).reshape(2, 3).float()
+    result = Tensor(a.uop.sum(axis=-1)).numpy()
+    assert result[0] == 3 and result[1] == 12
+
+  def test_uop_sum_multi_axis(self):
+    a = Tensor.arange(24).reshape(2, 3, 4).float()
+    ref = a.sum(axis=(0, 2)).numpy()
+    result = Tensor(a.uop.sum(axis=(0, 2))).numpy()
+    for i in range(3): self.assertAlmostEqual(result[i], ref[i])
+
+  def test_uop_sum_dtype(self):
+    a = Tensor([1.0, 2, 3], dtype=dtypes.float16)
+    result = Tensor(a.uop.sum(axis=0, dtype=dtypes.float32))
+    self.assertEqual(result.dtype, dtypes.float)
+    self.assertAlmostEqual(result.item(), 6.0, places=2)
+
+  def test_uop_prod(self):
+    a = Tensor([1.0, 2, 3, 4, 5])
+    self.assertAlmostEqual(Tensor(a.uop.prod(axis=0)).item(), 120.0)
+
+  def test_uop_max(self):
+    a = Tensor([1.0, 5, 3, 2, 4])
+    self.assertAlmostEqual(Tensor(a.uop.max(axis=0)).item(), 5.0)
+
+  def test_uop_max_2d(self):
+    a = Tensor([[1, 5, 3], [4, 2, 6]]).float()
+    result = Tensor(a.uop.max(axis=0)).numpy()
+    assert result[0] == 4 and result[1] == 5 and result[2] == 6
+
+  def test_uop_std(self):
+    a = Tensor([2.0, 4, 4, 4, 5, 5, 7, 9])
+    self.assertAlmostEqual(Tensor(a.uop.std()).item(), a.std().item(), places=5)
+
+class TestUOpWhere(unittest.TestCase):
+  def test_uop_where_both_const(self):
+    cond = Tensor([True, False, True])
+    result = Tensor(cond.uop.where(1, 0))
+    self.assertEqual(result.tolist(), [1, 0, 1])
+
+    result = Tensor(cond.uop.where(1.5, 0))
+    self.assertEqual(result.tolist(), [1.5, 0, 1.5])
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_device.py

@@ -0,0 +1,9 @@
+import unittest
+from tinygrad import Device
+
+class TestDeviceCount(unittest.TestCase):
+  def test_count(self):
+    self.assertGreaterEqual(Device[Device.DEFAULT].count(), 1)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/backend/test_dtype.py

@@ -0,0 +1,419 @@
+import contextlib, unittest, math
+import numpy as np
+import torch
+from typing import Any, List
+from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES, DEV
+from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype, truncate
+from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
+from tinygrad import Context, Device, Tensor, dtypes
+from hypothesis import given, settings, strategies as strat
+from test.helpers import rand_for_dtype
+from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
+import pytest
+pytestmark = pytest.mark.filterwarnings("ignore")
+
+settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+settings.load_profile("my_profile")
+
+supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
+
+def get_available_cast_dtypes(dtype: DType) -> List[DType]:
+  dts = [v for k, v in DTYPES_DICT.items() if v != dtype and v in supported_dtypes or v in dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)]
+  if dtype in (dtypes.long, dtypes.ulong) and (dtype not in supported_dtypes or dtypes.long in EMULATED_DTYPES.tolist(dtypes)):
+    return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
+  if dtype not in supported_dtypes and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
+  return dts
+
+def _to_torch_storage_type(dtype:DType):
+  if dtype == dtypes.bfloat16: return torch.float32
+  if dtype in dtypes.fp8s: return torch.float32
+  return _to_torch_dtype(dtype)
+
+def _test_to_np(a:Tensor, np_dtype, target):
+  if DEBUG >= 2: print(a)
+  na = a.numpy()
+  if DEBUG >= 2: print(na, na.dtype, a.uop.base.realized)
+  try:
+    assert na.dtype == np_dtype
+    np.testing.assert_allclose(na, target)
+  except AssertionError as e:
+    raise AssertionError(f"\ntensor {a.numpy()} does not match target {target} with np_dtype {np_dtype}") from e
+
+def _test_op(fxn, target_dtype:DType, target):
+  _assert_eq(fxn(), target_dtype, target)
+def _test_cast(a:Tensor, target_dtype:DType):
+  if a.is_floating_point() and dtypes.is_unsigned(target_dtype):
+    # converting negative float to unsigned integer is undefined
+    a = a.abs()
+
+  expected = list(a.numpy().astype(_to_np_dtype(target_dtype)))
+  if target_dtype in dtypes.fp8s: expected = [truncate[target_dtype](x) for x in expected]
+  _test_op(lambda: a.cast(target_dtype), target_dtype, expected)
+def _test_bitcast(a:Tensor, target_dtype:DType, target=None):
+  expected = torch.tensor(a.tolist(), dtype=_to_torch_storage_type(a.dtype)).view(_to_torch_dtype(target_dtype)).tolist()
+  if target_dtype in dtypes.fp8s: expected = [fp8_to_float(x, target_dtype) for x in expected]
+  _test_op(lambda: a.bitcast(target_dtype), target_dtype, target or expected)
+
+class TestDType(unittest.TestCase):
+  DTYPE: Any = None
+  DATA: Any = None
+  @classmethod
+  def setUpClass(cls):
+    if cls.DTYPE is None: raise unittest.SkipTest("base class")
+    cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=cls.DTYPE in supported_dtypes)
+
+  def test_to_np(self):
+    _test_to_np(Tensor(self.DATA, dtype=self.DTYPE), _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))
+
+  def test_casts_to(self):
+    for dtype in get_available_cast_dtypes(self.DTYPE):
+      _test_cast(Tensor(self.DATA, dtype=dtype), self.DTYPE)
+
+  def test_casts_from(self):
+    for dtype in get_available_cast_dtypes(self.DTYPE):
+      _test_cast(Tensor(self.DATA, dtype=self.DTYPE), dtype)
+
+  def test_same_size_ops(self):
+    for dtype in get_available_cast_dtypes(self.DTYPE):
+      if dtype.itemsize == self.DTYPE.itemsize:
+        _test_ops(a_dtype=self.DTYPE, b_dtype=dtype)
+
+  def test_upcast_ops(self):
+    for dtype in get_available_cast_dtypes(self.DTYPE):
+      if dtype.itemsize > self.DTYPE.itemsize:
+        _test_ops(a_dtype=self.DTYPE, b_dtype=dtype)
+
+  def test_upcast_to_ops(self):
+    for dtype in get_available_cast_dtypes(self.DTYPE):
+      if dtype.itemsize < self.DTYPE.itemsize:
+        _test_ops(a_dtype=dtype, b_dtype=self.DTYPE)
+
+  def test_bitcast(self):
+    if self.DTYPE == dtypes.bool: raise unittest.SkipTest("no bools in bitcast")
+    for dtype in get_available_cast_dtypes(self.DTYPE):
+      if dtype != dtypes.bool:
+        _test_bitcast(Tensor(self.DATA[:8], dtype=self.DTYPE), dtype)
+
+  @unittest.skipIf(Device.DEFAULT == "PYTHON", "skip for now")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "skip for now")
+  def test_uint_overflow(self):
+    if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
+    v = self.DTYPE.max
+    _test_to_np(Tensor(v, dtype=self.DTYPE)+2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))+2)
+    _test_to_np(Tensor(v, dtype=self.DTYPE)*2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))*2)
+
+def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
+  target_dtype = target_dtype or least_upper_dtype(a_dtype, b_dtype)
+  if a_dtype == dtypes.bool or b_dtype == dtypes.bool: return
+  _assert_eq(Tensor([1,2,3,4], dtype=a_dtype)+Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [2,4,6,8])
+  _assert_eq((Tensor([1], dtype=a_dtype).cast(b_dtype)+Tensor([1], dtype=a_dtype).cast(b_dtype)).cast(a_dtype), a_dtype, [2])
+  _assert_eq(Tensor([1,2,3,4], dtype=a_dtype)*Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [1,4,9,16])
+  _assert_eq(Tensor([[1,2],[3,4]], dtype=a_dtype)@Tensor.eye(2, dtype=b_dtype), target_dtype, [[1,2],[3,4]])
+  _assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype), target_dtype, 2*np.ones((4,4)))
+  _assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype).clone(), target_dtype, 2*np.ones((4,4)))
+  _assert_eq(Tensor.ones((4,4), dtype=b_dtype).clone(), b_dtype, np.ones((4,4)))
+
+class TestFp8sConversions(unittest.TestCase):
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
+  def test_float_to_fp8e4m3(self, x):
+    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
+
+  def test_float_to_fp8e4m3_extreme_values(self):
+    for x in [FP8E4M3_MAX, FP8E4M3_MAX*1.01, -FP8E4M3_MAX, -FP8E4M3_MAX*1.01, math.inf, -math.inf, math.nan, -math.nan]:
+      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
+
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2_MAX, max_value=FP8E5M2_MAX))
+  def test_float_to_fp8e5m2(self, x):
+    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.float8_e5m2).view(torch.uint8).item())
+
+  def test_float_to_fp8e5m2_extreme_values(self):
+    for x in [FP8E5M2_MAX, FP8E5M2_MAX*1.01, -FP8E5M2_MAX, -FP8E5M2_MAX*1.01, math.inf, -math.inf, math.nan, -math.nan]:
+      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.float8_e5m2).view(torch.uint8).item())
+
+  @given(strat.integers(min_value=0, max_value=255))
+  def test_fp8e4m3_to_float(self, x):
+    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fn).float().item())
+
+  @given(strat.integers(min_value=0, max_value=255))
+  def test_fp8e5m2_to_float(self, x):
+    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2).float().item())
+
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3FNUZ_MAX, max_value=FP8E4M3FNUZ_MAX))
+  def test_float_to_fp8e4m3fnuz(self, x):
+    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
+
+  def test_float_to_fp8e4m3fnuz_extreme_values(self):
+    for x in [FP8E4M3FNUZ_MAX, FP8E4M3FNUZ_MAX*1.01, -FP8E4M3FNUZ_MAX, -FP8E4M3FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
+      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
+
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2FNUZ_MAX, max_value=FP8E5M2FNUZ_MAX))
+  def test_float_to_fp8e5m2fnuz(self, x):
+    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
+
+  def test_float_to_fp8e5m2fnuz_extreme_values(self):
+    for x in [FP8E5M2FNUZ_MAX, FP8E5M2FNUZ_MAX*1.01, -FP8E5M2FNUZ_MAX, -FP8E5M2FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
+      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
+
+  @given(strat.integers(min_value=0, max_value=255))
+  def test_fp8e4m3fnuz_to_float(self, x):
+    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fnuz).float().item())
+
+  @given(strat.integers(min_value=0, max_value=255))
+  def test_fp8e5m2fnuz_to_float(self, x):
+    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item())
+
+class TestBFloat16DType(unittest.TestCase):
+  def test_bf16_to_float(self):
+    _test_cast(Tensor([100000], dtype=dtypes.bfloat16), dtypes.float32)
+
+  def test_float_to_bf16(self):
+    _test_cast(Tensor([100000], dtype=dtypes.float32), dtypes.bfloat16)
+
+  def test_bf16(self):
+    t = Tensor([10000, -1, -1000, -10000, 20]).cast(dtypes.bfloat16)
+    t.realize()
+    back = t.cast(dtypes.float32)
+    assert tuple(back.numpy().tolist()) == (9984., -1, -1000, -9984, 20)
+
+class TestBFloat16DTypeCast(unittest.TestCase):
+  def test_f16_to_bf16_conversion(self):
+    original_tensor = Tensor([1.0, 2.0, 3.0], dtype=dtypes.float16)
+    converted_tensor = original_tensor.cast(dtypes.bfloat16)
+    self.assertEqual(converted_tensor.dtype, dtypes.bfloat16)
+    back_to_float32 = converted_tensor.cast(dtypes.float32)
+    original_to_float32 = original_tensor.cast(dtypes.float32)
+    np.testing.assert_allclose(back_to_float32.numpy(), original_to_float32.numpy(), rtol=1e-2, atol=1e-3)
+
+  def test_f16_to_bf16_edge_cases(self):
+    edge_cases = Tensor([0.0, -0.0, float('inf'), float('-inf'), float('nan')], dtype=dtypes.float16)
+    converted = edge_cases.cast(dtypes.bfloat16).cast(dtypes.float32)
+    np.testing.assert_equal(converted.numpy(), edge_cases.cast(dtypes.float32).numpy())
+
+  def test_f16_to_bf16_range_precision(self):
+    large_value = Tensor([65504.0], dtype=dtypes.float16)  # Max representable in float16
+    small_value = Tensor([6.1035e-5], dtype=dtypes.float16)  # Smallest positive normal float16
+    large_converted = large_value.cast(dtypes.bfloat16).cast(dtypes.float32)
+    small_converted = small_value.cast(dtypes.bfloat16).cast(dtypes.float32)
+    np.testing.assert_allclose(large_converted.numpy(), large_value.cast(dtypes.float32).numpy(), rtol=1e-2, atol=1e-3)
+    np.testing.assert_equal(small_converted.numpy(), small_value.cast(dtypes.float32).numpy())
+
+  def test_f16_to_bf16_randomized(self):
+    np.random.seed(42)  # For reproducibility
+    random_values = Tensor(np.random.uniform(-65504, 65504, 1000), dtype=dtypes.float16)
+    converted = random_values.cast(dtypes.bfloat16).cast(dtypes.float32)
+    np.testing.assert_allclose(converted.numpy(), random_values.cast(dtypes.float32).numpy(), rtol=1e-2, atol=1e-3)
+
+class TestHalfDType(TestDType): DTYPE = dtypes.half
+
+class TestEmulatedHalf(TestHalfDType):
+  @classmethod
+  def setUpClass(cls):
+    cls.stack = contextlib.ExitStack()
+    cls.stack.enter_context(Context(EMULATED_DTYPES="half"))
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
+
+  @classmethod
+  def tearDownClass(cls): cls.stack.close()
+
+
+class TestFloatDType(TestDType):
+  DTYPE = dtypes.float
+
+  def test_float_to_uint(self):
+    _test_op(lambda: Tensor([-0.9, -0.3, 1.2], dtype=dtypes.float32).cast(dtypes.uint32), dtypes.uint32,
+             [0, 0, 1])
+
+@unittest.skipUnless(dtypes.double in supported_dtypes, f"no double on {Device.DEFAULT}")
+class TestDoubleDType(TestDType):
+  DTYPE = dtypes.double
+  @unittest.skipIf((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"CUDA", "NV"}) or \
+   isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "conversion not supported on CI CUDA, PTX, and NIR")  # TODO: why not?
+  def test_float64_increased_precision(self):
+    for func in [
+      lambda t: t.exp(),
+      lambda t: t.exp2(),
+      lambda t: t.log(),
+      lambda t: t.log2(),
+      lambda t: t.sqrt(),
+      lambda t: t.rsqrt(),
+      lambda t: t.sin(),
+      lambda t: t.cos(),
+      lambda t: t.tan(),
+      lambda t: t.sigmoid(),
+    ]:
+      a = [2, 3, 4]
+      np.testing.assert_allclose(func(Tensor(a, dtype=self.DTYPE)).numpy(), func(torch.tensor(a, dtype=torch.float64)), rtol=1e-12, atol=1e-12)
+
+  def test_float64_to_float32_cast_inf(self):
+    _test_op(lambda: Tensor([3.4e40, 3.4e38, 1, 0], dtype=dtypes.float64).cast(dtypes.float32),
+             dtypes.float32, [float('inf'), 3.4e38, 1, 0])
+
+
+class TestInt8DType(TestDType):
+  DTYPE = dtypes.int8
+  @unittest.skipIf(Device.DEFAULT == "CUDA" or isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "cuda saturation works differently")
+  def test_int8_to_uint8_negative(self):
+    _test_op(lambda: Tensor([-1, -2, -3, -4], dtype=dtypes.int8).cast(dtypes.uint8), dtypes.uint8, [255, 254, 253, 252])
+
+  def test_int8_to_uint16_negative(self):
+    _test_op(lambda: Tensor([-1, -2, -3, -4], dtype=dtypes.int8).cast(dtypes.uint16), dtypes.uint16, [2**16-1, 2**16-2, 2**16-3, 2**16-4])
+
+  def test_bitcast_alt(self):
+    a = Tensor([72, -90, 27, 40, -53, 70, 96, 51], dtype=dtypes.int8).bitcast(dtypes.short)
+    self.assertListEqual(a.tolist(), [-22968, 10267, 18123, 13152])
+
+class TestUint8DType(TestDType):
+  DTYPE = dtypes.uint8
+  @unittest.skipIf(Device.DEFAULT == "CUDA" or isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "cuda saturation works differently")
+  def test_uint8_to_int8_overflow(self):
+    _test_op(lambda: Tensor([255, 254, 253, 252], dtype=dtypes.uint8).cast(dtypes.int8), dtypes.int8, [-1, -2, -3, -4])
+
+class TestBitCast(unittest.TestCase):
+  @given(strat.sampled_from(dtype_ints + dtype_floats), strat.sampled_from(dtype_ints + dtype_floats))
+  def test_shape_change_bitcast(self, dt1, dt2):
+    data = rand_for_dtype(dt1, 32).reshape(2, 2, 8)
+    expected = torch.tensor(data.tolist(), dtype=_to_torch_storage_type(dt1)).view(_to_torch_dtype(dt2))
+    if dt2 in dtypes.fp8s:
+      expected = torch.tensor([fp8_to_float(x, dt2) for x in expected.view(-1).tolist()]).view_as(expected)
+    _test_op(lambda: Tensor(data, dtype=dt1).bitcast(dt2), dt2, expected.tolist())
+
+  def test_shape_change_bitcast_exceptions(self):
+    with self.assertRaises(RuntimeError):
+      # should fail because 3 int8 is 3 bytes but float16 is two and 3 isn't a multiple of 2
+      Tensor.empty((3,), dtype=dtypes.int8).bitcast(dtypes.float16)
+
+  def test_bitcast_float_to_int32(self):
+    a = Tensor([1.,2,3])
+    b = a.bitcast(dtypes.int32)
+    assert b.numpy()[0] == 0x3f800000
+
+  def test_bitcast_upcasted(self):
+    a = Tensor.zeros(100, 4, dtype=dtypes.int32).contiguous() + 0x3f800000
+    b = a.bitcast(dtypes.float32)
+    assert b.numpy()[0,0] == 1.
+
+class TestInt16DType(TestDType): DTYPE = dtypes.int16
+
+class TestUint16DType(TestDType):
+  DTYPE = dtypes.uint16
+
+  def test_uint16_to_int8_overflow(self):
+    _test_op(lambda: Tensor([2**16-1, 2**16-2, 1, 0], dtype=dtypes.uint16).cast(dtypes.int8), dtypes.int8, [-1, -2, 1, 0])
+
+class TestInt32DType(TestDType): DTYPE = dtypes.int32
+class TestUint32DType(TestDType): DTYPE = dtypes.uint32
+
+class TestInt64DType(TestDType): DTYPE = dtypes.int64
+
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
+class TestEmulatedInt64DType(TestInt64DType):
+  @classmethod
+  def setUpClass(cls):
+    cls.stack = contextlib.ExitStack()
+    cls.stack.enter_context(Context(EMULATED_DTYPES="long"))
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
+
+  @classmethod
+  def tearDownClass(cls): cls.stack.close()
+
+class TestUint64DType(TestDType):
+  DTYPE = dtypes.uint64
+  def test_uint64_load(self):
+    assert Tensor(2**64 - 1, dtype=dtypes.uint64).numpy() == 2**64 - 1
+
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
+class TestEmulatedUInt64DType(TestUint64DType):
+  @classmethod
+  def setUpClass(cls):
+    cls.stack = contextlib.ExitStack()
+    cls.stack.enter_context(Context(EMULATED_DTYPES="long"))
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
+
+  @classmethod
+  def tearDownClass(cls): cls.stack.close()
+
+class TestBoolDType(TestDType): DTYPE = dtypes.bool
+
+class TestBFloat16Type(TestDType): DTYPE = dtypes.bfloat16
+
+class TestEmulatedBFloat16Type(TestBFloat16Type):
+  @classmethod
+  def setUpClass(cls):
+    cls.stack = contextlib.ExitStack()
+    cls.stack.enter_context(Context(EMULATED_DTYPES="bfloat16"))
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
+
+  @classmethod
+  def tearDownClass(cls): cls.stack.close()
+
+class TestFp8e4m3(TestDType): DTYPE = dtypes.fp8e4m3
+
+class TestEmulatedFp8e4m3(TestFp8e4m3):
+  @classmethod
+  def setUpClass(cls):
+    cls.stack = contextlib.ExitStack()
+    cls.stack.enter_context(Context(EMULATED_DTYPES="fp8e4m3"))
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
+
+  @classmethod
+  def tearDownClass(cls): cls.stack.close()
+
+class TestFp8e5m2(TestDType): DTYPE = dtypes.fp8e5m2
+
+class TestEmulatedFp8e5m2(TestFp8e5m2):
+  @classmethod
+  def setUpClass(cls):
+    cls.stack = contextlib.ExitStack()
+    cls.stack.enter_context(Context(EMULATED_DTYPES="fp8e5m2"))
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
+
+  @classmethod
+  def tearDownClass(cls): cls.stack.close()
+
+class TestImplicitFunctionTypeChange(unittest.TestCase):
+  def test_functions(self):
+    result = []
+    for func in [
+      lambda t: t.exp(),
+      lambda t: t.exp2(),
+      lambda t: t.log(),
+      lambda t: t.log2(),
+      lambda t: t.sqrt(),
+      lambda t: t.sin(),
+    ]:
+      t = func(Tensor([4.0, 3.0])).max() == func(Tensor([4.0, 3.0]))
+      result.append(t.numpy().sum())
+    assert all(result)
+
+class TestTensorMethod(unittest.TestCase):
+  @given(strat.sampled_from(core_dtypes))
+  def test_abs_diff(self, dt):
+    if dt == dtypes.bool or dt not in supported_dtypes: return
+    a, b = Tensor([2], dtype=dt), Tensor([1], dtype=dt)
+    ret = (a - b).abs()
+    np.testing.assert_allclose(ret.numpy(), np.abs(a.numpy()-b.numpy()))
+
+class TestDtypeUsage(unittest.TestCase):
+  def test_max_w_alu(self):
+    for d in dtypes.ints:
+      if d in supported_dtypes:
+        t = Tensor([[1, 2], [3, 4]], dtype=d)
+        (t*t).max().item()
+
+@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
+class TestOpsBFloat16(unittest.TestCase):
+  def test_cast(self):
+    # TODO: helper_test_op breaks in unrelated part
+    data = [60000.0, 70000.0, 80000.0]
+    np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy())
+
+  # some CPUs there is no native bfloat16 sqrt
+  @unittest.skipIf(Device.DEFAULT == "CPU", "no approximation")
+  def test_no_approximation(self):
+    data = [326.0, 339.0, 10603200512.0]
+    expected = torch.tensor(data, dtype=torch.bfloat16).sqrt().float().numpy()
+    np.testing.assert_allclose(Tensor(data, dtype=dtypes.bfloat16).sqrt().numpy(), expected)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_dtype_alu.py

@@ -0,0 +1,383 @@
+import unittest, operator, math
+from tinygrad import Context, Tensor, dtypes, Device
+from tinygrad.dtype import DType, truncate, fp8_to_float
+from tinygrad.helpers import EMULATED_DTYPES, DEV, getenv
+from tinygrad.tensor import _to_np_dtype
+from tinygrad.runtime.ops_python import from_storage_scalar
+from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
+from tinygrad.uop import Ops
+import numpy as np
+import pytest
+from hypothesis import assume, given, strategies as strat, settings
+
+pytestmark = pytest.mark.filterwarnings("ignore")
+
+settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+settings.load_profile("my_profile")
+print(settings.default)
+
+dtypes_float = (dtypes.float16, dtypes.float32, dtypes.float64)
+dtypes_int = (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
+dtypes_bool = (dtypes.bool,)
+binary_operations = [operator.add, operator.sub, operator.mul, operator.lt, operator.eq]
+
+integer_binary_operations = binary_operations + [(Tensor.bitwise_xor, np.bitwise_xor), (Tensor.bitwise_and, np.bitwise_and),
+                                                 (Tensor.bitwise_or, np.bitwise_or), (Tensor.maximum, np.maximum), operator.mod]
+integer_unary_operations = [operator.neg]
+unary_operations = [(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin),
+                    (Tensor.sqrt, np.sqrt), (Tensor.reciprocal, np.reciprocal), (Tensor.cos, np.cos)]
+
+# TODO: enable this (this is a dtype issue)
+#binary_operations.append(operator.truediv)
+
+# TODO: CI CUDA segfaults on sin, WEBGPU and NIR sines are not precise enough for large numbers
+if ((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"})
+    or Device.DEFAULT == "WEBGPU" or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer)):
+  unary_operations.remove((Tensor.sin, np.sin))
+  unary_operations.remove((Tensor.cos, np.cos))
+
+# transcendental isn't accurate enough
+if Ops.SQRT not in Device[Device.DEFAULT].renderer.code_for_op: unary_operations.remove((Tensor.sqrt, np.sqrt))
+
+supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
+
+class ht:
+  float64 = strat.floats(width=64, allow_subnormal=False)
+  float32 = strat.floats(width=32, allow_subnormal=False)
+  float16 = strat.floats(width=16, allow_subnormal=False)
+  uint8 = strat.integers(0, 255)
+  uint16 = strat.integers(0, 65535)
+  uint32 = strat.integers(0, 2**32-1)
+  uint64 = strat.integers(0, 2**64-1)
+  int8 = strat.integers(-128, 127)
+  int16 = strat.integers(-32768, 32767)
+  int32 = strat.integers(-2147483648, 2147483647)
+  int64 = strat.integers(-9223372036854775808, 9223372036854775807)
+  bool = strat.booleans()
+ht.bfloat16 = ht.uint16.filter(lambda x: ((x >> 7) & 0xFF) != 0)  # filter subnormal bfloat16
+ht.fp8e4m3 = ht.uint8
+ht.fp8e5m2 = ht.uint8
+ht.fp8e4m3fnuz = ht.uint8
+ht.fp8e5m2fnuz = ht.uint8
+
+def universal_test(a, b, dtype, op):
+  if not isinstance(op, tuple): op = (op, op)
+  if op[0] == operator.mod and b == 0: return
+  # lt and max with nan is undefined in tinygrad
+  if op[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
+  ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype)
+  if dtype in dtypes.fp8s and op[0] not in (operator.lt, operator.eq):
+    tensor_value = fp8_to_float((op[0](ta.realize(), tb.realize())).bitcast(dtypes.uint8).item(), dtype)
+    numpy_value = truncate[dtype](op[1](ta.numpy(), tb.numpy()).item())
+  else: tensor_value, numpy_value = (op[0](ta, tb)).numpy(), op[1](ta.numpy(), tb.numpy())
+  if dtype in dtypes.floats:
+    if dtype not in supported_dtypes or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
+      fe, fm = dtypes.finfo(dtype)
+      atol, rtol = 2 ** (2 - (1 << (fe - 1))), 2 ** (-fm)
+    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1),
+                        dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz:(5e-1, 5e-1)}.get(dtype, (1e-10, 1e-7))
+    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
+  else: np.testing.assert_equal(tensor_value, numpy_value)
+
+def universal_test_unary(a, dtype, op):
+  if not isinstance(op, tuple): op = (op, op)
+  ta = Tensor([a], dtype=dtype)
+  # TODO: cos does not match for large input
+  if op[0] == Tensor.cos and abs(a) > 30: return
+  if op[0] == Tensor.log and a <= 0: return
+  if dtype in dtypes.fp8s:
+    # denormals are zero
+    if (dtype in EMULATED_DTYPES.tolist(dtypes) or dtype not in supported_dtypes
+        and abs(ta.numpy().item()) < 0.015625): return
+    tensor_value = fp8_to_float(op[0](ta.realize()).bitcast(dtypes.uint8).item(), dtype)
+    numpy_value = truncate[dtype](v:=op[1](ta.numpy()).item())
+    # cuda cast f32 inf to f8 MAX, amd cast it to nan(E4M3)/inf(E5M2)
+    if math.isinf(v): return
+  else: tensor_value, numpy_value = op[0](ta).numpy(), op[1](ta.numpy())
+  if dtype in dtypes.floats:
+    atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2),
+      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1),
+      dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz: (5e-1, 5e-1)}.get(dtype, (1e-6, 1e-5))
+    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
+  else: np.testing.assert_equal(tensor_value, numpy_value)
+
+def universal_test_cast(a, in_dtype, dtype):
+  tensor_value = Tensor([a], dtype=in_dtype).cast(dtype)
+  numpy_value = np.array([a], dtype=_to_np_dtype(in_dtype)).astype(_to_np_dtype(dtype))
+  np.testing.assert_equal(tensor_value.numpy(), numpy_value)
+
+@unittest.skipIf(Device.DEFAULT == "WEBGPU", "Inf and nan cases are wrong on WebGPU")
+def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType):
+  if not isinstance(op1, tuple): op1 = (op1, op1)
+  if not isinstance(op2, tuple): op2 = (op2, op2)
+  if op1[0] == operator.mod and b == 0: return
+  # lt and max with nan is undefined in tinygrad
+  if op1[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
+  if op2[0] in (operator.lt, Tensor.maximum) and math.isnan(c): return
+  at, bt, ct = Tensor([a], dtype=d1), Tensor([b], dtype=d1), Tensor([c], dtype=d2)
+  an, bn, cn = np.array([a]).astype(_to_np_dtype(d1)), np.array([b]).astype(_to_np_dtype(d1)), np.array([c]).astype(_to_np_dtype(d2))
+  tensor_value = op2[0](op1[0](at, bt).cast(d2), ct).numpy()
+  numpy_value = op2[1](op1[1](an, bn).astype(_to_np_dtype(d2)), cn)
+  np.testing.assert_allclose(tensor_value, numpy_value, rtol=1e-6 if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) else 1e-7)
+
+class TestDTypeALU(unittest.TestCase):
+  @unittest.skipUnless(dtypes.float64 in supported_dtypes, f"no float64 on {Device.DEFAULT}")
+  @given(ht.float64, ht.float64, strat.sampled_from(binary_operations))
+  def test_float64(self, a, b, op): universal_test(a, b, dtypes.float64, op)
+
+  @given(ht.float32, ht.float32, strat.sampled_from(binary_operations))
+  def test_float32(self, a, b, op): universal_test(a, b, dtypes.float32, op)
+
+  @unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
+  @given(ht.float16, ht.float16, strat.sampled_from(binary_operations))
+  def test_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)
+
+  @given(ht.float16, ht.float16, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="half")
+  def test_emulated_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)
+
+  @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
+  @given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
+  def test_bfloat16(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(b, dtypes.bfloat16), dtypes.bfloat16, op)
+
+  @given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="bfloat16")
+  def test_emulated_bfloat16(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(b, dtypes.bfloat16), dtypes.bfloat16, op)
+
+  @unittest.skipUnless(dtypes.fp8e4m3 in supported_dtypes, f"no fp8e4m3 on {Device.DEFAULT}")
+  @given(ht.fp8e4m3, ht.fp8e4m3, strat.sampled_from(binary_operations))
+  def test_fp8e4m3(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op)
+
+  @given(ht.fp8e4m3, ht.fp8e4m3, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="fp8e4m3")
+  def test_emulated_fp8e4m3(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op)
+
+  @unittest.skipUnless(dtypes.fp8e5m2 in supported_dtypes, f"no fp8e5m2 on {Device.DEFAULT}")
+  @given(ht.fp8e5m2, ht.fp8e5m2, strat.sampled_from(binary_operations))
+  def test_fp8e5m2(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)
+
+  @given(ht.fp8e5m2, ht.fp8e5m2, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="fp8e5m2")
+  def test_emulated_fp8e5m2(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)
+
+  @unittest.skipUnless(dtypes.fp8e4m3fnuz in supported_dtypes, f"no fp8e4m3fnuz on {Device.DEFAULT}")
+  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
+  def test_fp8e4m3fnuz(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
+
+  @unittest.skipUnless(dtypes.fp8e5m2fnuz in supported_dtypes, f"no fp8e5m2fnuz on {Device.DEFAULT}")
+  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
+  def test_fp8e5m2fnuz(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
+
+  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="fp8e4m3fnuz")
+  def test_emulated_fp8e4m3fnuz(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
+
+  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="fp8e5m2fnuz")
+  def test_emulated_fp8e5m2fnuz(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
+
+  @given(ht.float32, strat.sampled_from(unary_operations))
+  def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)
+
+  @unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
+  @given(ht.float16, strat.sampled_from(unary_operations))
+  def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
+
+  @given(ht.float16, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="half")
+  def test_emulated_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
+
+  @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
+  @given(ht.bfloat16, strat.sampled_from(unary_operations))
+  def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
+
+  @given(ht.bfloat16, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="bfloat16")
+  def test_emulated_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
+
+  @unittest.skipUnless(dtypes.fp8e4m3 in supported_dtypes, f"no fp8e4m3 on {Device.DEFAULT}")
+  @given(ht.fp8e4m3, strat.sampled_from(unary_operations))
+  def test_fp8e4m3_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3), dtypes.fp8e4m3, op)
+
+  @given(ht.fp8e4m3, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="fp8e4m3")
+  def test_emulated_fp8e4m3_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3), dtypes.fp8e4m3, op)
+
+  @unittest.skipUnless(dtypes.fp8e5m2 in supported_dtypes, f"no fp8e5m2 on {Device.DEFAULT}")
+  @given(ht.fp8e5m2, strat.sampled_from(unary_operations))
+  def test_fp8e5m2_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)
+
+  @given(ht.fp8e5m2, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="fp8e5m2")
+  def test_emulated_fp8e5m2_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)
+
+  @unittest.skipUnless(dtypes.fp8e4m3fnuz in supported_dtypes, f"no fp8e4m3fnuz on {Device.DEFAULT}")
+  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
+  def test_fp8e4m3fnuz_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
+
+  @unittest.skipUnless(dtypes.fp8e5m2fnuz in supported_dtypes, f"no fp8e5m2fnuz on {Device.DEFAULT}")
+  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
+  def test_fp8e5m2fnuz_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
+
+  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="fp8e4m3fnuz")
+  def test_emulated_fp8e4m3fnuz_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
+
+  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="fp8e5m2fnuz")
+  def test_emulated_fp8e5m2fnuz_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
+
+  @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
+  def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)
+
+  @unittest.skipUnless(dtypes.uint16 in supported_dtypes, f"no uint16 on {Device.DEFAULT}")
+  @given(ht.uint16, ht.uint16, strat.sampled_from(integer_binary_operations))
+  def test_uint16(self, a, b, op): universal_test(a, b, dtypes.uint16, op)
+
+  @unittest.skipUnless(dtypes.uint32 in supported_dtypes, f"no uint32 on {Device.DEFAULT}")
+  @given(ht.uint32, ht.uint32, strat.sampled_from(integer_binary_operations))
+  def test_uint32(self, a, b, op): universal_test(a, b, dtypes.uint32, op)
+
+  @unittest.skipUnless(dtypes.uint64 in supported_dtypes, f"no uint64 on {Device.DEFAULT}")
+  @given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
+  def test_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
+
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
+  @given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
+  @Context(EMULATED_DTYPES="long")
+  def test_emulated_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
+
+  @given(ht.int8, ht.int8, strat.sampled_from(integer_binary_operations))
+  def test_int8(self, a, b, op): universal_test(a, b, dtypes.int8, op)
+
+  @given(ht.int16, ht.int16, strat.sampled_from(integer_binary_operations))
+  def test_int16(self, a, b, op): universal_test(a, b, dtypes.int16, op)
+
+  @given(ht.int32, ht.int32, strat.sampled_from(integer_binary_operations))
+  def test_int32(self, a, b, op): universal_test(a, b, dtypes.int32, op)
+
+  @given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
+  def test_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
+
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
+  @given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
+  @Context(EMULATED_DTYPES="long")
+  def test_emulated_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
+
+  @given(ht.uint8, strat.sampled_from(integer_unary_operations))
+  def test_uint8_unary(self, a, op): universal_test_unary(a, dtypes.uint8, op)
+
+  @unittest.skipUnless(dtypes.uint16 in supported_dtypes, f"no uint16 on {Device.DEFAULT}")
+  @given(ht.uint16, strat.sampled_from(integer_unary_operations))
+  def test_uint16_unary(self, a, op): universal_test_unary(a, dtypes.uint16, op)
+
+  @unittest.skipUnless(dtypes.uint32 in supported_dtypes, f"no uint32 on {Device.DEFAULT}")
+  @given(ht.uint32, strat.sampled_from(integer_unary_operations))
+  def test_uint32_unary(self, a, op): universal_test_unary(a, dtypes.uint32, op)
+
+  @unittest.skipUnless(dtypes.uint64 in supported_dtypes, f"no uint64 on {Device.DEFAULT}")
+  @given(ht.uint64, strat.sampled_from(integer_unary_operations))
+  def test_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)
+
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
+  @given(ht.uint64, strat.sampled_from(integer_unary_operations))
+  @Context(EMULATED_DTYPES="long")
+  def test_emulated_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)
+
+  @given(ht.int8, strat.sampled_from(integer_unary_operations))
+  def test_int8_unary(self, a, op): universal_test_unary(a, dtypes.int8, op)
+
+  @given(ht.int16, strat.sampled_from(integer_unary_operations))
+  def test_int16_unary(self, a, op): universal_test_unary(a, dtypes.int16, op)
+
+  @given(ht.int32, strat.sampled_from(integer_unary_operations))
+  def test_int32_unary(self, a, op): universal_test_unary(a, dtypes.int32, op)
+
+  @given(ht.int64, strat.sampled_from(integer_unary_operations))
+  def test_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64, op)
+
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
+  @given(ht.int64, strat.sampled_from(integer_unary_operations))
+  @Context(EMULATED_DTYPES="long")
+  def test_emulated_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64, op)
+
+  @given(ht.bool, ht.bool, strat.sampled_from(((operator.add, operator.add), (operator.mul, operator.mul))))
+  def test_bool(self, a, b, op): universal_test(a, b, dtypes.bool, op)
+
+  @given(ht.int32, ht.int32, ht.float32, strat.sampled_from(integer_binary_operations), strat.sampled_from(binary_operations))
+  def test_int32_midcast_float(self, a, b, c, op1, op2): universal_test_midcast(a, b, c, op1, op2, dtypes.int32, dtypes.float32)
+
+  # Metal and (MOCK)CUDA and HIP and NIR behave differently than numpy for overflows
+  skip_overflow = ((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"AMD", "NV", "CUDA"})
+                   or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer))
+  @given(strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
+         strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
+         ht.int32, strat.sampled_from(binary_operations), strat.sampled_from(integer_binary_operations))
+  @unittest.skipIf(Device.DEFAULT == "PYTHON", "TODO: fix cast inf to int32 in PYTHON")
+  @unittest.skip("broken on Mac")
+  def test_float_midcast_int32(self, a, b, c, op1, op2): universal_test_midcast(a, b, c, op1, op2, dtypes.float32, dtypes.int32)
+
+  @unittest.skip("broken. TODO: fix it")
+  @given(ht.float32, strat.sampled_from(dtypes_float+dtypes_int+dtypes_bool))
+  def test_float_cast(self, a, dtype): universal_test_cast(a, dtypes.float32, dtype)
+
+  @unittest.skip("broken. TODO: fix it")
+  @given(ht.int32, strat.sampled_from(dtypes_float+dtypes_int+dtypes_bool))
+  def test_int32_cast(self, a, dtype): universal_test_cast(a, dtypes.int32, dtype)
+
+  @given(strat.floats(width=32, min_value=1.0, max_value=254.0, allow_subnormal=False),
+         strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
+  def test_float_cast_to_unsigned(self, a, float_dtype, unsigned_dtype):
+    if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
+    universal_test_cast(a, float_dtype, unsigned_dtype)
+
+  @unittest.skip("relied on hacks")
+  @given(strat.floats(width=32, min_value=256.0, max_value=65000.0, allow_subnormal=False),
+         strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
+  def test_float_cast_to_unsigned_overflow(self, a, float_dtype, unsigned_dtype):
+    if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
+    universal_test_cast(a, float_dtype, unsigned_dtype)
+
+  @unittest.skip("relied on hacks")
+  @given(strat.floats(width=32, min_value=-65000.0, max_value=-1.0, allow_subnormal=False),
+         strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
+  def test_float_cast_to_unsigned_underflow(self, a, float_dtype, unsigned_dtype):
+    if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
+    universal_test_cast(a, float_dtype, unsigned_dtype)
+
+  @unittest.expectedFailure
+  def test_unsafe_cast_float_to_int_failure(self):
+    val = float(dtypes.int32.max - 1)
+    t1 = Tensor([val], dtype=dtypes.float32).cast(dtypes.int32)
+    t2 = Tensor(val, dtype=dtypes.float32).cast(dtypes.int32)
+    np.testing.assert_equal(t1.item(), t2.item())
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_edgecases.py

@@ -0,0 +1,246 @@
+# end to end tests of tinygrad that you think might be edge cases.
+# using the documentation, write code you think should work.
+# you can compare the outputs to torch or numpy, or just tinygrad assert/raise while doing things that should be valid
+
+# i'm not interested in tests that currently pass, i'm only interested in tests that you think should pass but don't.
+# mark them with @unittest.expectedFailure
+# all the tests in here didn't pass until bugs were fixed
+# get creative! think about things that failed in pytorch or tensorflow for a long time until they were fixed.
+# every test should surface a unique bug. if tinygrad throws an error saying something is not supported, this is probably not a bug.
+# the tests don't have to test the same parts of the code that these current ones test, more diversity is better
+
+# focus on making tinygrad throw runtime errors or assertions for valid things, or find clear numerical mismatches from pytorch
+# confirm any bugs found are valid by doing the same thing in pytorch in the test.
+# for any failing tests, explain in a comment why tinygrad is wrong and what the desired behavior should be.
+# don't worry about running mypy or linters. focus on writing more of these tests and running them to confirm broken behavior.
+# surface level bugs, like issues with empty tensors, input validation, or nans, are not that interesting.
+# focus on bugs that would frustrate real users.
+
+# these are not bugs, these are desired behavior. don't add failing tests for them:
+#   tinygrad only accepts tinygrad dtypes or strings of the tinygrad dtype.
+#   boolean indexing, or anything with unknown output shape of tensor at compile time isn't supported.
+#   invalid indexing in things like gather and one_hot is not an error in tinygrad. nothing that depends on the value is
+#   repeat_interleave doesn't support a tensor as the dim. check tinygrad type signature before claiming something is a bug
+
+import unittest
+import numpy as np
+import torch
+from tinygrad import Tensor, dtypes, nn
+from tinygrad.device import Device
+from tinygrad.helpers import DEV
+from tinygrad.renderer.nir import NIRRenderer
+
+MOCKGPU = DEV.interface.startswith("MOCK")
+
+class TestNaNEdgeCases(unittest.TestCase):
+  # we don't need more of these. it's unclear if torch's behavior is desired here
+
+  @unittest.expectedFailure
+  def test_max_nan(self):
+    # Reductions with NaN should propagate NaN like PyTorch.
+    arr = [1.0, float('nan'), 3.0]
+    torch_out = torch.tensor(arr).max().item()
+    out = Tensor(arr).max().numpy()
+    if np.isnan(torch_out):
+      self.assertTrue(np.isnan(out))
+    else:
+      np.testing.assert_equal(out, torch_out)
+
+  @unittest.skip("passes on webgpu")
+  @unittest.expectedFailure
+  def test_argmax_nan(self):
+    # PyTorch returns the index of the NaN, tinygrad returns the index of the maximum value.
+    arr = [1.0, float('nan'), 3.0]
+    torch_idx = torch.tensor(arr).argmax().item()
+    idx = Tensor(arr).argmax().item()
+    self.assertEqual(idx, torch_idx)
+
+  @unittest.expectedFailure
+  def test_sort_with_nan(self):
+    # Sorting a tensor containing NaN should keep NaN at the end like PyTorch.
+    arr = [1.0, float('nan'), 3.0]
+    torch_vals, torch_idxs = torch.tensor(arr).sort()
+    vals, idxs = Tensor(arr).sort()
+    np.testing.assert_equal(vals.numpy(), torch_vals.numpy())
+    np.testing.assert_equal(idxs.numpy(), torch_idxs.numpy().astype(np.int32))
+
+class TestEmptyTensorEdgeCases(unittest.TestCase):
+  # we don't need more of these
+
+  def test_sort_empty(self):
+    # Sorting an empty tensor works in PyTorch and should return empty
+    # values and indices. tinygrad raises an error instead.
+    torch_vals, torch_idxs = torch.tensor([]).sort()
+    values, indices = Tensor([]).sort()
+    np.testing.assert_equal(values.numpy(), torch_vals.numpy())
+    np.testing.assert_equal(indices.numpy(), torch_idxs.numpy().astype(np.int32))
+
+  @unittest.expectedFailure
+  def test_max_empty(self):
+    # Max on an empty tensor should also raise an error.
+    with self.assertRaises(RuntimeError):
+      torch.tensor([]).max()
+    with self.assertRaises(RuntimeError):
+      Tensor([]).max()
+
+  @unittest.expectedFailure
+  def test_argmax_empty(self):
+    # Argmax on an empty tensor should raise an error like torch does.
+    with self.assertRaises(RuntimeError):
+      torch.tensor([]).argmax()
+    with self.assertRaises(RuntimeError):
+      Tensor([]).argmax()
+
+  def test_masked_select_empty(self):
+    # Masked select on empty tensors should return an empty tensor.
+    torch_out = torch.tensor([], dtype=torch.float32).masked_select(torch.tensor([], dtype=torch.bool))
+    out = Tensor([], dtype=dtypes.float32).masked_select(Tensor([], dtype=dtypes.bool))
+    np.testing.assert_equal(out.numpy(), torch_out.numpy())
+
+class TestDropoutProbabilityEdgeCases(unittest.TestCase):
+  # we don't need more of these
+
+  def test_dropout_rate_one(self):
+    with Tensor.train():
+      out = Tensor.ones(100).dropout(1.0)
+      np.testing.assert_allclose(out.numpy(), np.zeros(100))
+
+  def test_dropout_invalid_prob(self):
+    with self.assertRaises(ValueError):
+      torch.nn.functional.dropout(torch.ones(10), -0.1, True)
+    with self.assertRaises(ValueError):
+      with Tensor.train():
+        Tensor.ones(10).dropout(-0.1)
+
+class TestInputValidation(unittest.TestCase):
+  # we don't need more of these, input validation bugs are not very interesting, many are WONTFIX
+
+  @unittest.expectedFailure
+  def test_repeat_negative(self):
+    # repeating with a negative value should error like PyTorch
+    with self.assertRaises(RuntimeError):
+      torch.tensor([1, 2, 3]).repeat(-1, 2)
+    with self.assertRaises(RuntimeError):
+      Tensor([1, 2, 3]).repeat(-1, 2)
+
+  def test_negative_weight_decay(self):
+    with self.assertRaises(ValueError):
+      torch.optim.AdamW([torch.tensor([1.], requires_grad=True)], lr=0.1, weight_decay=-0.1)
+    with self.assertRaises(ValueError):
+      nn.optim.AdamW([Tensor([1.])], lr=0.1, weight_decay=-0.1)
+
+  def test_negative_lr(self):
+    with self.assertRaises(ValueError):
+      torch.optim.SGD([torch.tensor([1.], requires_grad=True)], lr=-0.1)
+    with self.assertRaises(ValueError):
+      nn.optim.SGD([Tensor([1.])], lr=-0.1)
+
+  def test_negative_momentum(self):
+    with self.assertRaises(ValueError):
+      torch.optim.SGD([torch.tensor([1.], requires_grad=True)], lr=0.1, momentum=-0.1)
+    with self.assertRaises(ValueError):
+      nn.optim.SGD([Tensor([1.])], lr=0.1, momentum=-0.1)
+
+class TestZeroFolding(unittest.TestCase):
+  # we don't need more of these
+
+  # folding rules treat x/x, x//x and x%x as constants even when x can be zero
+  @unittest.expectedFailure
+  def test_divide_by_self_with_zero(self):
+    x = Tensor([0.0, 1.0])
+    torch_out = torch.tensor([0.0, 1.0]) / torch.tensor([0.0, 1.0])
+    out = (x / x).numpy()
+    np.testing.assert_allclose(out, torch_out.numpy(), equal_nan=True)
+
+  @unittest.expectedFailure
+  def test_floordiv_by_self_with_zero(self):
+    x = Tensor([0])
+    with self.assertRaises(RuntimeError):
+      torch.tensor([0]) // torch.tensor([0])
+    with self.assertRaises(RuntimeError):
+      (x // x).numpy()
+
+  @unittest.expectedFailure
+  def test_mod_by_self_with_zero(self):
+    x = Tensor([0])
+    with self.assertRaises(RuntimeError):
+      torch.tensor([0]) % torch.tensor([0])
+    with self.assertRaises(RuntimeError):
+      (x % x).numpy()
+
+class TestAssignIssues(unittest.TestCase):
+  # these are good failures. i'm not sure we need more, but we need to fix these.
+
+  def test_assign_permuted_view_constant(self):
+    # assigning to a permuted view should modify the underlying tensor
+    arr = np.arange(6).reshape(2, 3).astype(np.float32)
+    torch_tensor = torch.tensor(arr)
+    torch_tensor.t().copy_(torch.tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]))
+    t = Tensor(arr).contiguous().realize()
+    t.permute(1, 0).assign(Tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]))
+    np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())
+
+  def test_assign_shrink_view_constant(self):
+    # assigning to a shrunk view should update the base tensor
+    arr = np.arange(9).reshape(3, 3).astype(np.float32)
+    torch_tensor = torch.tensor(arr)
+    torch_tensor[1:3, 1:3] = torch.ones(2, 2)
+    t = Tensor(arr).contiguous().realize()
+    t.shrink(((1, 3), (1, 3))).assign(Tensor.ones(2, 2))
+    np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())
+
+  @unittest.expectedFailure
+  def test_assign_broadcast(self):
+    # broadcasting during assign should behave like PyTorch
+    # NOTE: we don't want implicit dtype casting (int64 -> float32 loses precision), so this fails
+    torch_tensor = torch.zeros(3, 5)
+    torch_tensor[:] = torch.arange(5)
+    t = Tensor.zeros(3, 5)
+    t.assign(Tensor.arange(5))
+    np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())
+
+class TestUOpValidationIssue(unittest.TestCase):
+  # these fail with UOp verification error.
+  # we want more of these with diverse errors!
+
+  @unittest.skipIf(MOCKGPU or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer), "hangs gpuocelot, NIR cannot render")
+  def test_tensor_index_overflow(self):
+    val = Tensor([1])
+    big = val.expand(2**31 + 3)
+    idx = Tensor([0, 2**31 + 2])
+    np.testing.assert_equal(big[idx].numpy(), np.array([1, 1]))
+
+  def test_float_floordiv_scalar(self):
+    (Tensor.arange(4, dtype=dtypes.float32) // 2).realize()
+
+  def test_float_floordiv_tensor(self):
+    (Tensor.arange(4, dtype=dtypes.float32) // Tensor.ones(4, dtype=dtypes.float32)).realize()
+
+class TestEdgeCases(unittest.TestCase):
+  # add tests exposing new and diverse kinds of bugs that might impact real users here
+
+  def test_circular_pad_negative(self):
+    # negative pads with circular mode should wrap like PyTorch
+    arr = np.arange(9).reshape(1, 1, 3, 3).astype(np.float32)
+    torch_out = torch.nn.functional.pad(torch.tensor(arr), (1, -1, 1, -1), mode='circular')
+    out = Tensor(arr).pad((1, -1, 1, -1), mode='circular')
+    np.testing.assert_equal(out.numpy(), torch_out.numpy())
+
+  def test_arange_float_step(self):
+    # float steps should match PyTorch exactly
+    torch_out = torch.arange(0, 2, 0.3).numpy()
+    out = Tensor.arange(0, 2, 0.3).numpy()
+    np.testing.assert_allclose(out, torch_out, atol=1e-7)
+
+  @unittest.skip("this is flaky")
+  @unittest.expectedFailure
+  def test_topk_ties_indices(self):
+    # topk should match PyTorch tie-breaking behavior when values are equal
+    arr = [1.0, 1.0, 1.0, 1.0]
+    _, ti = torch.tensor(arr).topk(2)
+    _, i = Tensor(arr).topk(2)
+    np.testing.assert_equal(i.numpy(), ti.numpy().astype(np.int32))
+
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/backend/test_encodings.py

@@ -0,0 +1,149 @@
+import unittest
+from tinygrad import Device
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.dtype import dtypes
+from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg
+
+def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=op, dtype=dt, src=src, tag=tag)
+
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only on x86")
+class TestEncodingsX86(unittest.TestCase):
+  # NOTE: x86 supports a single displacement as memory address and index without base memory address
+  # these have no use cases so they aren't supported
+  def encode(self, u:UOp): return Device[Device.DEFAULT].renderer.render([u])
+
+  # displacement of 0 isn't emitted
+  def test_base_address(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RDI)
+    # mov edi, dword ptr [rdi]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 3F"))
+
+  # rsp/r12 require a sib byte when used as base memory address
+  def test_rsp_base_address(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RSP), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RSP)
+    # mov esp, dword ptr [rsp]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 24 24"))
+
+  # rbp/r13 require a displacement when used as base memory address
+  def test_rbp_base_address(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RBP), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RBP)
+    # mov ebp, dword ptr [rbp + 0]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 6D 00"))
+
+  # test [base + index*scale]
+  def test_base_index_address(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, RDX), imm(dtypes.int8, 0)), RAX)
+    # mov eax, dword ptr [rax + rdx*4]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 04 90"))
+
+  # rsp as index means no index
+  def test_rsp_index_address(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, RSP), imm(dtypes.int8, 0)), RAX)
+    # mov eax, dword ptr [rax]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 00"))
+
+  # however r12 is a valid index
+  def test_r12_index_address(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, GPR[12]), imm(dtypes.int8, 0)), RAX)
+    # mov eax, dword ptr [rax + r12*4]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("42 8B 04 A0"))
+
+  # test [base + index*scale + 8bit disp]
+  def test_complex_address_8bit_disp(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)), RDI)
+    # mov edi, dword ptr [rdi + rsi*4 + 0xa]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 7C B7 0A"))
+
+  # test [base + index*scale + 32bit disp]
+  def test_complex_address_32bit_disp(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10000)), RDI)
+    # mov edi, dword ptr [rdi + rsi*4 + 0x2710]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B BC B7 10 27 00 00"))
+
+  # 8bit variants of legacy instructions subtract 1 from opcode
+  def test_8bit_legacy_encoding(self):
+    cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(dtypes.int8, RDX),), RAX)
+    # movsx eax, dl
+    self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("0F BE C2"))
+
+  # accessing lower 8 bits of rsp, rbp, rsi, rdi requires rex prefix
+  def test_lower_8bits_reg(self):
+    cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(dtypes.int8, RDI),), RAX)
+    # movsx eax, dil
+    self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("40 0F BE C7"))
+
+  # test 16 bit variant of legacy instruction
+  def test_16bit_legacy_encoding(self):
+    cast = ins(X86Ops.MOVSX, dtypes.int16, (def_reg(dtypes.int8, RDX),), RAX)
+    # movsx ax, dl
+    self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("66 0F BE C2"))
+
+  # test 64 bit variant of legacy instruction
+  def test_64bit_legacy_encoding(self):
+    cast = ins(X86Ops.MOVSX, dtypes.int64, (def_reg(dtypes.int8, RDX),), RAX)
+    # movsx rax, dl
+    self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("48 0F BE C2"))
+
+  # test compact vex encoding
+  def test_compact_vex_encoding(self):
+    xmm0, xmm1 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1])
+    add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm1), XMM[0])
+    # vaddss xmm0, xmm0, xmm1
+    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FA 58 C1"))
+
+  # test long vex encoding
+  def test_long_vex_encoding(self):
+    xmm0, xmm8 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[8])
+    add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm8), XMM[0])
+    # vaddss xmm0, xmm0, xmm8
+    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C4 C1 7A 58 C0"))
+
+  # test ymm encoding
+  def test_ymm_encoding(self):
+    xmm0, xmm1 = def_reg(dtypes.float32.vec(8), XMM[0]), def_reg(dtypes.float32.vec(8), XMM[1])
+    add = ins(X86Ops.VADDPS, dtypes.float32.vec(8), (xmm0, xmm1), XMM[0])
+    # vaddps ymm0, ymm0, ymm1
+    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FC 58 C1"))
+
+  # test encoding where register is in the immediate field
+  def test_reg_in_imm_field(self):
+    xmm0, xmm1, xmm2 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1]), def_reg(dtypes.float32, XMM[2])
+    blend = ins(X86Ops.VBLENDVPS, dtypes.float32, (xmm0, xmm1, xmm2), XMM[0])
+    # vblendvps xmm0, xmm0, xmm1, xmm2
+    self.assertEqual(bytes.fromhex(self.encode(blend)), bytes.fromhex("C4 E3 79 4A C1 20"))
+
+  # when writting to mem the uop takes the store form where dtype is void and there's no definition
+  def test_write_mem(self):
+    base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)
+    xmm0 = def_reg(dtypes.float32, XMM[0])
+    extr = ins(X86Ops.VPEXTRD, dtypes.void, (base, index, disp, xmm0, imm(dtypes.uint8, 0)))
+    # vpextrd dword ptr [rdi + rsi*4 + 0xa], xmm0, 0
+    self.assertEqual(bytes.fromhex(self.encode(extr)), bytes.fromhex("C4 E3 79 16 44 B7 0A 00"))
+
+  # test two address instruction with fused load works
+  def test_two_address_load(self):
+    base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)
+    cmove = ins(X86Ops.CMOVE, dtypes.int32,  (base, index, disp), RAX)
+    # cmove eax, dword ptr [rdi + rsi*4 + 0xa]
+    self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 44 B7 0A"))
+
+  # test instruction where displacement and imm have the same value
+  def test_disp_imm_same_value(self):
+    base, index, disp = def_reg(dtypes.int8.ptr(), RDI), def_reg(dtypes.int8, RSI), imm(dtypes.int8, 10)
+    mov = ins(X86Ops.MOVi, dtypes.void, (base, index, disp, disp))
+    # mov byte ptr [rdi + rsi + 0xa], 0xa
+    self.assertEqual(bytes.fromhex(self.encode(mov)), bytes.fromhex("40 C6 44 37 0A 0A"))
+
+    base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10)
+    imul = ins(X86Ops.IMULi, dtypes.int32, (base, index, disp) + (imm(dtypes.int32, 10),), RDI)
+    # imul edi, dword ptr [rdi + rsi*4 + 0xa], 0xa
+    self.assertEqual(bytes.fromhex(self.encode(imul)), bytes.fromhex("69 BC B7 0A 00 00 00 0A 00 00 00"))
+
+  # cmoves have the cmp as the last src even though it is not explicitly used, the cmp doesn't define a reg and is ignored in the encoding
+  def test_cmove_ignore_cmp(self):
+    cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=X86Ops.CMP)), RDX)
+    # cmove edx, eax
+    self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0"))
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/backend/test_graph.py

@@ -0,0 +1,319 @@
+import numpy as np
+import functools, unittest, ctypes
+
+from tinygrad.device import Device, Buffer
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import Context, from_mv
+from tinygrad.dtype import dtypes
+from tinygrad.engine.jit import MultiGraphRunner
+from tinygrad.engine.realize import run_linear, compile_linear
+from tinygrad.uop.ops import UOp, Ops, buffers
+
+from test.helpers import needs_second_gpu
+
+np.random.seed(1337)
+Tensor.manual_seed(1337)
+BUF_SIZE = 4096
+RUN_CNT = 5
+
+# cache AST by (device, num_inputs)
+cached_asts: dict[tuple[str, int], UOp] = {}
+def get_ast(device:str, num_inputs:int) -> UOp:
+  if (device, num_inputs) not in cached_asts:
+    with Context(DEBUG=0):
+      fst = [Tensor.randn(BUF_SIZE, dtype=dtypes.int).realize() for _ in range(num_inputs)]
+      s = fst[0]
+      for i in range(1, num_inputs): s = s.bitwise_xor(fst[i])
+      cached_asts[(device, num_inputs)] = s.schedule_linear().src[-1].src[0]
+  return cached_asts[(device, num_inputs)]
+
+def make_buffer(device, size=BUF_SIZE, fill=False):
+  buf = Buffer(device, size, dtypes.int).ensure_allocated()
+  if fill:
+    with Context(DEBUG=0):
+      buf.copyin(Tensor(np.random.randint(-10000, 10000, size=size, dtype=np.int32)).realize().uop.base.realized.as_memoryview())
+  return buf
+
+def make_view(base, offset_elems, size_elems):
+  return Buffer(base.device, size_elems, base.dtype, base=base, offset=offset_elems * base.dtype.itemsize).ensure_allocated()
+
+def get_buf_uop(buf:Buffer, cache:dict[Buffer,UOp]) -> UOp:
+  if buf not in cache:
+    cache[buf] = u = UOp.new_buffer(buf.device, buf.size, buf.dtype)
+    buffers[u] = buf
+  return cache[buf]
+
+def make_graph(graph_cls, calls:list[UOp]):
+  linear = compile_linear(UOp(Ops.LINEAR, src=tuple(calls)))
+  cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(linear,), arg="graph")
+  return graph_cls(cf, [])
+
+def run_schedule(calls:list[UOp]):
+  run_linear(UOp(Ops.LINEAR, src=tuple(calls)))
+
+def zero_bufs(bufs):
+  for b in bufs:
+    mv = memoryview(bytearray(b.nbytes))
+    ctypes.memset(from_mv(mv), 0, len(mv))
+    b.copyin(mv)
+
+@unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
+class TestGraph(unittest.TestCase):
+  def skip_if_no_offset(self):
+    if not hasattr(Device[Device.DEFAULT].allocator, "_offset"): self.skipTest("device does not support _offset")
+
+  def skip_if_not_multigraph(self):
+    graph = g.func if isinstance(g:=(d:=Device[Device.DEFAULT]).graph, functools.partial) else g
+    if not issubclass(graph, MultiGraphRunner): self.skipTest("graph is not supported (not MultiGraphRunner)")
+    if not hasattr(d.allocator, '_transfer') or not d.allocator.supports_transfer: self.skipTest("device is not supported (no transfers)")
+
+  def test_order_2_writes_to_same_buf(self):
+    d0 = Device.DEFAULT
+    b = [make_buffer(d0, fill=True) for _ in range(5)]
+    c: dict[Buffer,UOp] = {}
+
+    calls = [
+      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c), metadata=()),
+    ]
+
+    zero_bufs([b[0]])
+    run_schedule(calls)
+    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
+
+    for _ in range(RUN_CNT):
+      zero_bufs([b[0]])
+      make_graph(Device[d0].graph, calls)([], {})
+      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
+
+  def test_order_read_write_same_buf(self):
+    d0 = Device.DEFAULT
+    b = [make_buffer(d0, fill=True) for _ in range(5)]
+    c: dict[Buffer,UOp] = {}
+
+    calls = [
+      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c), metadata=()),
+    ]
+
+    zero_bufs([b[0], b[1]])
+    run_schedule(calls)
+    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
+
+    for _ in range(RUN_CNT):
+      zero_bufs([b[0], b[1]])
+      make_graph(Device[d0].graph, calls)([], {})
+      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
+
+  def test_order_write_read_same_buf(self):
+    d0 = Device.DEFAULT
+    b = [make_buffer(d0, fill=True) for _ in range(5)]
+    c: dict[Buffer,UOp] = {}
+
+    calls = [
+      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c), get_buf_uop(b[4],c), metadata=()),
+    ]
+
+    zero_bufs([b[0], b[1]])
+    run_schedule(calls)
+    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
+
+    for _ in range(RUN_CNT):
+      zero_bufs([b[0], b[1]])
+      make_graph(Device[d0].graph, calls)([], {})
+      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
+
+  def test_order_copy_writed(self):
+    self.skip_if_not_multigraph()
+    d0 = Device.DEFAULT
+    b = [make_buffer(d0, fill=True) for _ in range(4)]
+    c: dict[Buffer,UOp] = {}
+
+    calls = [
+      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
+      UOp(Ops.COPY).call(get_buf_uop(b[3],c), get_buf_uop(b[0],c), metadata=()),
+    ]
+
+    zero_bufs([b[0], b[3]])
+    run_schedule(calls)
+    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
+
+    for _ in range(RUN_CNT):
+      zero_bufs([b[0], b[3]])
+      make_graph(Device[d0].graph, calls)([], {})
+      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
+
+  def test_order_copy_then_read(self):
+    self.skip_if_not_multigraph()
+    d0 = Device.DEFAULT
+    b = [make_buffer(d0, fill=True) for _ in range(4)]
+    c: dict[Buffer,UOp] = {}
+
+    calls = [
+      UOp(Ops.COPY).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
+    ]
+
+    zero_bufs([b[1], b[3]])
+    run_schedule(calls)
+    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
+
+    for _ in range(RUN_CNT):
+      zero_bufs([b[1], b[3]])
+      make_graph(Device[d0].graph, calls)([], {})
+      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
+
+  def test_read_write_several_graphs(self):
+    d0 = Device.DEFAULT
+    b = [make_buffer(d0, fill=True) for _ in range(8)]
+    c: dict[Buffer,UOp] = {}
+
+    calls1 = [get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=())]
+    calls2 = [get_ast(d0, 2).call(get_buf_uop(b[4],c), get_buf_uop(b[1],c), get_buf_uop(b[3],c), metadata=())]
+    calls3 = [get_ast(d0, 2).call(get_buf_uop(b[5],c), get_buf_uop(b[4],c), get_buf_uop(b[2],c), metadata=())]
+
+    out = [b[3], b[4], b[5]]
+    zero_bufs(out)
+    run_schedule(calls1 + calls2 + calls3)
+    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
+
+    for _ in range(RUN_CNT):
+      zero_bufs(out)
+      make_graph(Device[d0].graph, calls1)([], {})
+      make_graph(Device[d0].graph, calls2)([], {})
+      make_graph(Device[d0].graph, calls3)([], {})
+      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
+
+  @needs_second_gpu
+  def test_copies_2_devs(self):
+    self.skip_if_not_multigraph()
+    d0, d1 = Device.DEFAULT, f"{Device.DEFAULT}:1"
+    b0 = [make_buffer(d0, fill=True) for _ in range(3)]
+    b1 = [make_buffer(d1, fill=True)]
+    c: dict[Buffer,UOp] = {}
+
+    calls = [
+      UOp(Ops.COPY).call(get_buf_uop(b1[0],c), get_buf_uop(b0[0],c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(b0[2],c), get_buf_uop(b0[0],c), get_buf_uop(b0[1],c), metadata=()),
+    ]
+
+    out = [b1[0], b0[2]]
+    zero_bufs(out)
+    run_schedule(calls)
+    expected = {buf: np.frombuffer(buf.as_memoryview(), np.int32).copy() for buf in b0 + b1}
+
+    for _ in range(RUN_CNT):
+      zero_bufs(out)
+      make_graph(Device[d0].graph, calls)([], {})
+      for buf in b0 + b1: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
+
+  def test_graph_offset_bufs(self):
+    self.skip_if_not_multigraph()
+    d0 = Device.DEFAULT
+    if not hasattr(Device[d0].allocator, "_offset"): self.skipTest("device does not support _offset")
+
+    b0 = make_buffer(d0, fill=True)
+    b1 = make_view(b0, 0, b0.size)
+    b2 = make_view(b0, 0, b0.size)
+    c: dict[Buffer,UOp] = {}
+
+    calls = [
+      UOp(Ops.COPY).call(get_buf_uop(b0,c), get_buf_uop(b2,c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(b1,c), get_buf_uop(b0,c), get_buf_uop(b2,c), metadata=()),
+    ]
+
+    zero_bufs([b0])
+    run_schedule(calls)
+    expected = np.frombuffer(b0.as_memoryview(), np.int32).copy()
+
+    for _ in range(RUN_CNT):
+      zero_bufs([b0])
+      make_graph(Device[d0].graph, calls)([], {})
+      np.testing.assert_equal(expected, np.frombuffer(b0.as_memoryview(), np.int32))
+
+  def test_partial_write_preserves_write_dep(self):
+    self.skip_if_not_multigraph()
+    self.skip_if_no_offset()
+    d0 = Device.DEFAULT
+
+    base = make_buffer(d0, BUF_SIZE * 2, fill=True)
+    copy_src_full = make_buffer(d0, BUF_SIZE * 2, fill=True)
+    copy_src_lo = make_buffer(d0, fill=True)
+    v_lo, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE)
+    a, out = make_buffer(d0, fill=True), make_buffer(d0, fill=True)
+    c: dict[Buffer,UOp] = {}
+
+    calls = [
+      UOp(Ops.COPY).call(get_buf_uop(base,c), get_buf_uop(copy_src_full,c), metadata=()),
+      UOp(Ops.COPY).call(get_buf_uop(v_lo,c), get_buf_uop(copy_src_lo,c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(out,c), get_buf_uop(v_hi,c), get_buf_uop(a,c), metadata=()),
+    ]
+
+    zero_bufs([base, out])
+    run_schedule(calls)
+    expected = {base: np.frombuffer(base.as_memoryview(), np.int32).copy(), out: np.frombuffer(out.as_memoryview(), np.int32).copy()}
+
+    for _ in range(RUN_CNT):
+      zero_bufs([base, out])
+      make_graph(Device[d0].graph, calls)([], {})
+      for buf in [base, out]: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
+
+  def test_partial_write_preserves_read_dep(self):
+    self.skip_if_not_multigraph()
+    self.skip_if_no_offset()
+    d0 = Device.DEFAULT
+
+    base = make_buffer(d0, BUF_SIZE * 2, fill=True)
+    copy_dst = make_buffer(d0, BUF_SIZE * 2, fill=True)
+    copy_src_lo = make_buffer(d0, fill=True)
+    v_lo, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE)
+    a, b = make_buffer(d0, fill=True), make_buffer(d0, fill=True)
+    c: dict[Buffer,UOp] = {}
+
+    calls = [
+      UOp(Ops.COPY).call(get_buf_uop(copy_dst,c), get_buf_uop(base,c), metadata=()),
+      UOp(Ops.COPY).call(get_buf_uop(v_lo,c), get_buf_uop(copy_src_lo,c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(v_hi,c), get_buf_uop(a,c), get_buf_uop(b,c), metadata=()),
+    ]
+
+    zero_bufs([copy_dst, base])
+    run_schedule(calls)
+    expected = {copy_dst: np.frombuffer(copy_dst.as_memoryview(), np.int32).copy(), base: np.frombuffer(base.as_memoryview(), np.int32).copy()}
+
+    for _ in range(RUN_CNT):
+      zero_bufs([copy_dst, base])
+      make_graph(Device[d0].graph, calls)([], {})
+      for buf in [copy_dst, base]: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
+
+  def test_middle_write_splits_write_dep(self):
+    self.skip_if_not_multigraph()
+    self.skip_if_no_offset()
+    d0 = Device.DEFAULT
+
+    base = make_buffer(d0, BUF_SIZE * 3, fill=True)
+    copy_src_full = make_buffer(d0, BUF_SIZE * 3, fill=True)
+    copy_src_mid = make_buffer(d0, fill=True)
+    v_lo, v_mid, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE), make_view(base, BUF_SIZE * 2, BUF_SIZE)
+    a, out1, out2 = make_buffer(d0, fill=True), make_buffer(d0, fill=True), make_buffer(d0, fill=True)
+    c: dict[Buffer,UOp] = {}
+
+    calls = [
+      UOp(Ops.COPY).call(get_buf_uop(base,c), get_buf_uop(copy_src_full,c), metadata=()),
+      UOp(Ops.COPY).call(get_buf_uop(v_mid,c), get_buf_uop(copy_src_mid,c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(out1,c), get_buf_uop(v_lo,c), get_buf_uop(a,c), metadata=()),
+      get_ast(d0, 2).call(get_buf_uop(out2,c), get_buf_uop(v_hi,c), get_buf_uop(a,c), metadata=()),
+    ]
+
+    outs = [base, out1, out2]
+    zero_bufs(outs)
+    run_schedule(calls)
+    expected = {buf: np.frombuffer(buf.as_memoryview(), np.int32).copy() for buf in outs}
+
+    for _ in range(RUN_CNT):
+      zero_bufs(outs)
+      make_graph(Device[d0].graph, calls)([], {})
+      for buf in outs: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_interop.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+import unittest, os
+import torch
+import numpy as np
+
+from tinygrad.helpers import DEV
+from tinygrad.tensor import Tensor
+from tinygrad.device import Device
+from tinygrad.dtype import _from_torch_dtype, _to_torch_dtype
+
+MOCKGPU = DEV.interface.startswith("MOCK")
+
+@unittest.skipIf(Device.DEFAULT not in ["METAL", "CUDA"] or MOCKGPU, f"no support on {Device.DEFAULT}")
+class TestInterop(unittest.TestCase):
+  def setUp(self):
+    if Device.DEFAULT == "CUDA": self.torch_device = "cuda"
+    elif Device.DEFAULT == "METAL": self.torch_device = "mps"
+
+  def test_torch_interop(self):
+    inp = torch.rand(2, 2, 3, device=torch.device(self.torch_device))
+
+    if self.torch_device == "mps": torch.mps.synchronize()
+    else: torch.cuda.synchronize()
+
+    tg_data = Tensor.from_blob(inp.data_ptr(), inp.shape, dtype=_from_torch_dtype(inp.dtype))
+
+    tg_out = tg_data[:, :, 0] * 0.2989 + tg_data[:, :, 1] * 0.5870 + tg_data[:, :, 2] * 0.1140
+    tg_res = tg_out.numpy()
+
+    if self.torch_device == "mps" and os.getenv("CI", "") != "":
+      # MPS backend out of memory: https://discuss.pytorch.org/t/mps-back-end-out-of-memory-on-github-action/189773
+      # Calculate expected value on cpu.
+      inp = inp.cpu()
+    torch_out = inp[:, :, 0] * 0.2989 + inp[:, :, 1] * 0.5870 + inp[:, :, 2] * 0.1140
+
+    np.testing.assert_allclose(tg_res, torch_out.cpu().numpy(), atol=1e-5, rtol=1e-5)
+
+  def test_torch_interop_write(self):
+    tg_data = Tensor.randn((4, 4), device=Device.DEFAULT)
+
+    out = torch.empty(4, 4, device=torch.device(self.torch_device), dtype=_to_torch_dtype(tg_data.dtype))
+    tg_out = Tensor.from_blob(out.data_ptr(), out.shape, dtype=_from_torch_dtype(out.dtype))
+
+    tg_out.assign(tg_data).realize()
+    Device[Device.DEFAULT].synchronize()
+
+    torch_out_np = out.cpu().numpy()
+
+    np.testing.assert_allclose(tg_data.numpy(), torch_out_np, atol=1e-5, rtol=1e-5)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_isel.py

@@ -0,0 +1,147 @@
+import unittest
+from typing import cast
+from tinygrad import Device
+from tinygrad.uop import Ops
+from tinygrad.uop.ops import UOp, dtypes, graph_rewrite
+from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
+from tinygrad.renderer.isa import IselContext
+
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
+class TestIselX86(unittest.TestCase):
+  def isel_rewrite(self, x:UOp):
+    return graph_rewrite(x, cast(X86Renderer, Device[Device.DEFAULT].renderer).isel_matcher, IselContext(x), bottom_up=True)
+
+  def _check_op(self, dt_op, expr):
+    nargs = expr.__code__.co_argcount
+    for dt,op in dt_op:
+      with self.subTest(dtype=dt):
+        v = [UOp.variable(str(i), 0, 0, dt) for i in range(nargs)]
+        n = self.isel_rewrite(expr(*v))
+        self.assertIs(n.arg, op)
+
+  def test_cmove(self):
+    a = UOp.variable("a", 0, 0, dtypes.int32)
+    b = UOp.variable("b", 0, 0, dtypes.int32)
+    c = (a < b).where(a, b)
+    d = (a != b).where(a, b)
+    f = c + d
+    n = self.isel_rewrite(f)
+    self.assertTrue(n.src[0].arg is X86Ops.CMOVL and n.src[1].arg is X86Ops.CMOVNE)
+    # both comparisons become the same instruction
+    self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg is X86Ops.CMP)
+
+  def test_vmax(self):
+    dt_op = [(dtypes.float32, X86Ops.VMAXSS), (dtypes.float64, X86Ops.VMAXSD),
+             (dtypes.float32.vec(4), X86Ops.VMAXPS), (dtypes.float64.vec(4), X86Ops.VMAXPD)]
+    self._check_op(dt_op, lambda a,b: (a < b).where(b, a))
+
+  def test_vmin(self):
+    dt_op = [(dtypes.float32, X86Ops.VMINSS), (dtypes.float64, X86Ops.VMINSD),
+             (dtypes.float32.vec(4), X86Ops.VMINPS), (dtypes.float64.vec(4), X86Ops.VMINPD)]
+    self._check_op(dt_op, lambda a,b: (a < b).where(a, b))
+
+  def test_vfmadd(self):
+    dt_op = [(dtypes.float32, X86Ops.VFMADD213SS), (dtypes.float64, X86Ops.VFMADD213SD),
+             (dtypes.float32.vec(4), X86Ops.VFMADD213PS), (dtypes.float64.vec(4), X86Ops.VFMADD213PD)]
+    self._check_op(dt_op, lambda a,b,c: a * b + c)
+
+  # don't use fmadd if op being fused (mul) is used multiple times
+  def test_no_vfmadd(self):
+    dt_op = [(dtypes.float32, X86Ops.VADDSS), (dtypes.float64, X86Ops.VADDSD),
+             (dtypes.float32.vec(4), X86Ops.VADDPS), (dtypes.float64.vec(4), X86Ops.VADDPD)]
+    self._check_op(dt_op, lambda a,b: a * b + a * b)
+
+  def test_vpbroadcast(self):
+    a = UOp.variable("a", 0, 0, dtypes.int32)
+    n = self.isel_rewrite(a.broadcast(4))
+    # need to move src from gpr to xmm before broadcasting
+    self.assertTrue(n.arg is X86Ops.VPBROADCASTD and n.src[0].arg is X86Ops.VMOVD)
+    # if we can fuse a load we can skip the move and access memory directly
+    load = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load()
+    n = self.isel_rewrite(load.broadcast(4))
+    self.assertTrue(n.arg is X86Ops.VPBROADCASTD and len(n.src) == 3)
+
+  def test_vbroadcastss(self):
+    a = UOp.variable("a", 0, 0, dtypes.float32)
+    valid = [UOp.vectorize(a, a, a, a), UOp.vectorize(a, a, a, a, a, a, a, a)]
+    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VBROADCASTSS)
+
+  def test_vshufps(self):
+    a = UOp.variable("a", 0, 0, dtypes.float32.vec(8))
+    b = UOp.variable("b", 0, 0, dtypes.float32.vec(8))
+    c = UOp.variable("c", 0, 0, dtypes.float32)
+    d = UOp.variable("d", 0, 0, dtypes.float32)
+
+    valid = [UOp.vectorize(c, c, d, d),
+             UOp.vectorize(a.gep(0), a.gep(1), c, c),
+             UOp.vectorize(a.gep(0), a.gep(1), b.gep(2), b.gep(3)),
+             UOp.vectorize(a.gep(1), a.gep(2), a.gep(3), a.gep(0)),
+             UOp.vectorize(a.gep(3), a.gep(2), a.gep(1), a.gep(0), a.gep(7), a.gep(6), a.gep(5), a.gep(4)),
+             UOp.vectorize(a.gep(0), a.gep(0), b.gep(1), b.gep(1), a.gep(4), a.gep(4), b.gep(5), b.gep(5))]
+    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPS)
+
+    invalid = [UOp.vectorize(a.gep(0), a.gep(1), b.gep(4), b.gep(5)),
+               UOp.vectorize(a.gep(0), a.gep(5), b.gep(2), b.gep(3)),
+               UOp.vectorize(a.gep(0), a.gep(0), a.gep(0), a.gep(0), a.gep(4), a.gep(4), a.gep(4), a.gep(5)),
+               UOp.vectorize(a.gep(0), a.gep(0), b.gep(0), b.gep(0), a.gep(4), a.gep(4), b.gep(4), a.gep(4))]
+    for shuf in invalid: self.assertIsNot(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPS)
+
+  def test_vshufpd(self):
+    a = UOp.variable("a", 0, 0, dtypes.float64.vec(4))
+    b = UOp.variable("b", 0, 0, dtypes.float64.vec(4))
+    c = UOp.variable("c", 0, 0, dtypes.float64)
+    d = UOp.variable("d", 0, 0, dtypes.float64)
+
+    valid = [UOp.vectorize(c, d),
+             UOp.vectorize(a.gep(0), c),
+             UOp.vectorize(a.gep(1), b.gep(1)),
+             UOp.vectorize(a.gep(0), b.gep(1), a.gep(2), b.gep(3)),
+             UOp.vectorize(a.gep(1), a.gep(1), a.gep(3), a.gep(3))]
+    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPD)
+
+    invalid = [UOp.vectorize(c, c, c, c),
+               UOp.vectorize(a.gep(0), a.gep(1), b.gep(2), b.gep(3)),
+               UOp.vectorize(a.gep(2), b.gep(3), a.gep(2), b.gep(3)),
+               UOp.vectorize(a.gep(0), b.gep(1), a.gep(0), b.gep(1))]
+    for shuf in invalid: self.assertIsNot(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPD)
+
+  def test_vinsertps(self):
+    a = UOp.variable("a", 0, 0, dtypes.float32.vec(4))
+    b = UOp.variable("b", 0, 0, dtypes.float32.vec(4))
+    c = UOp.variable("c", 0, 0, dtypes.float32.vec(4))
+    d = UOp.variable("e", 0, 0, dtypes.float32)
+    # moving 0th element to position 0 does nothing so only 1 vinsertps is generated
+    n = self.isel_rewrite(UOp.vectorize(a.gep(0), d))
+    self.assertIs(n.arg, X86Ops.VINSERTPS)
+    self.assertIsNot(n.src[0].arg, X86Ops.VINSERTPS)
+
+    valid = [UOp.vectorize(a.gep(0), b.gep(1), a.gep(2), b.gep(3)),
+             UOp.vectorize(a.gep(3), b.gep(2), c.gep(1), d)]
+    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VINSERTPS)
+
+  # complex address is [base + index*scale + displacement]
+  def test_complex_address(self):
+    a = UOp.variable("a", 0, 0, dtypes.int32)
+    load = UOp.param(0, dtypes.int32.ptr()).index(a + 1, ptr=True).load()
+    n = self.isel_rewrite(load)
+    # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
+    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].arg == 4)
+
+  def test_fold_load(self):
+    load1 = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load()
+    load2 = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 1), ptr=True).load()
+    n = self.isel_rewrite(load1 + load2)
+    self.assertTrue(len(n.src) == 4)
+
+  # don't fold when used multiple times
+  def test_dont_fold_load(self):
+    load = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load()
+    # used by multiple users
+    n = self.isel_rewrite(load + 1 + load)
+    self.assertTrue(len(n.src) == 2)
+    # used mutiple times by same user
+    n = self.isel_rewrite(load * load)
+    self.assertTrue(len(n.src) == 2)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/backend/test_jit.py

@@ -0,0 +1,892 @@
+#!/usr/bin/env python
+import unittest
+import numpy as np
+
+from hypothesis import given, settings, strategies as strat
+from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu
+from tinygrad import Variable
+from tinygrad.tensor import Tensor
+from tinygrad.engine.jit import TinyJit, JitError, graph_class
+from tinygrad.device import Device
+from tinygrad.helpers import Context, JIT, DEV, GlobalCounters
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import Ops, UOp
+from extra.models.unet import ResBlock
+from tinygrad.renderer.isa.x86 import X86Renderer
+
+def _simple_test(add, extract=lambda x: x, N=10):
+  for _ in range(5):
+    a = Tensor.randn(N, N)
+    b = Tensor.randn(N, N)
+    c = add(a, b)
+    np.testing.assert_allclose(extract(c).numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
+  assert_jit_cache_len(add, 1)
+
+class TestJit(unittest.TestCase):
+
+  @settings(deadline=2e4)
+  @unittest.skipUnless(Device.DEFAULT in ["CPU"], f"no support on {Device.DEFAULT}")
+  @given(strat.sampled_from([Tensor.exp2, Tensor.log2, Tensor.sin]))
+  def test_approx_jit_timeout(self, op):
+    with Context(TRANSCENDENTAL=2):
+      model = [ResBlock(16, 24, 16) for _ in range(4)]
+      @TinyJit
+      def fw_approx(t, t2):
+        for l in model: t = l(t, t2)
+        return op(t).realize()
+      fw_approx(Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24))
+
+  def test_simple_jit(self):
+    @TinyJit
+    def add(a, b): return (a+b).realize()
+    _simple_test(add)
+
+  @unittest.skipUnless(Device.DEFAULT == "CPU", "core_id is a CPU runtimevar")
+  def test_hcq_core_id_runtimevar_merge(self):
+    N = 262144
+    @TinyJit
+    def f(x, st):
+      y = (x + 1).contiguous().realize()
+      z = x.shrink(((st, st + N),)).contiguous().realize()
+      return y, z
+    x = Tensor.arange(2*N).clone().realize()
+    for _ in range(3): y, z = f(x, Variable("a", 0, N).bind(0))
+    self.assertEqual(y.shape, (2*N,))
+    self.assertEqual(z.shape, (N,))
+
+  def test_jitbeam_triggers_beam(self):
+    from unittest.mock import patch
+    from tinygrad.helpers import getenv as _getenv
+    @TinyJit
+    def add(a, b): return (a+b).realize()
+    a, b = Tensor.ones(10, 10).contiguous().realize(), Tensor.ones(10, 10).contiguous().realize()
+    with patch("tinygrad.codegen.opt.search.beam_search", wraps=lambda k,*a,**kw: k) as mock_beam:
+      add(a, b)
+      assert mock_beam.call_count == 0
+      with patch("tinygrad.engine.jit.getenv", side_effect=lambda k, d=0: 1 if k == "JITBEAM" else _getenv(k, d)): add(a, b)
+      assert mock_beam.call_count == 1
+
+  def test_simple_jit_reset(self):
+    @TinyJit
+    def add(a, b): return (a+b).realize()
+    _simple_test(add)
+    add.reset()
+    _simple_test(add, N=20)
+
+  def test_simple_jit_norealize(self):
+    @TinyJit
+    def add(a, b): return (a+b)
+    _simple_test(add)
+
+  def test_simple_jit_norealize_list(self):
+    @TinyJit
+    def add(a, b): return [a+b]
+    _simple_test(add, extract=lambda x: x[0])
+
+  def test_simple_jit_norealize_dict(self):
+    @TinyJit
+    def add(a, b): return {"billy": a+b}
+    _simple_test(add, extract=lambda x: x["billy"])
+
+  def test_jit_input_view(self):
+    @TinyJit
+    def f(x): return (x[2:5].contiguous() + 1).realize()
+    for i in range(5):
+      x = (Tensor.arange(10).float() + i * 10).clone().realize()
+      np.testing.assert_allclose(f(x).numpy(), x.numpy()[2:5] + 1)
+
+  def test_jit_multiple_outputs(self):
+    @TinyJit
+    def f(a, b): return (a+b).realize(), (a-b).realize(), (a*b).realize()
+    for _ in range(5):
+      a = Tensor.randn(10, 10)
+      b = Tensor.randn(10, 10)
+      c, d, e = f(a, b)
+      np.testing.assert_allclose(c.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
+      np.testing.assert_allclose(d.numpy(), a.numpy()-b.numpy(), atol=1e-4, rtol=1e-5)
+      np.testing.assert_allclose(e.numpy(), a.numpy()*b.numpy(), atol=1e-4, rtol=1e-5)
+    assert_jit_cache_len(f, 3)
+
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "estimates are wrong for x86")
+  def test_global_counters_jit(self):
+    @TinyJit
+    def f(a, b):
+      c = (a + b).realize()
+      d = (c * 2).realize()
+      return (d - a).realize()
+    a, b = Tensor.randn(64, 64).realize(), Tensor.randn(64, 64).realize()
+    for _ in range(4):
+      GlobalCounters.reset()
+      f(a, b)
+      Device[a.device].synchronize()
+      self.assertGreater(GlobalCounters.global_mem, 0)
+      self.assertGreater(GlobalCounters.global_ops, 0)
+
+  def test_nothing_jitted(self):
+    @TinyJit
+    def add(a, b): return None
+    with self.assertRaises(JitError):
+      for _ in range(5):
+        a = Tensor.randn(10, 10)
+        b = Tensor.randn(10, 10)
+        add(a, b)
+
+  def test_jit_zero_does_not_jit(self):
+    @TinyJit
+    def add(a, b): return (a+b).realize()
+    with Context(JIT=0):
+      for i in range(5):
+        a = Tensor([i])
+        b = Tensor([i])
+        c = add(a, b)
+        np.testing.assert_allclose(c.numpy(), 2*i)
+      assert_jit_cache_len(add, 0)
+
+  def test_jit_not_capturing(self):
+    @TinyJit
+    def add(a, b):
+      Tensor.zeros(4, 4).contiguous().realize()  # no-op kernel is captured
+      return (a+b).realize()
+    for i in range(5):
+      a = Tensor([i])
+      b = Tensor([i])
+      c = add(a, b)
+      np.testing.assert_allclose(c.numpy(), 2*i)
+    assert_jit_cache_len(add, 2)
+
+    @TinyJit
+    def add2(a, b):
+      with Context(CAPTURING=0):  # not captured
+        Tensor.zeros(4, 4).contiguous().realize()
+      return (a+b).realize()
+    for i in range(5):
+      a = Tensor([i])
+      b = Tensor([i])
+      c = add2(a, b)
+      np.testing.assert_allclose(c.numpy(), 2*i)
+    assert_jit_cache_len(add2, 1)
+
+  def test_jit_shape_mismatch(self):
+    @TinyJit
+    def add(a, b): return (a+b).realize()
+    for _ in range(5):
+      a = Tensor.randn(10, 10)
+      b = Tensor.randn(10, 10)
+      add(a, b)
+    bad = Tensor.randn(20, 20)
+    with self.assertRaises(JitError):
+      add(a, bad)
+
+  def test_jit_shape_views_mismatch(self):
+    @TinyJit
+    def add(a): return (a+1).realize()
+    with self.assertRaises(JitError):
+      for i in range(1,5):
+        # a has an offset that the kernel doesn't know about
+        a = Tensor.randn(10, 10).realize()[:, i:i+2]
+        add(a)
+
+  def test_jit_duplicate_fail(self):
+    # the jit doesn't support duplicate arguments
+    @TinyJit
+    def add(a, b): return (a+b).realize()
+    a = Tensor.randn(10, 10)
+    with self.assertRaises(JitError):
+      add(a, a)
+
+  def test_jit_assign(self, dtype=dtypes.float32):
+    @TinyJit
+    def add(a):
+      a += 1
+      a.realize()
+    a = Tensor.zeros(1, dtype=dtype).contiguous().realize()
+    for _ in range(5): add(a)
+    self.assertEqual(a.item(), 5)
+
+  def test_jit_assign_int8(self): self.test_jit_assign(dtypes.int8)
+
+  def test_kwargs_jit(self):
+    @TinyJit
+    def add_kwargs(first, second): return (first+second).realize()
+    for _ in range(5):
+      a = Tensor.randn(10, 10)
+      b = Tensor.randn(10, 10)
+      c = add_kwargs(first=a, second=b)
+      np.testing.assert_allclose(c.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
+    assert_jit_cache_len(add_kwargs, 1)
+
+  def test_reorder_kwargs_jit(self):
+    @TinyJit
+    def add_kwargs(first, second): return (first/second).realize()
+    for _ in range(2):
+      a = Tensor.randn(10, 10)
+      b = Tensor.randn(10, 10)
+      c = add_kwargs(second=b, first=a)
+      np.testing.assert_allclose(c.numpy(), a.numpy()/b.numpy(), atol=1e-4, rtol=1e-5)
+    for _ in range(2):
+      a = Tensor.randn(10, 10)
+      b = Tensor.randn(10, 10)
+      c = add_kwargs(first=a, second=b)
+      np.testing.assert_allclose(c.numpy(), a.numpy()/b.numpy(), atol=1e-4, rtol=1e-5)
+    assert_jit_cache_len(add_kwargs, 1)
+
+  def test_array_jit(self):
+    @TinyJit
+    def add_array(a, arr): return (a+arr[0]).realize()
+    for _ in range(5):
+      a, b = Tensor.randn(10, 10).realize(), Tensor.randn(10, 10).realize()
+      np.testing.assert_allclose(add_array(a, [b]).numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
+    assert_jit_cache_len(add_array, 1)
+
+  def test_jit_copyin(self):
+    @TinyJit
+    def f(a):
+      return a + Tensor([1,2,3])
+    for _ in range(5):
+      b = Tensor.randn(3)
+      c = f(b)
+      np.testing.assert_allclose(c.numpy(), b.numpy()+[1,2,3], atol=1e-4, rtol=1e-5)
+
+  def test_method_jit(self):
+    class Fun:
+      def __init__(self):
+        self.a = Tensor.randn(10, 10)
+      @TinyJit
+      def __call__(self, b:Tensor) -> Tensor:
+        return (self.a+b).realize()
+    fun = Fun()
+    for _ in range(5):
+      b = Tensor.randn(10, 10)
+      c = fun(b)
+      np.testing.assert_allclose(c.numpy(), fun.a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
+    assert_jit_cache_len(fun.__call__.func.__self__, 1)
+
+  def test_jit_size1_input(self):
+    @TinyJit
+    def f(a, b): return (a+b).realize()
+    a = Tensor([1, 2, 3])
+    for i in range(5):
+      np.testing.assert_allclose(f(a, Tensor([i])).numpy(), (a+i).numpy(), atol=1e-4, rtol=1e-5)
+    assert_jit_cache_len(f, 1)
+
+  def test_jit_output_non_tensor_fail(self):
+    @TinyJit
+    def f(a, b, i): return (a+b).realize(), i
+    with self.assertRaises(JitError):
+      for i in range(3):
+        f(Tensor.randn(10, 10), Tensor.randn(10, 10), i)
+
+  def test_jit_random_regen(self):
+    def f(a, b):
+      rn = Tensor.randn(*a.shape)
+      return ((a+b)*rn).realize()
+    a = Tensor.randn(10, 10).realize()  # realize these before resetting the random seed
+    b = Tensor.randn(10, 10).realize()
+
+    Tensor.manual_seed(1234)
+    jf = TinyJit(f)
+    res = set()
+    for _ in range(5):
+      o1 = jf(a, b)
+      res.add(o1.numpy()[0][0])
+    assert len(res) == 5, "All values should be different, rand works in jit."
+
+    Tensor.manual_seed(1234)
+    jf2 = TinyJit(f)
+    res2 = set()
+    for _ in range(5):
+      o1 = jf2(a, b)
+      res2.add(o1.numpy()[0][0])
+    assert len(res2) == 5, "All values should be different, rand works in jit."
+    assert res == res2, "Jit rand is not reproducible with the same seed"
+
+    Tensor.manual_seed(3421)
+    jf3 = TinyJit(f)
+    res3 = set()
+    for _ in range(5):
+      o1 = jf3(a, b)
+      res3.add(o1.numpy()[0][0])
+    assert len(res3) == 5, "All values should be different, rand works in jit."
+    assert res3 != res2, "Jit rand is diff with diff seeds"
+
+  def test_jit_v_nojit_random_regen(self):
+    def f(a, b):
+      rn = Tensor.randn(*a.shape)
+      rn = rn * a
+      rn2 = Tensor.randn(*a.shape)
+      rn2 = rn2 * b
+      rn = rn + rn2
+      rn2 = rn2 + Tensor.randn(*a.shape)
+      return ((a+b)*rn).realize(), ((a+b)*rn2).realize()
+    Tensor.manual_seed(0)
+    a = Tensor.randn(10, 10).realize()  # realize these before resetting the random seed
+    b = Tensor.randn(10, 10).realize()
+
+    Tensor.manual_seed(1234)
+    without_jit = set()
+    for _ in range(5):
+      o1, o2 = f(a, b)
+      without_jit.add(o1.numpy()[0][0])
+      without_jit.add(o2.numpy()[0][0])
+    assert len(without_jit) == 10, "All values should be different."
+
+    Tensor.manual_seed(1234)
+    jf = TinyJit(f)
+    with_jit = set()
+    for _ in range(5):
+      o1, o2 = jf(a, b)
+      with_jit.add(o1.numpy()[0][0])
+      with_jit.add(o2.numpy()[0][0])
+    assert len(with_jit) == 10, "All values should be different."
+    assert with_jit == without_jit, "jit and non-jit should produce the same random values with the same seed"
+
+  def test_jit_multiple_random_regen(self):
+    def f(a, b):
+      rn = Tensor.randn(*a.shape)
+      rn = rn * a
+      rn2 = Tensor.randn(*a.shape)
+      rn2 = rn2 * b
+      rn = rn + rn2
+      rn2 = rn2 + Tensor.randn(*a.shape)
+      return ((a+b)*rn).realize(), ((a+b)*rn2).realize()
+    a = Tensor.randn(10, 10).realize()  # realize these before resetting the random seed
+    b = Tensor.randn(10, 10).realize()
+
+    Tensor.manual_seed(1234)
+    jf = TinyJit(f)
+    res = set()
+    for _ in range(5):
+      o1, o2 = jf(a, b)
+      res.add(o1.numpy()[0][0])
+      res.add(o2.numpy()[0][0])
+    assert len(res) == 10, "All values should be different, rand works in jit."
+
+    Tensor.manual_seed(1234)
+    jf2 = TinyJit(f)
+    res2 = set()
+    for _ in range(5):
+      o1, o2 = jf2(a, b)
+      res2.add(o1.numpy()[0][0])
+      res2.add(o2.numpy()[0][0])
+    assert len(res2) == 10, "All values should be different, rand works in jit."
+    assert res == res2, "Jit rand is not reproducible with the same seed"
+
+    Tensor.manual_seed(3421)
+    jf3 = TinyJit(f)
+    res3 = set()
+    for _ in range(5):
+      o1, o2 = jf3(a, b)
+      res3.add(o1.numpy()[0][0])
+      res3.add(o2.numpy()[0][0])
+    assert len(res3) == 10, "All values should be different, rand works in jit."
+    assert res3 != res2, "Jit rand is diff with diff seeds"
+
+  def test_jit_random_after_unrealized_random(self):
+    @TinyJit
+    def f(): return Tensor.rand()
+    Tensor.manual_seed(1234)
+    Tensor.rand()
+    res = [f().numpy() for _ in range(3)]
+    assert res[1] != res[2]
+
+  def test_jit_realization_and_sampling(self):
+    w = Tensor.eye(5)
+
+    @TinyJit
+    def foo (x): return w.dot(x).realize()
+
+    arg  = [
+        Tensor([1,2,3,4,5]),
+        Tensor([1,3,3,4,6]),
+        Tensor([1,2,5,4,7]),
+        Tensor([0,2,3,1,0]),
+    ]
+
+    Y = [foo(e).numpy() for e in arg]
+
+    foo(Tensor([7,7,7,7,7]))
+    want = [[1., 2., 3., 4., 5.],
+            [1., 3., 3., 4., 6.],
+            [1., 2., 5., 4., 7.],
+            [0., 2., 3., 1., 0.]]
+    np.testing.assert_allclose(want, Y)
+
+  def test_jit_buffer_behavior(self):
+    @TinyJit
+    def foo(x) -> Tensor: return x.sum().realize()
+
+    result_1 = foo(Tensor([1] * 2))
+    result_2 = foo(Tensor([2] * 2))
+    result_3 = foo(Tensor([3] * 2))
+
+    # expect the buffer to share underlying buffer
+    np.testing.assert_allclose(result_1.numpy(), [2], atol=1e-4, rtol=1e-5)
+    np.testing.assert_allclose(result_2.numpy(), [6], atol=1e-4, rtol=1e-5)
+    np.testing.assert_allclose(result_3.numpy(), [6], atol=1e-4, rtol=1e-5)
+
+  def test_jit_batch_split(self):
+    if Device[Device.DEFAULT].graph is None or JIT >= 2: raise unittest.SkipTest("only test graphs")
+
+    # Create long jit with 83 kernels.
+    def f(a, b, c, d, e):
+      for _ in range(80):
+        a = (a+b).realize()
+      y = (a*c).realize()
+      z = (y*d).realize()
+      w = (z*e)
+      return w.realize()
+
+    a = Tensor.randn(10, 10).realize()
+    b = Tensor.randn(10, 10).realize()
+    c = Tensor.randn(10, 10).realize()
+    d = Tensor.randn(10, 10).realize()
+    e = Tensor.randn(10, 10).realize()
+
+    jf = TinyJit(f)
+    prev = None
+    for _ in range(5):
+      o = jf(a, b, c, d, e).numpy()
+      if prev is not None: np.testing.assert_allclose(o, prev, atol=1e-4, rtol=1e-5)
+      prev = o
+
+    # Checking that 2 graphs are inited.
+    assert len(jf.captured.linear.src) == 2
+    for si in jf.captured.linear.src:
+      assert call_is_graph(si)
+
+  def test_jitted_clone(self):
+    def f(a): return a.clone().realize()
+    jf = TinyJit(f)
+    for _ in range(5):
+      a = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+      ja = jf(a)
+      np.testing.assert_allclose(a.numpy(), ja.numpy(), atol=1e-4, rtol=1e-5)
+
+  @needs_second_gpu
+  @unittest.skipIf(not_support_multi_device(), "no multi")
+  def test_jitted_transfers(self):
+    d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
+
+    def f(a, b):
+      x = a.to(d1)
+      y = b.to(d1)
+      return x.realize(), y.realize()
+
+    jf = TinyJit(f)
+    for _ in range(5):
+      a = Tensor.randn(10, 10, device=d0).realize()
+      b = Tensor.randn(10, 10, device=d0).realize()
+      xc, yc = jf(a, b)
+      np.testing.assert_allclose(a.numpy(), xc.numpy(), atol=1e-4, rtol=1e-5)
+      np.testing.assert_allclose(b.numpy(), yc.numpy(), atol=1e-4, rtol=1e-5)
+
+  def test_jit_several_devs(self):
+    d0, d1 = f"{Device.DEFAULT}:0", "CPU"
+
+    def f(a, b):
+      x = a.to(d0).realize()
+      y = b.to(d0).realize()
+      return x+y.realize(), x*y.realize()
+
+    jf = TinyJit(f)
+    for _ in range(5):
+      a = Tensor.randn(10, 10, device=d1).realize()
+      b = Tensor.randn(10, 10, device=d1).realize()
+      zc, wc = jf(a, b)
+      np.testing.assert_allclose((a.numpy()+b.numpy()), zc.numpy(), atol=1e-4, rtol=1e-5)
+      np.testing.assert_allclose((a.numpy()*b.numpy()), wc.numpy(), atol=1e-4, rtol=1e-5)
+
+  @needs_second_gpu
+  @unittest.skipIf(not_support_multi_device(), "no multi")
+  def test_jitted_view(self):
+    d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
+
+    def f(a):
+      x1 = a.sum(axis=(1,))
+      x = (x1 + 5).bitcast(dtypes.int32)
+      y = x.to(d1)
+      return y.realize()
+
+    jf = TinyJit(f)
+    for _ in range(5):
+      a = Tensor.randn(10, 1000, device=d0).realize()
+      xc = jf(a)
+      np.testing.assert_allclose((a.numpy().sum(axis=(1,)) + 5).view(np.int32), xc.numpy(), atol=1e-4, rtol=5e-5)
+
+  def test_jit_output_clone(self):
+    @TinyJit
+    def f(x:Tensor) -> Tensor: return (x + 1).realize()
+
+    f(Tensor([0.0]))
+    f(Tensor([0.0]))
+
+    a = f(Tensor([1.0])).clone().realize()
+    b = f(Tensor([2.0]))
+    assert abs((a - b).item()) > 0.5
+
+  def test_jit_init_empty(self):
+    @TinyJit
+    def f(x:Tensor) -> Tensor: return (x + 1).realize()
+
+    f(Tensor.empty(1))
+    f(Tensor.empty(1))
+    # scalar const input is not allowed
+    with self.assertRaises(JitError):
+      f(Tensor(2.0)).item()
+    # self.assertEqual(f(Tensor([2.0])).item(), 1.0) # TODO: wrong output, should be 3.0. currently depends on empty value
+
+  def test_jit_const_input(self):
+    @TinyJit
+    def f(x:Tensor) -> Tensor: return (x + 1).realize()
+    with self.assertRaises(JitError):
+      f(Tensor(UOp.const(dtypes.float, 2.0))).item()
+
+  def test_jit_deviceless_compute_input(self):
+    @TinyJit
+    def f(x:Tensor) -> Tensor: return (x + 1).realize()
+    with self.assertRaises(JitError):
+      f(Tensor(UOp.const(dtypes.float, 2.0) + UOp.const(dtypes.float, 1.0))).item()
+
+  def test_jit_init_empty_alt(self):
+    @TinyJit
+    def f(a:Tensor, b:Tensor) -> Tensor: return b.assign(a+1)
+    for i in range(4):
+      a = Tensor([i])
+      b = Tensor.empty_like(a)
+      c = f(a, b)
+      self.assertEqual(c.item(), i+1)
+
+@unittest.skip("Pending multioutput implementation #3607")
+class TestMultioutputJit(unittest.TestCase):
+  def _test(self, f):
+    for _ in range(5):
+      a, b = Tensor.randn(10, 10), Tensor.randn(10, 10)
+      out0, out1, out2 = f(a, b)
+      np.testing.assert_allclose(out0.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
+      np.testing.assert_allclose(out1.numpy(), a.numpy()-b.numpy(), atol=1e-4, rtol=1e-5)
+      np.testing.assert_allclose(out2.numpy(), a.numpy()*b.numpy(), atol=1e-4, rtol=1e-5)
+
+  def test_jit_multioutput_realize(self):
+    @TinyJit
+    def fxn(a, b): return (a+b).realize(), (a-b).realize(), (a*b).realize()
+    self._test(fxn)
+    assert_jit_cache_len(fxn, 3)
+
+  def test_jit_multioutput_norealize(self):
+    @TinyJit
+    def fxn(a, b): return a+b, a-b, a*b
+    self._test(fxn)
+    assert_jit_cache_len(fxn, 1)
+
+  def test_jit_multioutput_mix(self):
+    @TinyJit
+    def fxn(a, b): return a+b, a-b, (a*b).realize()
+    self._test(fxn)
+    assert_jit_cache_len(fxn, 2)
+
+class TestJitInsideJit(unittest.TestCase):
+  def test_jit_jit_error(self):
+    @TinyJit
+    def f(t): return t + 1
+
+    @TinyJit
+    def g(t): return f(t) * 3
+
+    # NOTE: first does not raise
+    g(Tensor([1])).realize()
+    with self.assertRaisesRegex(RuntimeError, "having TinyJit inside another TinyJit is not supported"):
+      g(Tensor([1])).realize()
+
+class TestCopyInsideJit(unittest.TestCase):
+  def test_copy_inside_jit(self):
+    @TinyJit
+    def add(x,y) -> Tensor: return x.to(Device.DEFAULT)+y
+    for _ in range(5):
+      # create a Tensor on CPU
+      a = Tensor.rand(16,16,device="CPU").realize()
+      b = Tensor.rand(16,16).realize()
+      out = add(a,b)
+      np.testing.assert_allclose(out.flatten().tolist(), [x+y for x,y in zip(a.flatten().tolist(), b.flatten().tolist())])
+
+class TestJitPrune(unittest.TestCase):
+  def test_simple_prune(self):
+    weights = Tensor.rand(16).realize()
+    def w2(x) -> Tensor: return (weights*2).contiguous() + x
+    w2_noprune = TinyJit(w2)
+    w2_prune = TinyJit(w2, prune=True)
+
+    for _ in range(3):
+      a = Tensor.rand(16).realize()
+      out = w2_noprune(a)
+      np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
+    assert_jit_cache_len(w2_noprune, 2)
+
+    for _ in range(3):
+      a = Tensor.rand(16).realize()
+      out = w2_prune(a)
+      np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
+    assert_jit_cache_len(w2_prune, 1)
+
+  def test_prune_w_copy_correct(self):
+    weights = Tensor.rand(16).realize()
+    def w2(x) -> Tensor: return (weights*2).contiguous() + x.to(Device.DEFAULT)
+    w2_noprune = TinyJit(w2)
+    w2_prune = TinyJit(w2, prune=True)
+
+    for _ in range(3):
+      a = Tensor.rand(16, device="CPU").realize()
+      out = w2_noprune(a)
+      np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
+
+    for _ in range(3):
+      a = Tensor.rand(16, device="CPU").realize()
+      out = w2_prune(a)
+      np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
+
+  def test_prune_w_independent_copy_correct(self):
+    weights = Tensor.rand(16, device="CPU").realize()
+    def w2(x) -> Tensor: return (weights*2).contiguous().to(Device.DEFAULT) + x
+    w2_noprune = TinyJit(w2)
+    w2_prune = TinyJit(w2, prune=True)
+
+    for _ in range(3):
+      a = Tensor.rand(16).realize()
+      out = w2_noprune(a)
+      np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
+
+    for _ in range(3):
+      a = Tensor.rand(16).realize()
+      out = w2_prune(a)
+      np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
+
+    assert_jit_cache_len(w2_prune, 1)
+
+class TestJitFree(unittest.TestCase):
+  def test_free_intermediates(self):
+    ext_tensor = Tensor([1,24,23,45,1])
+    @TinyJit
+    def fxn(x:Tensor):
+      t1 = (x * 2).contiguous().realize()
+      t2 = (t1 + ext_tensor).contiguous().realize()
+      out = (t2.sum()).contiguous().realize()
+      return out
+    for i in range(5):
+      out = fxn(inp:=Tensor([i,1,2,3,4]))
+      self.assertEqual(out.item(), 114+2*i)
+    pre_free = GlobalCounters.mem_used
+    fxn.captured.free_intermediates()
+    savings_after_free = pre_free - GlobalCounters.mem_used
+
+    expected_savings = (len(inp) * inp.dtype.itemsize * 2) + dtypes.float32.itemsize # (t1 and t2) + out
+
+    self.assertGreaterEqual(savings_after_free, expected_savings)
+    out = fxn(Tensor([11,1,2,3,4]))
+    self.assertEqual(out.item(), 136)
+
+    # Try one more time...
+    pre_free = GlobalCounters.mem_used
+    fxn.captured.free_intermediates()
+    fxn.captured.free_intermediates() # 2nd time to validate
+    savings_after_free = pre_free - GlobalCounters.mem_used
+
+    self.assertGreaterEqual(savings_after_free, expected_savings)
+    out = fxn(Tensor([11,1,2,3,4]))
+    self.assertEqual(out.item(), 136)
+
+  def test_updated_not_freed(self):
+    x = Tensor([1]).realize()
+    @TinyJit
+    def fxn(y):
+      nonlocal x
+      x += y
+      return x
+    for _ in range(5): fxn(Tensor([1]))
+    self.assertEqual(x.item(), 6)
+    pre_free = GlobalCounters.mem_used
+    fxn.captured.free_intermediates()
+    savings_after_free = pre_free - GlobalCounters.mem_used
+    self.assertEqual(savings_after_free, 0)
+    fxn(Tensor([2]))
+    self.assertEqual(x.item(), 8)
+
+class TestJitGraphSplit(unittest.TestCase):
+  def compute(self, device, inp):
+    assert inp.device == device, f"Input device {inp.device} does not match expected {device}"
+    return (inp + 1.0).contiguous().realize()
+
+  def copy(self, device, to_device, inp):
+    assert inp.device == device, f"Input device {inp.device} does not match expected {device}"
+    return inp.to(to_device).realize()
+
+  def expect(self, f, *args, graph=None, multigraph=None, hcqgraph=None):
+    def _numpies(tpl): return tpl.numpy() if tpl.__class__ is Tensor else tuple([t.numpy() for t in tpl])
+
+    expected = _numpies(f(*args))
+    for i in range(4):
+      res = _numpies(f(*args))
+      np.testing.assert_allclose(res, expected, atol=1e-4, rtol=1e-5)
+
+    dev = Device[Device.DEFAULT]
+    graph_t = graph_class(dev)
+    if graph_t is None: return
+
+    got = f.captured.linear.src
+    from tinygrad.runtime.graph.hcq import HCQGraph
+    from tinygrad.engine.jit import MultiGraphRunner
+    if graph_t is HCQGraph:
+      validate = hcqgraph
+    elif issubclass(graph_t, MultiGraphRunner):
+      validate = multigraph
+    else:
+      validate = graph
+
+    assert len(got) == len(validate), f"Expected {len(validate)} operations, got {len(got)}"
+    for expected, si in zip(validate, got):
+      ast = si.src[0]
+      if expected["type"] == "graph":
+        assert call_is_graph(si), f"Expected graph, got {ast.op}"
+        inner_cnt = len(ast.src[0].src)
+        assert inner_cnt == expected["cnt"], f"Expected {expected['cnt']} operations in graph, got {inner_cnt}"
+      elif expected["type"] == "comp":
+        assert ast.op in (Ops.SINK, Ops.PROGRAM), f"Expected kernel, got {ast.op}"
+      elif expected["type"] in ("copy", "xfer"):
+        assert ast.op is Ops.COPY, f"Expected COPY, got {ast.op}"
+
+  def ji_graph(self, cnt): return {"type": "graph", "cnt": cnt}
+  def ji_comp(self): return {"type": "comp"}
+  def ji_copy(self): return {"type": "copy"}
+  def ji_xfer(self): return {"type": "xfer"}
+
+  def test_jit_split_simple(self):
+    @TinyJit
+    def f(inp):
+      op0 = self.compute(Device.DEFAULT, inp)
+      op1 = self.compute(Device.DEFAULT, op0)
+      op2 = self.compute(Device.DEFAULT, op1)
+      return op2
+
+    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+    self.expect(f, inp,
+      graph=[self.ji_graph(3)],
+      multigraph=[self.ji_graph(3)],
+      hcqgraph=[self.ji_graph(3)])
+
+  def test_jit_cpu_simple(self):
+    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
+
+    @TinyJit
+    def f(inp, inp_cpu):
+      op0 = self.compute(Device.DEFAULT, inp)
+      op1 = self.compute(Device.DEFAULT, op0)
+      op2 = self.compute("CPU", inp_cpu)
+      op3 = self.compute(Device.DEFAULT, op1)
+      return op2, op3
+
+    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+    inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
+    self.expect(f, inp, inp_cpu,
+      graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
+      multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
+      hcqgraph=[self.ji_graph(4)])
+
+  def test_jit_cpu_several(self):
+    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
+
+    @TinyJit
+    def f(inp, inp_cpu):
+      op0 = self.compute(Device.DEFAULT, inp)
+      op1 = self.compute(Device.DEFAULT, op0)
+      op2 = self.compute("CPU", inp_cpu)
+      op3 = self.compute("CPU", op2)
+      op4 = self.compute(Device.DEFAULT, op1)
+      return op3, op4
+
+    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+    inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
+    self.expect(f, inp, inp_cpu,
+      graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
+      multigraph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
+      hcqgraph=[self.ji_graph(5)])
+
+  def test_jit_multidev(self):
+    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
+
+    try: Device[f"{Device.DEFAULT}:1"]
+    except Exception: raise unittest.SkipTest("no multidevice")
+
+    @TinyJit
+    def f(inp, inp_d1):
+      op0 = self.compute(Device.DEFAULT, inp)
+      op1 = self.compute(Device.DEFAULT, op0)
+      op2 = self.compute(f"{Device.DEFAULT}:1", inp_d1)
+      op3 = self.compute(f"{Device.DEFAULT}:1", op2)
+      op4 = self.compute(Device.DEFAULT, op1)
+      return op3, op4
+
+    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+    inp_d1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:1").realize()
+    self.expect(f, inp, inp_d1,
+      graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
+      multigraph=[self.ji_graph(5)],
+      hcqgraph=[self.ji_graph(5)])
+
+  def test_jit_multidev_xfer(self):
+    if Device.DEFAULT in {"CPU"}: raise unittest.SkipTest("CPU is not a valid default device for this test (zero-copies)")
+    if Device.DEFAULT == "METAL": raise unittest.SkipTest("Metal is flaky, with multidevice (same as metal llama 4gpu?)")
+
+    try: Device[f"{Device.DEFAULT}:1"]
+    except Exception: raise unittest.SkipTest("no multidevice")
+
+    @TinyJit
+    def f(inp, inp_d1):
+      op0 = self.compute(Device.DEFAULT, inp)
+      op1 = self.compute(Device.DEFAULT, op0)
+      op2 = self.compute(f"{Device.DEFAULT}:1", inp_d1)
+      op3 = self.copy(f"{Device.DEFAULT}:1", Device.DEFAULT, op2)
+      op4 = self.compute(f"{Device.DEFAULT}:1", op2)
+      op5 = self.compute(Device.DEFAULT, op3)
+      return op1, op4, op5
+
+    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+    inp_d1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:1").realize()
+    self.expect(f, inp, inp_d1,
+      graph=[self.ji_graph(2), self.ji_comp(), self.ji_xfer(), self.ji_comp(), self.ji_comp()],
+      multigraph=[self.ji_graph(6)],
+      hcqgraph=[self.ji_graph(6)])
+
+  @unittest.skip("this fails if you don't have SDMA or are using AMD_DISABLE_SDMA=1")
+  @unittest.skipIf(DEV.interface.startswith("MOCK"), "MockGPU does not support parallel copies")
+  def test_jit_multidev_copy(self):
+    if Device.DEFAULT in {"CPU"}: raise unittest.SkipTest("CPU/LLVM is not a valid default device for this test (zero-copies)")
+
+    @TinyJit
+    def f(inp):
+      op0 = self.compute(Device.DEFAULT, inp)
+      op1 = self.compute(Device.DEFAULT, op0)
+      op2 = self.copy(Device.DEFAULT, "CPU", op1)
+      op3 = self.compute("CPU", op2)
+      return op3
+
+    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+    self.expect(f, inp,
+      graph=[self.ji_graph(2), self.ji_copy(), self.ji_comp()],
+      multigraph=[self.ji_graph(2), self.ji_copy(), self.ji_comp()],
+      hcqgraph=[self.ji_graph(4)])
+
+class TestJitRandom(unittest.TestCase):
+  def test_jit_rangeify(self):
+    tst = {0:[], 1:[]}
+    for r in [0,1]:
+      Tensor.manual_seed(1337)
+      with Context(JIT=r):
+        _ = Tensor.randint(4, high=3)
+        # this second one makes the behavior different
+        _ = Tensor.randint(4, high=3)
+        @TinyJit
+        def f(): return Tensor.randint(20, high=5)
+        for _ in range(5): tst[r].append(f().tolist())
+    for i, (t0, t1) in enumerate(zip(tst[0], tst[1])):
+      self.assertListEqual(t0, t1, msg=f"mismatch at list {i}")
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_jit_cases.py

@@ -0,0 +1,78 @@
+import unittest
+from tinygrad import TinyJit, Tensor
+
+# The JIT functions as a "capturing" JIT.
+# Whatever kernels ran in the JIT the second run through the function will be the kernels that will run from then on.
+# Explicit inputs to the function are updated in the JIT graph to the new inputs.
+
+# JITs have four tensor types
+#  1. Tensors that are explicit in the input, aka what's passed in. TODO: support lists/dicts/classes, anything get_state works on
+#  2. Tensors that are explicit in the output, aka what's returned. TODO: same as above
+#  3. Tensors that are implicit in the input as a closure.
+#  4. Tensors that are implicit in the output because they were assigned to and realized.
+
+# explicit inputs and outputs are realized on their way in and out of the JIT
+# there's a whole bunch of edge cases and weirdness here that needs to be tested and clarified.
+
+class TestJitCases(unittest.TestCase):
+  def test_explicit(self):
+    # this function has an explicit input and an explicit output
+    @TinyJit
+    def f(x:Tensor):
+      ret:Tensor = x*2
+      return ret
+
+    for i in range(5):
+      out = f(Tensor([i]))
+      self.assertEqual(out.item(), i*2)
+
+  def test_implicit_input(self):
+    # x is the implicit input (like a weight)
+    x = Tensor([0])
+
+    # this function has an implicit input and an explicit output
+    @TinyJit
+    def f():
+      ret:Tensor = x*2
+      return ret
+
+    for i in range(5):
+      # NOTE: this must be realized here, otherwise the update doesn't happen
+      # if we were explicitly tracking the implicit input Tensors, we might not need this realize
+      x.assign(Tensor([i])).realize()
+      out = f()
+      self.assertEqual(out.item(), i*2)
+
+  def test_implicit_output(self):
+    # out is the implicit output (it's assigned to)
+    out = Tensor([0])
+
+    # this function has an explicit input and an implicit output
+    @TinyJit
+    def f(x:Tensor):
+      # NOTE: this must be realized here
+      # if we were explicitly tracking the implicit output Tensors, we might not need this realize
+      out.assign(x*2).realize()
+
+    for i in range(5):
+      f(Tensor([i]))
+      self.assertEqual(out.item(), i*2)
+
+  def test_implicit_io(self):
+    # x is the implicit input (like a weight)
+    # out is the implicit output (it's assigned to)
+    x = Tensor([0])
+    out = Tensor([0])
+
+    # this function has an implicit input and an implicit output
+    @TinyJit
+    def f():
+      out.assign(x*2).realize() # NOTE: this must be realized here
+
+    for i in range(5):
+      x.assign(Tensor([i])).realize()
+      f()
+      self.assertEqual(out.item(), i*2)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_jit_footguns.py

@@ -0,0 +1,414 @@
+#!/usr/bin/env python
+"""
+JIT Footguns: Documenting unexpected behavior changes when using @TinyJit
+
+Each test shows behavior that works without JIT but changes with JIT.
+Comments marked "should be X!" indicate the intuitively expected value.
+
+SILENT MISMATCHES (highest priority - wrong results, no error):
+  class_method_shared_across_instances EASY could check if first arg is self and warn
+  slice_assign_requires_realize      MED    assign graph not connected to read during JIT replay
+  output_buffer_reuse                MED    performance tradeoff, could add option or better docs
+  symbolic_pad_view_frozen           MED    pad view BIND values baked in at capture time
+  python_constants_frozen            HARD   inherent to tracing JITs
+  conditional_branches_frozen        HARD   inherent to tracing JITs
+
+ERRORS RAISED (lower priority - at least users know):
+  item_bakes_in_values               EASY   raises JitError if .item()/.data() accessed during capture
+  unrealized_const_input_error       EASY   raises JitError for unrealized const inputs
+  non_tensor_outputs_error           EASY   raises JitError if return contains non-Tensor values
+  positional_kwargs_cannot_mix       EASY   normalize positional args to kwargs using function signature
+  duplicate_inputs_fail              MED    would need to handle aliasing in input_replace
+  nested_jit_fails_on_second_call    MED    could fail on first call instead of second
+"""
+import unittest
+import numpy as np
+from tinygrad import Tensor, TinyJit, Device
+from tinygrad.engine.jit import JitError
+from tinygrad.helpers import JIT
+
+class TestJitFootguns(unittest.TestCase):
+
+  def test_output_buffer_reuse(self):
+    """Output tensors share buffer after capture - old references get overwritten."""
+    @TinyJit
+    def f(x): return x.sum().realize()
+
+    r1 = f(Tensor([1, 1]))  # warmup
+    r2 = f(Tensor([2, 2]))  # capture
+    r3 = f(Tensor([3, 3]))  # jit exec
+
+    self.assertEqual(r1.item(), 2)  # warmup result independent
+    self.assertEqual(r3.item(), 6)  # latest is correct
+    self.assertEqual(r2.item(), 6)  # should be 4! (overwritten by r3)
+
+  def test_output_buffer_workaround(self):
+    """Use .clone().realize() to get independent copies."""
+    @TinyJit
+    def f(x): return x.sum().realize()
+
+    r1 = f(Tensor([1, 1])).clone().realize()
+    r2 = f(Tensor([2, 2])).clone().realize()
+    r3 = f(Tensor([3, 3])).clone().realize()
+
+    self.assertEqual([r1.item(), r2.item(), r3.item()], [2, 4, 6])
+
+  def test_graph_input_output_aliasing(self):
+    """Test that JIT handles input=output aliasing correctly, simulating LLM generate pattern.
+
+    The LLM generate pattern:
+    1. First "session": multiple iterations where output becomes next input
+    2. Second "session": starts with a NEW input tensor (not the previous output)
+
+    The bug: GraphRunner computes input_replace during _first_run. If at that time input buffer == output buffer
+    (aliasing), it incorrectly includes the output position in input_replace. Later, when a DIFFERENT input
+    is passed, the output position gets overwritten with the input, corrupting the computation.
+
+    This requires multiple kernels to trigger because single-kernel JITs don't get graphed ("only one kernel doesn't graph").
+    """
+    if Device[Device.DEFAULT].graph is None or JIT != 1:
+      self.skipTest("test requires JIT graph support")
+
+    # Multiple operations to create multiple kernels that get batched into a GraphRunner
+    @TinyJit
+    def step(x):
+      y = (x + 1).realize()  # kernel 1
+      z = (y * 2).realize()  # kernel 2
+      return z
+
+    # Phase 1: warmup and capture
+    a = Tensor([10]).contiguous().realize()
+    step(a)  # warmup (cnt=0)
+    b = Tensor([20]).contiguous().realize()
+    x = step(b)  # capture (cnt=1), x = (20+1)*2 = 42
+
+    # Phase 2: first "session" - iterations where output becomes input (triggers _first_run with aliasing)
+    for _ in range(3):
+      x = step(x)  # (42+1)*2=86, (86+1)*2=174, (174+1)*2=350
+    self.assertEqual(x.item(), 350)
+
+    # Phase 3: second "session" - NEW input tensor (simulates new generate() call)
+    # The bug: GraphRunner's input_replace incorrectly includes the output position
+    # When new input y is passed, it overwrites the output buffer, using old value (350) instead of new (100)
+    y = Tensor([100]).contiguous().realize()
+    for _ in range(3):
+      y = step(y)  # should be (100+1)*2=202, (202+1)*2=406, (406+1)*2=814
+    self.assertEqual(y.item(), 814)  # fails with 1406 if bug exists (uses 350 instead of 100)
+
+  def test_multiple_outputs_same_intermediate(self):
+    """Multiple outputs derived from the same intermediate - JIT copies aliased inputs to prevent hazard."""
+    @TinyJit
+    def f(buf, frame):
+      new_buf = buf[1:].cat(frame, dim=0)
+      return new_buf.contiguous(), new_buf[:1].contiguous()
+
+    buf = Tensor([[0], [1], [2]]).contiguous().realize()
+    for i in range(4):
+      frame = Tensor([[10+i]]).contiguous().realize()
+      expected_first = buf[1:2].numpy().item()
+      new_buf, first = f(buf, frame)
+      self.assertEqual(first.numpy().item(), expected_first)
+      buf = new_buf
+
+  def test_intra_kernel_output_input_aliasing(self):
+    """JIT must copy aliased input when output buffer is fed back as input (read-write race in same kernel)."""
+    N = 1 << 20
+    f = TinyJit(lambda buf, new: buf[N//2:].cat(new), prune=True)
+    buf = Tensor.zeros(N, dtype='int32').contiguous().realize()
+    for i in range(10):
+      buf = f(buf, Tensor(np.ones(N//2, dtype=np.int32)*(i+1)))
+      np.testing.assert_array_equal(buf[:N//2].numpy(), np.full(N//2, i, dtype=np.int32))
+
+  def test_slice_assign_works_without_realize(self):
+    """Slice assign then read from same buffer - pending assigns are side-realized."""
+    from tinygrad import Variable
+    v_pos = Variable("pos", 0, 3)
+    cache = Tensor.zeros(4, 4).contiguous().realize()
+    @TinyJit
+    def f(pos):
+      cache[pos:pos+1, :].assign(Tensor.ones(1, 4))
+      return cache.sum().realize()
+    for i in range(4):
+      cache.assign(Tensor.zeros(4, 4)).realize()
+      self.assertEqual(f(v_pos.bind(i)).item(), 4.0)
+
+  def test_symbolic_pad_view_frozen(self):
+    """Symbolic pad view has BIND values baked in at capture time. TODO: pad should be captured in jit."""
+    from tinygrad import Variable
+    a = Tensor.rand(3, 10).realize()
+
+    # broken: pad is a view, BIND values frozen at capture (i=2)
+    @TinyJit
+    def f_broken(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
+    for i in range(1, 5): f_broken(a[:, :Variable("i", 1, 10).bind(i)])
+    self.assertEqual(int((f_broken(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 2)  # should be 4!
+
+    # workaround: contiguous fuses pad into kernel
+    @TinyJit
+    def f_fixed(a): return (a+1).pad((None, (0, 10-a.shape[1]))).contiguous().realize()
+    for i in range(1, 5): f_fixed(a[:, :Variable("i", 1, 10).bind(i)])
+    self.assertEqual(int((f_fixed(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 4)
+
+  def test_non_tensor_outputs_error(self):
+    @TinyJit
+    def f(x, mult): return (x * 2).realize(), mult * 10
+    with self.assertRaises(JitError):
+      for i in range(3): f(Tensor([i]), i)
+
+  def test_duplicate_inputs_fail(self):
+    """JIT cannot handle the same tensor passed as multiple arguments."""
+    @TinyJit
+    def f(a, b): return (a + b).realize()
+
+    x = Tensor([1, 2, 3])
+    with self.assertRaises(JitError):
+      f(x, x)
+
+  def test_tensors_in_containers(self):
+    @TinyJit
+    def f(a, arr): return (a + arr[0]).realize()
+    for i in range(4):
+      a, b = Tensor([1, 1, 1]).realize(), Tensor([i, i, i]).realize()
+      np.testing.assert_array_equal(f(a, [b]).numpy(), [1+i, 1+i, 1+i])
+
+  def test_nested_jit_fails_on_second_call(self):
+    """Nested JIT works on first call but fails on second."""
+    @TinyJit
+    def inner(t): return t + 1
+    @TinyJit
+    def outer(t): return inner(t) * 3
+
+    self.assertEqual(outer(Tensor([1])).realize().item(), 6)  # works!
+    with self.assertRaises(RuntimeError):
+      outer(Tensor([2])).realize()  # fails
+
+  def test_implicit_inputs_need_realize(self):
+    """Closure tensors must be realized before JIT call."""
+    x = Tensor([0])
+
+    @TinyJit
+    def f(): return (x * 2).realize()
+
+    for i in range(5):
+      x.assign(Tensor([i])).realize()  # must realize!
+      self.assertEqual(f().item(), i * 2)
+
+  def test_views_with_different_offsets_fail(self):
+    """JIT requires consistent tensor views across calls."""
+    @TinyJit
+    def f(a): return (a + 1).realize()
+
+    base = Tensor.randn(10, 10).realize()
+    with self.assertRaises(JitError):
+      for i in range(1, 5):
+        f(base[:, i:i+2])  # different offset each time
+
+  def test_shape_change_after_capture_fails(self):
+    """Shapes are locked at capture time."""
+    @TinyJit
+    def f(a, b): return (a + b).realize()
+
+    f(Tensor.randn(10, 10), Tensor.randn(10, 10))  # warmup
+    f(Tensor.randn(10, 10), Tensor.randn(10, 10))  # capture
+
+    with self.assertRaises(JitError):
+      f(Tensor.randn(20, 20), Tensor.randn(20, 20))
+
+  def test_python_constants_frozen(self):
+    """Python variables inside JIT use capture-time values."""
+    mult = 1
+
+    @TinyJit
+    def f(x): return (x * mult).realize()
+
+    results = []
+    for i in range(5):
+      mult = i + 1
+      results.append(f(Tensor([10])).item())
+
+    self.assertEqual(results[0], 10)   # warmup, mult=1
+    self.assertEqual(results[1], 20)   # capture, mult=2
+    self.assertEqual(results[2], 20)   # should be 30!
+    self.assertEqual(results[3], 20)   # should be 40!
+
+  def test_unrealized_const_input_error(self):
+    """Const tensors have no buffer to replace, so JIT raises an error. Even explicit .realize() doesn't help."""
+    @TinyJit
+    def f(a, b): return (a * b).realize()
+
+    # unrealized const fails
+    with self.assertRaises(JitError):
+      f(Tensor([1, 2, 3]).realize(), Tensor(2))
+
+    # explicit .realize() on const still fails - const cannot be realized to have a buffer
+    @TinyJit
+    def g(a, b): return (a * b).realize()
+    with self.assertRaises(JitError):
+      g(Tensor([1, 2, 3]).realize(), Tensor(2).realize())
+
+  def test_conditional_branches_frozen(self):
+    """Only the branch taken during capture runs thereafter."""
+    @TinyJit
+    def f(x, use_square):
+      if use_square:
+        return (x * x).realize()
+      return (x * 2).realize()
+
+    f(Tensor([3]), True)   # warmup
+    f(Tensor([3]), False)  # capture (False branch)
+
+    result = f(Tensor([3]), True)  # passing True but False branch runs
+    self.assertEqual(result.item(), 6)  # should be 9!
+
+  def test_positional_kwargs_cannot_mix(self):
+    """Must use same calling convention after capture."""
+    @TinyJit
+    def f(a, b): return (a + b).realize()
+
+    f(Tensor([1]), Tensor([2]))  # warmup with positional
+    f(Tensor([1]), Tensor([2]))  # capture with positional
+
+    with self.assertRaises(JitError):
+      f(a=Tensor([3]), b=Tensor([4]))  # kwargs fail
+
+  def test_class_method_shared_across_instances(self):
+    """JIT on instance methods is shared at class level."""
+    class Model:
+      def __init__(self, scale):
+        self.scale = Tensor([scale])
+      @TinyJit
+      def forward(self, x):
+        return (x * self.scale).realize()
+
+    m1, m2 = Model(2), Model(3)
+
+    m1.forward(Tensor([5]))  # warmup
+    m1.forward(Tensor([5]))  # capture with m1.scale=2
+
+    self.assertEqual(m1.forward(Tensor([5])).item(), 10)
+    self.assertEqual(m2.forward(Tensor([5])).item(), 10)  # should be 15!
+
+  def test_side_effects_only_during_capture(self):
+    """Function body not executed during JIT replay."""
+    call_count = [0]
+
+    @TinyJit
+    def f(x):
+      call_count[0] += 1
+      return (x * 2).realize()
+
+    f(Tensor([1]))  # warmup
+    f(Tensor([2]))  # capture
+    self.assertEqual(call_count[0], 2)
+
+    f(Tensor([3]))
+    f(Tensor([4]))
+    f(Tensor([5]))
+    self.assertEqual(call_count[0], 2)  # still 2, not 5!
+
+  def test_nothing_realized_fails(self):
+    """Must JIT at least one kernel."""
+    @TinyJit
+    def f(a, b): return None
+
+    with self.assertRaises(JitError):
+      for _ in range(3):
+        f(Tensor([1]), Tensor([2]))
+
+  def test_item_creates_unrealized_return(self):
+    """.item() in shape computation raises error during JIT capture."""
+    @TinyJit
+    def f(x): return Tensor.zeros(x.sum().item())
+
+    f(Tensor([1, 1, 1]))  # warmup
+    with self.assertRaises(JitError):
+      f(Tensor([1, 1, 1]))  # capture - .item() raises
+
+  def test_item_bakes_in_values(self):
+    """.item() during JIT capture raises error (would bake in value)."""
+    @TinyJit
+    def f(x, mask): return x.masked_select(mask)
+
+    f(Tensor([1, 2, 3, 4]), Tensor([True, False, True, False]))  # warmup
+    with self.assertRaises(JitError):
+      f(Tensor([1, 2, 3, 4]), Tensor([True, False, True, False]))  # capture - .item() raises
+
+  def test_masked_select_static_size_jittable(self):
+    @TinyJit
+    def f(x, mask): return x.masked_select(mask, size=4, fill_value=-1).realize()
+
+    for _ in range(3):
+      np.testing.assert_equal(f(Tensor([1, 2, 3, 4]), Tensor([True, False, True, False])).numpy(), [1, 3, -1, -1])
+      np.testing.assert_equal(f(Tensor([5, 6, 7, 8]), Tensor([False, True, True, True])).numpy(), [6, 7, 8, -1])
+      np.testing.assert_equal(f(Tensor([9, 8, 7, 6]), Tensor([True, True, True, True])).numpy(), [9, 8, 7, 6])
+      np.testing.assert_equal(f(Tensor([1, 1, 1, 1]), Tensor([False, False, False, False])).numpy(), [-1, -1, -1, -1])
+
+  def test_nonzero_static_size_jittable(self):
+    @TinyJit
+    def f(x): return x.nonzero(size=3, fill_value=-1).realize()
+
+    for _ in range(3):
+      np.testing.assert_equal(f(Tensor([1, 0, 2, 0, 3])).numpy(), [[0], [2], [4]])
+      np.testing.assert_equal(f(Tensor([0, 0, 5, 0, 0])).numpy(), [[2], [-1], [-1]])
+      np.testing.assert_equal(f(Tensor([0, 0, 0, 0, 0])).numpy(), [[-1], [-1], [-1]])
+
+  def test_tolist_bakes_in_values(self):
+    """.tolist() raises error during JIT capture (would bake in values)."""
+    @TinyJit
+    def f(x): return Tensor(x.tolist())
+
+    f(Tensor([1, 2, 3]))  # warmup
+    with self.assertRaises(JitError):
+      f(Tensor([1, 2, 3]))  # capture - .tolist() raises
+
+
+class TestJitCorrectBehavior(unittest.TestCase):
+  """Behaviors that work correctly - documented for clarity."""
+
+  def test_random_regenerates(self):
+    """Random tensors regenerate each call."""
+    @TinyJit
+    def f(x):
+      return (x + Tensor.rand(3)).realize()
+
+    f(Tensor([0, 0, 0]))  # warmup
+    f(Tensor([0, 0, 0]))  # capture
+
+    results = {tuple(f(Tensor([0, 0, 0])).numpy().tolist()) for _ in range(5)}
+    self.assertEqual(len(results), 5)
+
+  def test_unrealized_return_auto_realized(self):
+    """Unrealized return tensors are auto-realized."""
+    @TinyJit
+    def f(a, b): return a + b  # no explicit realize
+
+    for _ in range(5):
+      a, b = Tensor.randn(10), Tensor.randn(10)
+      np.testing.assert_allclose(f(a, b).numpy(), a.numpy() + b.numpy(), atol=1e-5)
+
+  def test_kwargs_order_doesnt_matter(self):
+    """Kwargs are sorted by name, so order doesn't matter."""
+    @TinyJit
+    def f(first, second): return (first / second).realize()
+
+    for _ in range(3):
+      a, b = Tensor.randn(10), Tensor.randn(10) + 1
+      np.testing.assert_allclose(f(second=b, first=a).numpy(), a.numpy() / b.numpy(), atol=1e-4)
+      np.testing.assert_allclose(f(first=a, second=b).numpy(), a.numpy() / b.numpy(), atol=1e-4)
+
+  def test_input_mutation_consistent(self):
+    """Input mutation via assign works consistently."""
+    @TinyJit
+    def f(x):
+      x += 1
+      x.realize()
+      return x
+
+    a = Tensor([0]).contiguous().realize()
+    for _ in range(5):
+      f(a)
+    self.assertEqual(a.item(), 5)
+
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_kernel_cache.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python
+import unittest
+from tinygrad.tensor import Tensor
+from tinygrad import Device
+
+class TestKernelCache(unittest.TestCase):
+  def test_kernel_cache_in_action(self):
+    if Device.DEFAULT not in ["CPU"]:
+      self.skipTest("No custom kernel cache is implemented")
+
+    const_value = 0.6765677269
+    a = Tensor.rand(4,4).realize()
+    b = Tensor.rand(4,4).realize()
+    x = a + b + const_value
+    x.realize()
+
+    a1 = Tensor.rand(4,4).realize()
+    b1 = Tensor.rand(4,4).realize()
+    orig_compile_func = Device['CPU'].compiler.compile_cached
+    Device['CPU'].compiler.compile_cached = None # making it not callable
+
+    try:
+      x1 = a1 + b1 + const_value
+      x1.realize() # Same kernel should be from cache.
+    finally:
+      Device['CPU'].compiler.compile_cached = orig_compile_func
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/backend/test_linearizer.py

@@ -0,0 +1,459 @@
+import numpy as np
+import unittest
+
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType, buffers
+from tinygrad.device import Device, Buffer
+from tinygrad.tensor import Tensor, _to_np_dtype
+from tinygrad.engine.realize import run_linear
+from tinygrad.codegen import to_program
+from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, DEV
+from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace
+from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.cstyle import CUDARenderer
+from tinygrad.renderer.isa import ISARenderer
+from test.helpers import replace_opts
+MOCKGPU = DEV.interface.startswith("MOCK")
+
+from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
+
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, ISARenderer), "isa backends don't preserve the op spec when lowering")
+class TestLinearizer(unittest.TestCase):
+  def test_arg_dedup(self):
+    # NOTE: this realize exists because Tensor.numpy calls .contiguous() internally
+    # without contiguous folding, rand.to("CPU") and rand.contiguous().to("CPU") are different UOps.
+    # this test asserts they are the identical Buffer
+    # having different buffers is fine for correctness, because the outputs match.
+    a, b = Tensor.randn(4).realize(), Tensor.randn(4).realize()
+    np_a, np_b = a.numpy(), b.numpy()
+    c = ((a.shrink(((0, 2),)) - a.shrink(((2, 4),))) - (b.shrink(((0, 2),)) - b.shrink(((2, 4),))))
+    linear = c.schedule_linear()
+    run_linear(linear)
+    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if s.op is not Ops.BIND]
+    assert len(rawbufs) == 3 and set(rawbufs[1:]) == {a.uop.base.realized, b.uop.base.realized}
+    np_c = (np_a[:2] - np_a[2:]) - (np_b[:2] - np_b[2:])
+    np.testing.assert_allclose(np_c, c.numpy(), atol=1e-4, rtol=1e-4)
+
+  def test_load_removed(self):
+    a = Tensor.rand(1).realize()
+    b = Tensor.rand(1).realize()
+    ta = Tensor.where(Tensor(True), a, b).numpy()
+    tb = Tensor.where(Tensor(False), a, b).numpy()
+    np.testing.assert_equal(a.numpy(), ta)
+    np.testing.assert_equal(b.numpy(), tb)
+
+  @unittest.skip("TODO: some backends insert more casts")
+  def test_cast_there_and_back(self):
+    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
+    out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
+    ast = helper_linearizer_opt(out)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
+
+  @unittest.expectedFailure
+  def test_cast_back_and_there(self):
+    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
+    out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
+    ast = helper_linearizer_opt(out)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
+
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
+  def test_late_bias_load(self):
+    img = Tensor.empty(1, 3, 16, 16)
+    w = Tensor.empty(16, 3, 3, 3)
+    b = Tensor.empty(16)
+    out = img.conv2d(w, b)
+    ast = helper_linearizer_opt(out)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    # slice at the last loop end
+    uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
+    # only valid test if outermost range is the reduce
+    if uops[uslice].src[-1].arg[-1] == AxisType.REDUCE:
+      load_types = [u.src[0].dtype for u in uops[uslice+1:] if u.op == Ops.LOAD]
+      # assert that there is a global load after the reduce ends
+      assert any(dt.addrspace == AddrSpace.GLOBAL for dt in load_types)
+
+  def _test_no_nested_ranges(self, lins, skip=None):
+    for l in lins:
+      range_in_acc = flatten([[x for x in u.src if x.op is Ops.RANGE] for u in l.uops if u.op is Ops.DEFINE_REG])
+      ranges = [u.op for u in l.uops if (u.op is Ops.RANGE and u in range_in_acc) or (u.op is Ops.END and u.src[0] in range_in_acc)]
+      for i,u in enumerate(ranges):
+        if skip and i in skip: continue
+        assert ranges[i-1] != u, f"multireduce nested the ranges! {ranges[i-1], {u}}"
+
+  def test_two_nested_range(self):
+    a = Tensor.randn(2, ).realize()
+    out = a.reshape(2, 1).expand(2, 3).sum()
+    ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()])
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
+    assert len(ranges) == 1 # NOTE: it collapses now
+
+  def test_three_nested_range(self):
+    a = Tensor.randn(2, ).realize()
+    out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum()
+    ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()])
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
+    assert len(ranges) == 1 # NOTE: it collapses now
+
+  def test_two_nested_range_alt_indexing(self):
+    a = Tensor([2, 2]).realize()
+    out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
+    ast = helper_linearizer_opt(out, wanna_output=[24])
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
+    # RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
+    assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
+    # the index of the load doesnt depend on the second range
+    assert any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]])
+    assert any(x.op in {*GroupOp.ALU, Ops.LOAD} for x in uops[ranges[1]:])
+
+  def test_range_outer_op_before_phi(self):
+    a = Tensor.randn(4, 1).realize()
+    b = Tensor.randn(1, 1).realize()
+    out = (a + b[0]).sum() + b[0]
+    ast = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
+    # LOAD -> RANGE -> LOAD -> STORE
+    assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
+
+  def test_range_outer_op_before_phi_nested_range(self):
+    a = Tensor.randn(2, ).realize()
+    b = Tensor.randn(1, 1).realize()
+    out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0]
+    ast = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()])
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
+    assert len(ranges) == 1 # NOTE: it collapses now
+
+  def test_load_dedup(self):
+    # for different leaves in the AST, the same loads may occur.
+
+    a = Tensor.randn(4).realize()
+    # these are of size 3 to avoid float4 coalesce
+    r = a[:-1] + a[1:]
+
+    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
+                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
+    assert num_loads <= 4, "more load uops than needed"
+    assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
+
+  @unittest.skip("this is handled at higher level now")
+  def test_upcast_cse(self):
+    # when upcasting, within a subtree, there may be common expressions.
+
+    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
+    r = a.expand([2]) + b.expand([2])
+
+    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
+                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
+    assert num_ops <= 1, "more alu uops than needed"
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
+  def test_reduce_upcast(self):
+    x, w = Tensor.randn((1,1,3)).realize(), Tensor.randn((1,1,2)).realize()
+    r = Tensor.conv2d(x,w,padding=1).relu()
+
+    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0],
+      [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    accs = [u for u in uops if u.op is Ops.DEFINE_REG]
+    stores = [u for u in uops if u.op is Ops.STORE]
+    assert len(accs) == 0  # it's removed now
+    assert len(stores) == 1
+    assert stores[0].src[1].dtype == dtypes.float.vec(4)
+
+  # NOTE: can reenable, it does work. it just makes BEAM slow
+  @unittest.expectedFailure
+  @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
+  def test_upcast_with_locals_cpu(self):
+    out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
+    prg = to_program(replace_opts(out.schedule_linear().src[-1].src[0], [Opt(OptOps.LOCAL, axis=0, arg=4)]),
+                      renderer=Device[Device.DEFAULT].renderer)
+    self.assertEqual(len(prg.src[3].arg.split("for")), 5)
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason")
+  def test_upcast_with_locals(self):
+    x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
+    r = (x@y).relu()
+    opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
+    program = to_program(replace_opts(r.schedule_linear().src[-1].src[0], opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
+
+    stores = [u for u in tuple(program.src[2].src) if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
+
+    # the first store is to lds and can be upcasted
+    assert stores[0].src[1].dtype == dtypes.float.vec(4)
+    assert any(x.op is Ops.DEFINE_LOCAL for x in stores[0].toposort())
+    # the second store is to gds with no upcasts
+    assert stores[1].src[1].dtype == dtypes.float
+    assert any(x.op is Ops.PARAM for x in stores[1].toposort())
+
+  def test_zero_fold(self):
+    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
+    r = Tensor.stack(a, b)
+    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
+                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
+    assert num_ops == 0, "more alu uops than needed"
+
+  def test_sum_acc_dtype(self):
+    for tensor_dtype, acc_dtype in (
+      (dtypes.bool, dtypes.int), (dtypes.int16, dtypes.int), (dtypes.float16, dtypes.float), (dtypes.bfloat16, dtypes.float)):
+      if tensor_dtype in (dts:=Device[Device.DEFAULT].renderer.supported_dtypes()) and acc_dtype in dts:
+        a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
+        realized_ast = a.schedule_linear().src[-1].src[0]
+        program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
+        local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.DEFINE_REG]
+        assert local[0].dtype.base == acc_dtype
+
+  def test_arg_acc_dtype(self):
+    def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
+      realized_ast = c.schedule_linear().src[-1].src[0]
+      program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
+      local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.DEFINE_REG]
+      self.assertEqual(local[0].dtype.base, expected_dtype)
+
+    tests = (
+      (dtypes.float16, None, dtypes.float),
+      (dtypes.bfloat16, None, dtypes.float),
+      (dtypes.float, None, dtypes.float),
+      (dtypes.float16, dtypes.float16, dtypes.float16),
+      (dtypes.bfloat16, dtypes.bfloat16, dtypes.bfloat16),
+      (dtypes.float, dtypes.float16, dtypes.float16),
+    )
+    for tensor_dtype, acc_dtype, expected_dtype in tests:
+      if tensor_dtype in (dts:=Device[Device.DEFAULT].renderer.supported_dtypes()) and acc_dtype in dts and expected_dtype in dts:
+        a, b = Tensor.rand(8, 8, dtype=tensor_dtype), Tensor.rand(8, 8, dtype=tensor_dtype)
+        helper_arg_acc_dtype(a.sum(dtype=acc_dtype), expected_dtype)
+        helper_arg_acc_dtype(a.matmul(b, dtype=acc_dtype), expected_dtype)
+        helper_arg_acc_dtype(Tensor.einsum("ki,ij->kj", a, b, dtype=acc_dtype), expected_dtype)
+        d, w = Tensor.rand(4, 8, 8, 8, dtype=tensor_dtype), Tensor.rand(8, 8, 2, 2, dtype=tensor_dtype)
+        helper_arg_acc_dtype(d.conv2d(w, dtype=acc_dtype), expected_dtype)
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
+  def test_simple_unroll_no_between_phi_dependencies(self):
+    x, y = Tensor.empty(64, 64), Tensor.empty(64, 64)
+    r = (x@y).relu()
+    opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]
+    ast = helper_linearizer_opt(r, [opt])
+    # the uops graph is DEFINE_REG -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
+    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
+    end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
+    for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
+    for u in uops:
+      if u.op is Ops.STORE and isinstance(dt:=u.src[0].dtype, PtrDType) and dt.addrspace is AddrSpace.REG:
+        if uops.index(u) < begin_range:
+          assert u.src[1].op is Ops.CONST
+        else:
+          assert u.src[1].op in GroupOp.ALU
+          assert begin_range < uops.index(u) < end_range
+      # children of END are placed after ENDRANGE
+      if any(x.op is Ops.END and x.src[1].op in GroupOp.ALU for x in u.src):
+        assert end_range < uops.index(u)
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  def test_default_global_reversed(self):
+    # shrink so that the dims do not collapse
+    t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
+    ast = helper_linearizer_opt(t+1)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
+    idxs = sorted(idxs, key=lambda uop: uop.arg)
+    assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0]
+    assert (idxs[1].arg, idxs[1].src[0].arg) == ('gidx1', 5), idxs[1].arg
+    assert (idxs[2].arg, idxs[2].src[0].arg) == ('gidx2', 4), idxs[2].arg
+
+  def test_sum_collapse(self):
+    t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()
+    sched = [si for si in t.schedule_linear().src if si.src[0].op is Ops.SINK]
+    # sum_collapse is a full collapse now
+    assert len(sched) == 1
+    assert not any(u.op is Ops.REDUCE and len(u.arg[1]) > 0 for u in sched[0].src[0].toposort()), "found reduce in sum collapse"
+    #lin = Kernel(sched[0].ast)
+    #assert not any(u.op is Ops.RANGE for u in lin.linearize().uops), "found loop in sum collapse"
+
+  def test_assign_fold(self):
+    a = Tensor.ones(4, 4).contiguous().realize()
+    m = Tensor.ones(4, 4).shrink(((1, 2), None)).pad(((1, 2), None))
+    a.assign(a+m)
+    a.realize()
+    np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
+
+  @unittest.skipIf(MOCKGPU and isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CUDARenderer)), "PTX indexes differently. might be ok?")
+  def test_where_fold(self):
+    a = Tensor.ones(4, 4).contiguous().realize()
+    b = a.shrink(((1, 2), None)).pad(((1, 2), None))
+    a.assign(b.where(2, a))
+    linear, var_vals = a.linear_with_vars()
+    assert len(linear.src) == 1
+    run_linear(linear, var_vals)
+    np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
+    program = to_program(replace_opts(linear.src[-1].src[0], []), renderer=Device[Device.DEFAULT].renderer)
+    assert not any(u.op == Ops.WHERE for u in tuple(program.src[2].src)), "found where where where should be folded"
+
+  def test_phi_simplification(self):
+    def helper(t, max_ops=0):
+      ast = helper_linearizer_opt(t)
+      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
+      # ignore kernel optimized IF statements for now
+      if if_op:=next((u for u in uops if u.op is Ops.IF), None):
+        uops = uops[:uops.index(if_op)]
+      assert len(set([u.op for u in uops if u.op in {Ops.RANGE, Ops.SPECIAL}])) == 1, "has either specials or ranges, not both"
+      reg_stores = [u for u in uops if u.op is Ops.STORE and isinstance(dt:=u.src[0].dtype, PtrDType) and dt.addrspace == AddrSpace.REG]
+      assert len(reg_stores) == 0, "STORE to reg should have been simplified"
+      assert len([u for u in uops if u.op is Ops.MAX]) <= max_ops, "no unnecessary MAX ops"
+
+    helper(Tensor.arange(5.5, (3.5*300), 3.5), max_ops=2)
+    helper(Tensor.arange(-1, -100, -5), max_ops=2)
+    # NOTE: both of these split the reduce (this just wasn't tracked before)
+    #helper(Tensor.arange(-3.2, 6.7, 0.64), max_ops=2)
+    #helper(Tensor.arange(256), max_ops=2)
+    helper(Tensor.arange(255), max_ops=2)
+
+  @unittest.skip("test implicitly depends on certain optimizations")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason")
+  def test_grouped_store_phis(self):
+    """
+    float4 acc0 = float4(0.0,0.0,0.0,0.0);
+    {
+      acc0 = // ...
+    }
+    *((device float4*)(data0+alu2)) = float4(acc0.x,acc0.y,acc0.z,acc0.w);
+    simplifies to:
+    *((device float4*)(data0+alu2)) = acc0;
+    """
+    x, y = Tensor.empty(64,64), Tensor.empty(64,64)
+    out = x.matmul(y)
+    with Context(TC=0):
+      ast = helper_linearizer_opt(out)
+      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    # check that the float4 cast collapses
+    store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
+    for val in store_vals:
+      assert val.dtype == dtypes.float.vec(4) # and val.op is not Ops.VECTORIZE
+
+  @unittest.skip("test implicitly depends on certain optimizations")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
+  def test_grouped_store_values(self):
+    x = Tensor.randn((4,3,6,6)).realize()
+    out = x.flip((0,1)).contiguous()
+    ast = helper_linearizer_opt(out)
+    store_val = [u.src[1] for u in tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src) if u.op is Ops.STORE][0]
+    assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.STACK
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
+  def test_grouped_store_locals_and_globals(self):
+    x, y = Tensor.empty(64, 64), Tensor.empty(64, 64)
+    out = x@y
+    opt = [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUPTOP, 0, 8),
+            Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
+    ast = helper_linearizer_opt(out, opts=[opt])
+    def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
+    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
+    global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.PARAM for x in get_recursive(u.src[0]))]
+    barrier = [u for u in uops if u.op is Ops.BARRIER]
+    assert len(barrier) == 1
+    # check that the float4 cast collapses for all stores
+    for store in local_stores+global_stores:
+      assert store.src[1].dtype.count > 1 # and store.src[2].op is not Ops.VECTORIZE
+    # # check the children's vins
+    # TODO: src ALU are not the same, should it?
+    # assert barrier.src == tuple(local_stores)
+    assert len([u for u in uops if u.op is Ops.IF])
+
+  @unittest.skip("test implicitly depends on certain optimizations")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason")
+  def test_grouped_store_local_only(self):
+    x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
+    r = (x@y).relu()
+    ast = helper_linearizer_opt(r)
+    uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    stores = [u for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
+
+    # the float4 value stores directly in lds and we skip upcast
+    self.assertEqual(stores[0].src[1].dtype, dtypes.float.vec(4))
+    #assert stores[0].src[-1].op is not Ops.VECTORIZE
+
+    # the global store doesn't change
+    assert stores[1].src[1].dtype == dtypes.float
+
+# *** helpers ***
+
+def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
+  if isinstance(r, Tensor): r = [r]
+  linear, var_vals = Tensor.linear_with_vars(*r)
+  run_linear(UOp(Ops.LINEAR, src=linear.src[:-1]), var_vals)  # run all kernels except the last one
+  last_call = linear.src[-1]
+  ast = last_call.src[0]
+  assert ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {last_call}"
+  last_bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
+  # now all input buffers in last_call should be realized
+  # create fresh buffers for the outputs
+  bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(ast.src) else x for i,x in enumerate(last_bufs)]
+  # ensure buffers are allocated
+  for b in bufs: b.ensure_allocated()
+  return ast, bufs
+
+def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
+  assert isinstance(ast, UOp), "ast must be UOp"
+  inbufs = [x.uop.base.buffer for x in inputs]
+  outbufs = [Buffer(inbufs[-1].device if inbufs else Device.DEFAULT, out.size, out.src[1].dtype).allocate() for out in ast.src]
+  _helper_linearizer_opt_ast(ast, outbufs+inbufs, *args, **kwargs)
+
+def helper_linearizer_opt(r:Tensor|list[Tensor], *args, **kwargs):
+  realized_ast, real_bufs = helper_realized_ast(r)
+  _helper_linearizer_opt_ast(realized_ast, real_bufs, *args, **kwargs)
+  return realized_ast
+
+def copyout_outputs(outbufs:list[Buffer]) -> list[np.ndarray]:
+  return [np.frombuffer(x.as_memoryview(), _to_np_dtype(x.dtype)) for x in outbufs]
+
+def reset_bufs(bufs:list[Buffer]):
+  for buf in bufs: buf.copyin(np.zeros((buf.size*buf.dtype.itemsize,), dtype=np.uint8).data)
+
+def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
+                               apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
+  outbufs = real_bufs[:len(realized_ast.src)]
+  wanna_output = [np.array(x).flatten() for x in wanna_output]
+  buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in real_bufs]
+  for u,b in zip(buf_uops, real_bufs): buffers[u] = b
+
+  def run_prg(opts):
+    ast = realized_ast if opts is None else replace_opts(realized_ast, list(opts))
+    run_linear(UOp(Ops.LINEAR, src=(ast.call(*buf_uops),)))
+
+  def check_opt(opts):
+    reset_bufs(outbufs)
+    run_prg(opts)
+    for x,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(x, want, atol=atol, rtol=rtol)
+
+  # Get baseline if it is not provided, which is not optimized at all.
+  run_prg(opts=())
+  if len(wanna_output) == 0: wanna_output = copyout_outputs(outbufs)
+  else:
+    for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
+
+  # Check correctness of handcoded optimiztions.
+  reset_bufs(outbufs)
+  run_prg(opts=None)
+  for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
+  for x in opts: # Check custom transformations if any.
+    check_opt(([Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, 1))] if apply_tc else [])+x)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_linearizer_dumb.py

@@ -0,0 +1,31 @@
+# ruff: noqa: E501
+# tests where the Linearizer is doing something dumb
+# like test_linearizer_failures, but they don't have to fail
+
+import unittest
+from tinygrad import Device, dtypes
+from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
+from tinygrad.codegen.opt.search import Opt, OptOps
+from tinygrad.codegen import to_program
+
+class TestLinearizerFailure(unittest.TestCase):
+  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
+  def test_failure_beam_mnist(self):
+    c0 = UOp.param(0, dtypes.uchar.ptr(4014080))
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 0, AxisType.GLOBAL)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 784), 1, AxisType.GLOBAL)
+    c3 = UOp.range(UOp.const(dtypes.weakint, 10), 3, AxisType.GLOBAL)
+    c4 = UOp.param(1, dtypes.int.ptr(512))
+    c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True)))
+    c6 = UOp.range(UOp.const(dtypes.weakint, 6000), 1004, AxisType.REDUCE)
+    c7 = UOp.range(UOp.const(dtypes.weakint, 3750), 2006, AxisType.REDUCE)
+    c8 = UOp.range(UOp.const(dtypes.weakint, 16), 2007, AxisType.GROUP_REDUCE)
+    c9 = UOp.param(2, dtypes.uchar.ptr(47040000))
+    c10 = c9.index((((c3*UOp.const(dtypes.weakint, 4704000))+c2)+(c6*UOp.const(dtypes.weakint, 784))).valid(UOp.const(dtypes.bool, True)))
+    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
+    c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
+    ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
+    _ = to_program(ast, Device["METAL"].renderer)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_llama_kernels.py

@@ -0,0 +1,96 @@
+import unittest
+from tinygrad import Tensor, Device, dtypes, Context, GlobalCounters
+from tinygrad.helpers import getenv
+from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
+from extra.llama_kernels.fused_ce import fused_ce_loss
+from extra.llama_kernels import local_abs_max
+from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar
+from test.helpers import needs_second_gpu
+
+def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None:
+  Tensor.manual_seed(0)
+  logits_rand = Tensor.randn(bs, seqlen, vocab).cast(dtypes.bfloat16)
+  targets = Tensor.randint(bs, seqlen, high=vocab, dtype=dtypes.int32)
+  logits, logits_ref = logits_rand.clone(), logits_rand.detach().float().contiguous()
+  with Context(DEBUG=0):
+    Tensor.realize(logits, logits_ref, targets)
+
+  loss = fused_ce_loss(logits, targets, label_smoothing=label_smoothing)
+  loss.backward()
+  Tensor.realize(loss, logits.grad)
+
+  ref = logits_ref.sparse_categorical_crossentropy(targets, label_smoothing=label_smoothing)
+  ref.backward()
+  Tensor.realize(ref, logits_ref.grad)
+
+  assert logits.grad.shape == (bs, seqlen, vocab)
+  with Context(DEBUG=0):
+    assert loss.allclose(ref, atol=2e-3, rtol=2e-3).item(), "forward mismatch"
+    assert logits.grad.allclose(logits_ref.grad, atol=2e-3, rtol=2e-3).item(), "grad mismatch"
+
+class TestFusedCE(unittest.TestCase):
+  def setUp(self):
+    if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("need bfloat16")
+
+  def test_fused_ce_1_2_16(self): run_fused_ce(1, 2, 16, label_smoothing=0.2)
+  def test_fused_ce_2_16_128(self): run_fused_ce(2, 16, 128)
+  def test_fused_ce_4_128_1024(self): run_fused_ce(4, 128, 1024, label_smoothing=0.2)
+
+  # note: this is the shape used in llama 8b
+  #def test_fused_ce_smoothing_16_1024_128256(self): run_fused_ce(16, 1024, 128256, label_smoothing=0.2)
+
+def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
+  Tensor.manual_seed(0)
+  x = Tensor.randn(*shape).cast(dtypes.bfloat16).contiguous()
+  amax_state = Tensor.full((), 2.0, dtype=dtypes.float32).contiguous()
+  with Context(DEBUG=0): Tensor.realize(x, amax_state)
+
+  if delayed:
+    fp8, inv_scale, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
+    ref_fp8, ref_inv_scale, ref_new_amax = quantize_fp8(x, amax_state=amax_state)
+    Tensor.realize(fp8, inv_scale, new_amax)
+    Tensor.realize(ref_fp8, ref_inv_scale, ref_new_amax)
+  else:
+    fp8 = quantize_fp8_scalar(x, amax_state, FP8_DTYPE)
+    ref_fp8, _, _ = quantize_fp8(x, amax_state=amax_state)
+    Tensor.realize(fp8)
+    Tensor.realize(ref_fp8)
+
+  with Context(DEBUG=0):
+    assert fp8.cast(dtypes.float).allclose(ref_fp8.cast(dtypes.float), atol=0, rtol=0).item(), "fp8 mismatch"
+    if delayed:
+      assert inv_scale.allclose(ref_inv_scale, atol=0, rtol=0).item(), "inv_scale mismatch"
+      assert new_amax.allclose(ref_new_amax, atol=0, rtol=0).item(), \
+        f"amax mismatch: got={new_amax.item()} ref={ref_new_amax.item()} diff={abs(new_amax.item()-ref_new_amax.item())}"
+
+class TestQuantizeFP8(unittest.TestCase):
+  def setUp(self):
+    ren = Device[Device.DEFAULT].renderer
+    if dtypes.bfloat16 not in ren.supported_dtypes(): self.skipTest("need bfloat16")
+    if not ren.has_local or not ren.has_shared: self.skipTest("need local/shared")
+
+  def test_scalar(self): run_quantize_fp8((getenv("N", 1024), 32), delayed=False)
+  def test_delayed(self): run_quantize_fp8((getenv("N", 2048), 1024))
+
+  @needs_second_gpu
+  def test_multi(self):
+    devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8))
+    x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.multi(0)
+    x = Tensor(x, device=devs)
+    amax_state = Tensor.full((), 2.0, dtype=dtypes.float32, device=devs).contiguous()
+    fp8, _, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
+    Tensor.realize(fp8, new_amax)
+    assert fp8.uop.shape == x.uop.shape
+    assert new_amax.shape == ()
+
+class TestLocalAmax(unittest.TestCase):
+  def test_multi_tensor_local_shard_amax(self):
+    devices = ("CPU:0", "CPU:1")
+    x = Tensor.arange(16, device=devices[0]).reshape(4, 4).cast(dtypes.float).contiguous().realize().shard(devices, axis=0).realize()
+    GlobalCounters.reset()
+    out = (x * local_abs_max(x)).contiguous().realize()
+    self.assertEqual(GlobalCounters.kernel_count, 4)
+    self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_nn.py

@@ -0,0 +1,624 @@
+#!/usr/bin/env python
+import unittest
+import numpy as np
+import torch
+from tinygrad import Tensor, Device, TinyJit, dtypes
+from tinygrad.uop.ops import Ops
+from tinygrad.helpers import GlobalCounters, Context
+from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding
+from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell
+from tinygrad.nn.state import load_state_dict
+from tinygrad.engine.realize import run_linear
+from test.helpers import not_support_multi_device, needs_second_gpu, slow
+
+@slow
+class TestNN(unittest.TestCase):
+  def test_batchnorm2d(self, training=False, threed=False, track_running_stats=True):
+    with Tensor.train(training):
+      szs = [4, 8, 16, 32]
+      for sz in szs:
+        # create in tinygrad
+        bn = BatchNorm(sz, eps=1e-5, track_running_stats=track_running_stats)
+        bn.weight = Tensor.randn(sz)
+        bn.bias = Tensor.randn(sz)
+        if track_running_stats:
+          bn.running_mean = Tensor.randn(sz)
+          bn.running_var = Tensor.randn(sz)
+          bn.running_var.numpy()[bn.running_var.numpy() < 0] = 0
+
+        # create in torch
+        with torch.no_grad():
+          if threed:
+            tbn = torch.nn.BatchNorm3d(sz, track_running_stats=track_running_stats).eval()
+          else:
+            tbn = torch.nn.BatchNorm2d(sz, track_running_stats=track_running_stats).eval()
+          tbn.training = training
+          tbn.weight[:] = torch.tensor(bn.weight.numpy())
+          tbn.bias[:] = torch.tensor(bn.bias.numpy())
+          if track_running_stats:
+            tbn.running_mean[:] = torch.tensor(bn.running_mean.numpy())
+            tbn.running_var[:] = torch.tensor(bn.running_var.numpy())
+
+        if track_running_stats:
+          np.testing.assert_allclose(bn.running_mean.numpy(), tbn.running_mean.detach().numpy(), rtol=1e-5, atol=1e-6)
+          np.testing.assert_allclose(bn.running_var.numpy(), tbn.running_var.detach().numpy(), rtol=1e-5, atol=1e-6)
+
+        # trial
+        if threed:
+          inn = Tensor.randn(2, sz, 3, 3, 3)
+        else:
+          inn = Tensor.randn(2, sz, 3, 3)
+
+        # in tinygrad
+        outt = bn(inn)
+
+        # in torch
+        toutt = tbn(torch.tensor(inn.numpy()))
+
+        # close
+        np.testing.assert_allclose(outt.numpy(), toutt.detach().numpy(), rtol=5e-4, atol=1e-6)
+        if track_running_stats:
+          np.testing.assert_allclose(bn.running_mean.numpy(), tbn.running_mean.detach().numpy(), rtol=1e-5, atol=1e-6)
+          np.testing.assert_allclose(bn.running_var.numpy(), tbn.running_var.detach().numpy(), rtol=1e-5, atol=1e-6)
+
+  def test_batchnorm2d_training(self): self.test_batchnorm2d(True, False, True)
+  def test_batchnorm2d_no_running_stats(self): self.test_batchnorm2d(False, False, False)
+  def test_batchnorm2d_training_no_running_stats(self): self.test_batchnorm2d(True, False, False)
+  def test_batchnorm3d(self): self.test_batchnorm2d(False, True, True)
+  def test_batchnorm3d_training(self): self.test_batchnorm2d(True, True, True)
+  def test_batchnorm3d_no_running_stats(self): self.test_batchnorm2d(False, True, False)
+  def test_batchnorm3d_training_no_running_stats(self): self.test_batchnorm2d(True, True, False)
+
+  def test_batchnorm_axis(self):
+    sz = (2, 4, 3, 2, 2)
+    x = Tensor.randn(sz)
+    weight = Tensor.randn(2, 3)
+    bias = Tensor.randn(2, 3)
+    mean = Tensor.randn(2, 3)
+    invstd = Tensor.randn(2, 3)
+    a = (x.batchnorm(weight, bias, mean, invstd, axis=(0, 2))
+         .permute(1, 0, 2, 3, 4).reshape(4, 6, 2, 2))
+    b = (x.permute(1, 0, 2, 3, 4).reshape(4, 6, 2, 2)
+         .batchnorm(weight.flatten(), bias.flatten(), mean.flatten(), invstd.flatten()))
+    t_x = torch.tensor(x.permute(1, 0, 2, 3, 4).reshape(4, 6, 2, 2).numpy())
+    t_weight, t_bias = torch.tensor(weight.flatten().numpy()), torch.tensor(bias.flatten().numpy())
+    t_mean, t_invstd = torch.tensor(mean.flatten().numpy()), torch.tensor(invstd.flatten().numpy())
+    torch.nn.functional.batch_norm(t_x, t_mean, 1.0 / t_invstd**2, t_weight, t_bias)
+
+    np.testing.assert_allclose(a.numpy(), b.numpy())
+
+  def test_linear(self):
+    def _test_linear(x, in_dim, out_dim):
+      # create in tinygrad
+      model = Linear(in_dim, out_dim)
+      z = model(x)
+
+      # create in torch
+      with torch.no_grad():
+        torch_layer = torch.nn.Linear(in_dim, out_dim).eval()
+        torch_layer.weight[:] = torch.tensor(model.weight.numpy(), dtype=torch.float32)
+        torch_layer.bias[:] = torch.tensor(model.bias.numpy(), dtype=torch.float32)
+        torch_x = torch.tensor(x.numpy(), dtype=torch.float32)
+        torch_z = torch_layer(torch_x)
+
+      # test
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
+
+    BS, T, in_dim, out_dim = 4, 2, 8, 16
+    _test_linear(Tensor.randn(BS, in_dim), in_dim, out_dim)
+    _test_linear(Tensor.randn(BS, T, in_dim), in_dim, out_dim) # test with more dims
+
+  def _test_conv(self, tiny_conv, torch_conv, BS, C1, DIMS, C2, K, S, P, D=1):
+    # create in tinygrad
+    layer = tiny_conv(C1, C2, kernel_size=K, stride=S, padding=P, dilation=D)
+
+    # create in torch
+    with torch.no_grad():
+      torch_layer = torch_conv(C1, C2, kernel_size=K, stride=S, padding=P, dilation=D).eval()
+      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
+      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
+
+    # test
+    x = Tensor.uniform(BS, C1, *DIMS)
+    z = layer(x)
+    torch_x = torch.tensor(x.numpy())
+    torch_z = torch_layer(torch_x)
+    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
+
+  def test_conv1d(self): self._test_conv(Conv1d, torch.nn.Conv1d, BS=4, C1=16, DIMS=[224//4], C2=64, K=7, S=2, P=1)
+  def test_conv2d(self): self._test_conv(Conv2d, torch.nn.Conv2d, BS=4, C1=16, DIMS=[224//4, 224//4], C2=64, K=7, S=2, P=1)
+
+  def test_conv1d_same_padding(self):
+    self._test_conv(Conv1d, torch.nn.Conv1d, BS=8, C1=3, DIMS=[32], C2=16, K=3, S=1, P='same')
+  def test_conv2d_same_padding_odd_input(self):
+    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[29, 31], C2=32, K=5, S=1, P='same')
+  def test_conv2d_same_padding_large_kernel(self):
+    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[28, 33], C2=32, K=9, S=1, P='same')
+  def test_conv2d_same_padding_with_dilation(self):
+    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 28], C2=32, K=3, S=1, P='same', D=3)
+
+  def test_conv2d_same_padding_invalid_stride(self):
+    self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=2, padding='same')
+  def test_conv2d_same_padding_invalid_padding_str(self):
+    self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=1, padding='not_same')
+
+  @unittest.skip("Takes too long to compile for Compiled backends")
+  def test_conv2d_winograd(self):
+    BS, C1, H, W = 2, 8, 16, 16
+    C2, K, S, P = 8, 3, 1, 1
+
+    # create in tinygrad
+    layer = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
+
+    # create in torch
+    torch_layer = torch.nn.Conv2d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
+    torch_layer.weight = torch.nn.Parameter(torch.tensor(layer.weight.numpy(), dtype=torch.float32))
+    torch_layer.bias = torch.nn.Parameter(torch.tensor(layer.bias.numpy(), dtype=torch.float32))
+
+    # test
+    x = Tensor.uniform(BS, C1, H, W)
+
+    with Context(WINO=1):
+      z = layer(x)
+
+    m = z.mean()
+    m.backward()
+
+    torch_x = torch.tensor(x.numpy(), requires_grad=True)
+    torch_z = torch_layer(torch_x)
+    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
+
+    gw = layer.weight.grad.realize()
+    gb = layer.bias.grad.realize()
+    gx = x.grad.realize()
+
+    torch_z.mean().backward()
+    np.testing.assert_allclose(gw.numpy(), torch_layer.weight.grad.numpy(), atol=5e-4, rtol=1e-5)
+    np.testing.assert_allclose(gb.numpy(), torch_layer.bias.grad.numpy(), atol=5e-4, rtol=1e-5)
+    np.testing.assert_allclose(gx.numpy(), torch_x.grad.numpy(), atol=5e-4, rtol=1e-5)
+
+  def test_conv_transpose1d(self):
+    self._test_conv(ConvTranspose1d, torch.nn.ConvTranspose1d, BS=4, C1=16, DIMS=[224//4], C2=64, K=7, S=2, P=1)
+  def test_conv_transpose2d(self):
+    self._test_conv(ConvTranspose2d, torch.nn.ConvTranspose2d, BS=4, C1=16, DIMS=[224//4, 224//4], C2=64, K=7, S=2, P=1)
+
+  def test_groupnorm(self):
+    BS, H, W, C, G = 20, 10, 10, 6, 3
+
+    # create in torch
+    torch_layer = torch.nn.GroupNorm(G, C).eval()
+
+    # create in tinygrad
+    layer = GroupNorm(G, C)
+    layer.weight = Tensor(torch_layer.weight.detach().numpy())
+    layer.bias = Tensor(torch_layer.bias.detach().numpy())
+
+    for _ in range(10):
+      # forward
+      x = Tensor.randn(BS, C, H, W)
+      z = layer(x)
+      z.sum().backward()
+
+      torch_x = torch.tensor(x.numpy(), requires_grad=True)
+      torch_z = torch_layer(torch_x)
+      torch_z.sum().backward()
+
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
+      np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
+      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
+      np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
+
+  def test_layernorm_forward(self):
+    N, C, H, W = 20, 5, 10, 10
+
+    # create in torch
+    torch_layer = torch.nn.LayerNorm([H, W]).eval()
+
+    # create in tinygrad
+    layer = LayerNorm([H, W])
+    layer.weight = Tensor(torch_layer.weight.detach().numpy())
+    layer.bias = Tensor(torch_layer.bias.detach().numpy())
+
+    x = Tensor.empty(N, C, H, W)
+    z = layer(x)
+    z.realize()
+
+    torch_x = torch.tensor(x.numpy(), requires_grad=True)
+    torch_z = torch_layer(torch_x)
+    torch_z.sum().backward()
+
+    # TODO: why is torch numbers all 0?
+    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=5e-6)
+
+  def test_layernorm(self):
+    N, C, H, W = 20, 5, 10, 10
+
+    # create in torch
+    torch_layer = torch.nn.LayerNorm([H, W]).eval()
+
+    # create in tinygrad
+    layer = LayerNorm([H, W])
+    layer.weight = Tensor(torch_layer.weight.detach().numpy())
+    layer.bias = Tensor(torch_layer.bias.detach().numpy())
+
+    for _ in range(10):
+      # forward
+      x = Tensor.randn(N, C, H, W)
+      z = layer(x)
+      z.sum().backward()
+
+      torch_x = torch.tensor(x.numpy(), requires_grad=True)
+      torch_z = torch_layer(torch_x)
+      torch_z.sum().backward()
+
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
+      np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
+      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
+      np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
+
+  def test_layernorm_2d(self):
+    N, C, H, W = 20, 5, 10, 10
+
+    # create in torch
+    torch_layer = torch.nn.LayerNorm([C]).eval()
+
+    # create in tinygrad
+    layer = LayerNorm2d(C)
+    layer.weight = Tensor(torch_layer.weight.detach().numpy())
+    layer.bias = Tensor(torch_layer.bias.detach().numpy())
+
+    for _ in range(10):
+      # forward
+      x = Tensor.randn(N, C, H, W)
+      z = layer(x)
+      z.sum().backward()
+
+      torch_x = torch.tensor(x.numpy(), requires_grad=True)
+      torch_z = torch_layer(torch_x.permute(0,2,3,1)).permute(0,3,1,2)
+      torch_z.sum().backward()
+
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
+      np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
+      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
+      np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
+
+  def test_instancenorm_2d(self):
+    N, C, H, W = 20, 10, 10, 10
+
+    # create in torch
+    torch_layer = torch.nn.InstanceNorm2d(C, affine=True).eval()
+
+    # create in tinygrad
+    layer = InstanceNorm(C)
+    layer.weight = Tensor(torch_layer.weight.detach().numpy())
+    layer.bias = Tensor(torch_layer.bias.detach().numpy())
+
+    for _ in range(10):
+      # forward
+      x = Tensor.randn(N, C, H, W)
+      z = layer(x)
+      z.sum().backward()
+
+      torch_x = torch.tensor(x.numpy(), requires_grad=True)
+      torch_z = torch_layer(torch_x)
+      torch_z.sum().backward()
+
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
+      np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
+      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
+      np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
+
+  def test_instancenorm_3d(self):
+    N, C, D, H, W = 20, 10, 10, 10, 10
+
+    # create in torch
+    torch_layer = torch.nn.InstanceNorm3d(C, affine=True).eval()
+
+    # create in tinygrad
+    layer = InstanceNorm(C)
+    layer.weight = Tensor(torch_layer.weight.detach().numpy())
+    layer.bias = Tensor(torch_layer.bias.detach().numpy())
+
+    for _ in range(10):
+      # forward
+      x = Tensor.randn(N, C, D, H, W)
+      z = layer(x)
+      z.sum().backward()
+
+      torch_x = torch.tensor(x.numpy(), requires_grad=True)
+      torch_z = torch_layer(torch_x)
+      torch_z.sum().backward()
+
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
+      np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
+      # TODO: is this numerical issue or a bug? RANGEIFY big reduce kernel amplifies numerical issue
+      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=1e-2, rtol=1e-3)
+      np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
+
+  def test_rmsnorm(self):
+    class TorchRMSNorm(torch.nn.Module):
+      # https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L34C1-L77C36
+      def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True):
+        super().__init__()
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        self.weight = torch.nn.Parameter(torch.ones(dim)) if elementwise_affine else None
+
+      def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+      def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output if self.weight is None else output * self.weight
+
+    B, T, embed_size = 4, 10, 20
+    torch_layer = TorchRMSNorm(embed_size)
+    layer = RMSNorm(embed_size)
+
+    for _ in range(10):
+      # forward
+      x = Tensor.randn(B, T, embed_size)
+      z = layer(x)
+      z.sum().backward()
+
+      torch_x = torch.tensor(x.numpy(), requires_grad=True)
+      torch_z = torch_layer(torch_x)
+      torch_z.sum().backward()
+
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
+      np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
+      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=2e-3, rtol=1e-3)
+
+    torch_layer = TorchRMSNorm(embed_size, elementwise_affine=False)
+    layer = RMSNorm(embed_size, elementwise_affine=False)
+
+    for _ in range(10):
+      # forward
+      x = Tensor.randn(B, T, embed_size)
+      z = layer(x)
+      z.sum().backward()
+
+      torch_x = torch.tensor(x.numpy(), requires_grad=True)
+      torch_z = torch_layer(torch_x)
+      torch_z.sum().backward()
+
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
+      np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
+
+  def test_embedding(self):
+    B, T, embed_size, vocab_size = 4, 10, 20, 28
+
+    # create in tinygrad
+    layer = Embedding(vocab_size, embed_size)
+
+    with torch.no_grad():
+      torch_layer = torch.nn.Embedding(vocab_size, embed_size).eval()
+      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
+
+    # test
+    x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
+    z = layer(x)
+    torch_x = torch.tensor(x.numpy())
+    torch_z = torch_layer(torch_x)
+    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
+
+    # test with empty input length
+    x = Tensor(np.random.randint(0, vocab_size, (B, 0), dtype=np.int32))
+    z = layer(x)
+    torch_x = torch.tensor(x.numpy())
+    torch_z = torch_layer(torch_x)
+    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
+
+    # test with jit enabled
+    @TinyJit
+    def layer_jit(x):
+      return layer(x).realize()
+
+    for _ in range(3):
+      x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
+      z = layer_jit(x)
+      torch_x = torch.tensor(x.numpy())
+      torch_z = torch_layer(torch_x)
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
+
+  def test_embedding_one_kernel(self, ops=612000, kcount=2):
+    GlobalCounters.reset()
+    layer = Embedding(20, 30)
+    layer.weight = Tensor.zeros_like(layer.weight).contiguous()
+    a = Tensor([[1, 5, 9, 11],
+                [12, 19, 8, 1]])
+    result = layer(a)
+    linear, var_vals = result.linear_with_vars()
+    self.assertEqual(len([call for call in linear.src if call.src[0].op is Ops.SINK]), kcount,
+                     "first run realizes weight and embedding")
+    run_linear(linear, var_vals)
+
+    b = Tensor([[1, 2, 3],
+                [4, 5, 6],
+                [7, 8, 9]])
+    result = layer(b)
+    linear, var_vals = result.linear_with_vars()
+    self.assertEqual(1, len([call for call in linear.src if call.src[0].op is Ops.SINK]),
+                     "second run realizes embedding only")
+    run_linear(linear, var_vals)
+    print(f"Embedding used {GlobalCounters.global_ops} ops")
+    self.assertLessEqual(GlobalCounters.global_ops, ops)
+
+  # TODO: fused with opts uses more ops
+  def test_embedding_one_kernel_fused(self):
+    with Context(NOOPT=0):
+      self.test_embedding_one_kernel(ops=612_000, kcount=2)
+
+  def test_embedding_one_kernel_fused_noopt(self):
+    with Context(NOOPT=1):
+      self.test_embedding_one_kernel(ops=0, kcount=2)
+
+  def test_embedding_shape(self):
+    vocab_size, embed_size = 10, 16
+    layer = Embedding(vocab_size, embed_size)
+    for rank in range(5):
+      shp = (1,) * rank
+      a = Tensor([3]).reshape(shp)
+      result = layer(a)
+      self.assertEqual(result.shape, shp + (embed_size,))
+
+  def test_embedding_regression(self):
+    # used to fail bounds check
+    embedding = Embedding(100, 1024)
+    input_ids = Tensor.empty(16, 16, dtype=dtypes.int)
+    embedding(input_ids).realize()
+
+  def test_load_state_dict(self):
+    layer = Conv2d(3, 5, kernel_size=3)
+
+    state_dict = {
+      'weight': Tensor.randn(5, 3, 3, 3),
+      'bias': Tensor.randn(5),
+    }
+    load_state_dict(layer, state_dict)
+
+    np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
+    np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
+
+  #https://github.com/pytorch/pytorch/blob/d38164a545b4a4e4e0cf73ce67173f70574890b6/torch/nn/modules/module.py#L2425
+  def test_load_conv_num_batches_tracked(self):
+    layer = BatchNorm(sz=1, track_running_stats=False)
+    state_dict = {
+      'weight': Tensor.ones(1),
+      'bias': Tensor.ones(1),
+      'num_batches_tracked': Tensor.ones(1),
+    }
+    load_state_dict(layer, state_dict)
+    state_dict['num_batches_tracked'] = Tensor.empty()
+    load_state_dict(layer, state_dict)
+    layer.num_batches_tracked = Tensor.ones(1)
+    load_state_dict(layer, state_dict)
+
+  @needs_second_gpu
+  @unittest.skipIf(not_support_multi_device(), "no multi")
+  def test_load_state_dict_sharded_model(self):
+    devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
+
+    layer = Conv2d(3, 5, kernel_size=3)
+    layer.weight.shard_(devices, 3)
+    layer.bias.shard_(devices, None)
+    state_dict = {
+      'weight': Tensor.randn(5, 3, 3, 3).realize(),
+      'bias': Tensor.randn(5).realize(),
+    }
+    load_state_dict(layer, state_dict)
+
+    # sharded model shards the state_dict
+    self.assertEqual(layer.weight.device, devices)
+    self.assertEqual(layer.weight.uop.axis, 3)
+    self.assertEqual(layer.bias.device, devices)
+    self.assertEqual(layer.bias.uop.axis, None)
+    np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
+    np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
+
+  @unittest.skipIf(not_support_multi_device, "no multi")
+  def test_load_state_dict_sharded_dict(self):
+    devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
+
+    layer = Conv2d(3, 5, kernel_size=3)
+    state_dict = {
+      'weight': Tensor.randn(5, 3, 3, 3).shard(devices, 3),
+      'bias': Tensor.randn(5).shard(devices, None),
+    }
+    load_state_dict(layer, state_dict)
+
+    # NOTE: model is not sharded, still not sharded after load_state_dict
+    self.assertEqual(layer.weight.device, Device.DEFAULT)
+    self.assertEqual(layer.bias.device, Device.DEFAULT)
+    np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
+    np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
+
+  @needs_second_gpu
+  @unittest.skipIf(not_support_multi_device(), "no multi")
+  def test_load_state_dict_sharded_model_dict_same_axis(self):
+    devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
+
+    layer = Conv2d(3, 5, kernel_size=3)
+    layer.weight.shard_(devices, 3)
+    layer.bias.shard_(devices, None)
+
+    state_dict = {
+      'weight': Tensor.randn(5, 3, 3, 3).shard(devices, 3),
+      'bias': Tensor.randn(5).shard(devices, None),
+    }
+    load_state_dict(layer, state_dict)
+
+    self.assertEqual(layer.weight.device, devices)
+    self.assertEqual(layer.weight.uop.axis, 3)
+    self.assertEqual(layer.bias.device, devices)
+    self.assertEqual(layer.bias.uop.axis, None)
+    np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
+    np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
+
+  @unittest.skipIf(not_support_multi_device, "no multi")
+  def test_load_state_dict_sharded_model_dict_different_axis(self):
+    devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
+    devices5 = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3", f"{Device.DEFAULT}:4", f"{Device.DEFAULT}:5")
+
+    layer = Conv2d(3, 5, kernel_size=3)
+    layer.weight.shard_(devices, 3)
+    layer.bias.shard_(devices, None)
+
+    # different shard axis
+    state_dict = {
+      'weight': Tensor.randn(5, 3, 3, 3).shard(devices, None),
+      'bias': Tensor.randn(5).shard(devices5, 0),
+    }
+    load_state_dict(layer, state_dict)
+
+    # NOTE: model and state_dict shard differently, use the state_dict sharding  # TODO: revisit this?
+    self.assertEqual(layer.weight.device, devices)
+    self.assertEqual(layer.weight.uop.axis, None)
+    self.assertEqual(layer.bias.device, devices5)
+    self.assertEqual(layer.bias.uop.axis, 0)
+    np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
+    np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
+
+  def test_load_state_dict_shape_mismatch(self):
+    d1, d2 = 2, 4
+    layer = Linear(d1, d1, bias=False)
+    state_dict = {'weight': Tensor.randn(d2, d2)}
+    with self.assertRaisesRegex(ValueError, r'Shape mismatch in layer `weight`: Expected shape \(2, 2\), but found \(4, 4\) in state dict.'):
+      load_state_dict(layer, state_dict)
+
+  def test_lstm_cell(self):
+    layer = LSTMCell(32, 16)
+    with torch.no_grad():
+      torch_layer = torch.nn.LSTMCell(32, 16)
+      layer.weight_hh.assign(torch_layer.weight_hh.numpy())
+      layer.weight_ih.assign(torch_layer.weight_ih.numpy())
+      layer.bias_hh.assign(torch_layer.bias_hh.numpy())
+      layer.bias_ih.assign(torch_layer.bias_ih.numpy())
+
+      inp = Tensor.randn(1, 32)
+      out_h, out_c = layer(inp)
+      torch_out_h, torch_out_c = torch_layer(torch.tensor(inp.numpy()))
+      np.testing.assert_allclose(out_h.numpy(), torch_out_h.numpy(), atol=1e-6)
+      np.testing.assert_allclose(out_c.numpy(), torch_out_c.numpy(), atol=1e-6)
+
+      out_h, out_c = layer(inp, (out_h, out_c))
+      torch_out_h, torch_out_c = torch_layer(torch.tensor(inp.numpy()), (torch_out_h, torch_out_c))
+      np.testing.assert_allclose(out_h.numpy(), torch_out_h.numpy(), atol=1e-6)
+      np.testing.assert_allclose(out_c.numpy(), torch_out_c.numpy(), atol=1e-6)
+
+  def test_lstm_cell_no_bias(self):
+    layer = LSTMCell(32, 16, bias=False)
+    inp = Tensor.randn(1, 32)
+    out_h, out_c = layer(inp)
+    out_h.realize()
+    out_c.realize()
+    h = Tensor.randn(1, 16)
+    c = Tensor.randn(1, 16)
+    out_h, out_c = layer(inp, (h, c))
+    out_h.realize()
+    out_c.realize()
+    assert layer.bias_hh is None
+    assert layer.bias_ih is None
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_opt_gemm.py

@@ -0,0 +1,46 @@
+import numpy as np
+import unittest
+from tinygrad import Tensor
+from tinygrad.helpers import get_single_element
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.engine.realize import run_linear
+from tinygrad.uop.ops import Ops, UOp
+from test.helpers import replace_opts
+
+class TestOptGemm(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    N = 64
+    cls.a = Tensor.randn(N, N).contiguous().realize()
+    cls.b = Tensor.randn(N, N).contiguous().realize()
+    with np.errstate(all='ignore'):
+      cls.res = cls.a.T.numpy() @ cls.b.T.numpy()
+
+  def _test_gemm_unrolled_permute_l(self, opts=[]):
+    t = self.a.T @ self.b.T
+    # TODO: this should be a generic test helper
+    call = get_single_element(t.schedule_linear().src)
+    new_call = call.replace(src=(replace_opts(call.src[0], opts), *call.src[1:]))
+    run_linear(UOp(Ops.LINEAR, src=(new_call,)))
+    test = call.src[1].buffer.numpy().reshape(self.res.shape)
+    np.testing.assert_allclose(self.res, test, atol=1e-4)
+
+  def test_gemm_unrolled_permute_l_44(self):
+    opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4)]
+    self._test_gemm_unrolled_permute_l(opts)
+
+  def test_gemm_unrolled_permute_l_424(self):
+    # was failing with LLVM
+    opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
+    self._test_gemm_unrolled_permute_l(opts)
+
+  def test_gemm_unrolled_permute_l_42(self):
+    opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]
+    self._test_gemm_unrolled_permute_l(opts)
+
+  def test_gemm_unrolled_permute_l_22(self):
+    opts = [Opt(op=OptOps.UPCAST, axis=0, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=2)]
+    self._test_gemm_unrolled_permute_l(opts)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_optim.py

@@ -0,0 +1,195 @@
+import numpy as np
+import torch
+import unittest
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.nn.optim import Adam, SGD, AdamW, Muon, LAMB
+from test.helpers import needs_second_gpu, slow
+
+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)
+
+def _param(tensor, val):
+  return tensor(val, requires_grad=True) if tensor is torch.tensor else tensor(val)
+
+class TeenyNet:
+  def __init__(self, tensor):
+    self.x = _param(tensor, x_init.copy())
+    self.W = _param(tensor, W_init.copy())
+  def forward(self):
+    return (self.x * self.W).sum()
+
+class TinyNet:
+  def __init__(self, tensor):
+    self.x = _param(tensor, x_init.copy())
+    self.W = _param(tensor, W_init.copy())
+    self.m = tensor(m_init.copy())
+
+  def forward(self):
+    out = self.x.matmul(self.W).relu()
+    # print(out.detach().numpy())
+    out = out.log_softmax(1)
+    out = out.mul(self.m).add(self.m).sum()
+    return out
+
+def step(tensor, optim, steps=1, teeny=False, **kwargs):
+  net = TeenyNet(tensor) if teeny else TinyNet(tensor)
+  optim = optim([net.x, net.W], **kwargs)
+  for _ in range(steps):
+    out = net.forward()
+    optim.zero_grad()
+    out.backward()
+    optim.step()
+  return net.x.detach().numpy(), net.W.detach().numpy()
+
+@slow
+class TestOptim(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, torch_optim, steps, opts, atol, rtol):
+    for x,y in zip(step(Tensor, tinygrad_optim, steps, **opts),
+                   step(torch.tensor, torch_optim, steps, **opts)):
+      np.testing.assert_allclose(x, y, atol=atol, rtol=rtol)
+
+  def _test_sgd(self, steps, opts, atol, rtol): self._test_optim(SGD, torch.optim.SGD, steps, opts, atol, rtol)
+  def _test_adam(self, steps, opts, atol, rtol): self._test_optim(Adam, torch.optim.Adam, steps, opts, atol, rtol)
+  def _test_adamw(self, steps, opts, atol, rtol): self._test_optim(AdamW, torch.optim.AdamW, steps, opts, atol, rtol)
+  def _test_muon(self, steps, opts, atol, rtol): self._test_optim(Muon, torch.optim.Muon, steps, opts, atol, rtol)
+
+  def test_multistep_sgd_high_lr_teeny(self): self._test_sgd(2, {'lr': 1.1, 'teeny': True}, 1e-6, 1e-5)
+  def test_multistep_adam_high_lr_teeny(self): self._test_adam(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)
+  def test_multistep_muon_high_lr_teeny(self): self._test_muon(2, {'lr': 1.1, 'teeny': True}, 1e-2, 5e-4)
+
+  def test_sgd(self): self._test_sgd(1, {'lr': 0.001}, 1e-6, 0)
+  def test_sgd_high_lr(self): self._test_sgd(1, {'lr': 10}, 1e-6, 1e-5)
+  def test_sgd_wd(self): self._test_sgd(1, {'lr': 0.001, 'weight_decay': 0.1}, 1e-6, 0)
+  def test_sgd_high_lr_wd(self): self._test_sgd(1, {'lr': 10, 'weight_decay': 0.1}, 1e-6, 1e-5)
+
+  def test_multistep_sgd(self): self._test_sgd(10, {'lr': 0.001}, 1e-6, 0)
+  def test_multistep_sgd_high_lr(self): self._test_sgd(10, {'lr': 10}, 1e-6, 3e-4)
+  def test_multistep_sgd_wd(self): self._test_sgd(10, {'lr': 0.001, 'weight_decay': 0.1}, 1e-6, 0)
+  def test_multistep_sgd_high_lr_wd(self): self._test_sgd(10, {'lr': 9, 'weight_decay': 0.1}, 1e-6, 3e-4)
+
+  def test_multistep_sgd_momentum(self): self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9}, 1e-6, 0)
+  def test_multistep_sgd_high_lr_momentum(self): self._test_sgd(10, {'lr': 10, 'momentum': 0.9}, 1e-5, 3e-4)
+  def test_multistep_sgd_momentum_wd(self): self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9, 'weight_decay': 0.1}, 1e-6, 0)
+  def test_multistep_sgd_high_lr_momentum_wd(self): self._test_sgd(10, {'lr': 10, 'momentum': 0.9, 'weight_decay': 0.1}, 1e-5, 3e-4)
+
+  def test_multistep_sgd_nesterov_momentum(self): self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9, 'nesterov': True}, 1e-5, 0)
+  def test_multistep_sgd_high_lr_nesterov_momentum(self): self._test_sgd(10, {'lr': 10, 'momentum': 0.9, 'nesterov': True}, 1e-5, 3e-4)
+  def test_multistep_sgd_nesterov_momentum_wd(self):
+    self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9, 'nesterov': True, 'weight_decay': 0.1}, 1e-5, 0)
+  def test_multistep_sgd_high_lr_nesterov_momentum_wd(self):
+    self._test_sgd(10, {'lr': 9, 'momentum': 0.9, 'nesterov': True, 'weight_decay': 0.1}, 1e-5, 3e-4)
+
+  def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-3, 0)
+  # TODO: disabled due to big atol
+  # def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
+  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-3, 3e-4)
+  # TODO: disabled due to big atol
+  # def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)
+
+  # NOTE: momentum set to 0.95 by default, nesterov set to True by default
+  def test_multistep_muon_momentum_wd(self): self._test_muon(10, {'lr': 0.001, 'weight_decay': 0.01}, 3e-3, 0)
+  # ns defaults are numerically unstable, but it is tolerable in real training (see nsteps/nparam tests)
+  # TODO: disabled due to big atol
+  # def test_multistep_muon_high_lr_momentum_wd(self): self._test_muon(10, {'lr': 10, 'weight_decay': 0.01}, 1e-1, 3e-4)
+  def test_multistep_muon_no_nesterov_momentum(self): self._test_muon(10, {'lr': 0.001, 'nesterov': False}, 1e-3, 0)
+  # TODO: disabled due to big atol
+  # def test_multistep_muon_high_lr_no_nesterov_momentum(self): self._test_muon(10, {'lr': 10, 'nesterov': False}, 5e-2, 1e-1)
+
+  def test_muon_ns_steps(self): self._test_muon(1, {'lr': 0.001, 'ns_steps': 3}, 1e-4, 0)
+  # TODO: disabled due to big atol
+  # def test_muon_high_lr_ns_steps(self): self._test_muon(1, {'lr': 10, 'ns_steps': 3}, 1e-5, 3e-4)
+  def test_muon_ns_coefficients(self): self._test_muon(1, {'lr': 0.001,'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+  # TODO: disabled due to big atol
+  # def test_muon_high_lr_ns_coefficients(self): self._test_muon(1, {'lr': 10,'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+
+  def test_muon_momentum_wd_ns_steps_ns_coefficients(self):
+    self._test_muon(10, {'lr': 0.001, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_coefficients': (2.0,-1.5,0.5)}, 1e-4, 0)
+  # TODO: disabled due to big atol
+  # def test_multistep_muon_high_lr_momentum_wd_ns_steps_ns_coefficients(self):
+  #   self._test_muon(10, {'lr': 10, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+
+  def test_adam(self): self._test_adam(1, {'lr': 0.001}, 1e-5, 0)
+  def test_adam_high_lr(self): self._test_adam(1, {'lr': 10}, 1e-4, 1e-4)
+  def test_adamw(self): self._test_adamw(1, {'lr': 0.001}, 1e-5, 0)
+  def test_adamw_high_lr(self): self._test_adamw(1, {'lr': 10}, 1e-4, 1e-4)
+
+  def test_multistep_adam(self): self._test_adam(10, {'lr': 0.001}, 1e-5, 0)
+  def test_multistep_adam_high_lr(self): self._test_adam(10, {'lr': 10}, 2e-3, 5e-4)
+
+  def test_multistep_adamw(self): self._test_adamw(10, {'lr': 0.001}, 1e-5, 0)
+  def test_multistep_adamw_high_lr(self): self._test_adamw(10, {'lr': 10}, 5e-4, 2e-3)
+
+  def test_duped_weights(self):
+    for Opt in [Adam, AdamW, SGD]:
+      losses = []
+      for i in range(2):
+        w = Tensor(x_init.copy())
+        opt = Opt([w], lr=0.1) if i == 0 else Opt([w, w], lr=0.1)
+
+        loss = None
+        for _ in range(3):
+          loss = w.sum()
+          opt.zero_grad()
+          loss.backward()
+          opt.step()
+        losses.append(loss.numpy())
+
+      np.testing.assert_allclose(losses[0], losses[1], atol=1e-4, rtol=0)
+
+  @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
+  def test_mixed_precision(self):
+    old_default_float, dtypes.default_float = dtypes.default_float, dtypes.half
+    # weight update would overflow without upcasting
+    self._test_sgd(10, {'lr': 1e10}, 1e-6, 3e-4)
+    self._test_adam(1, {'lr': 1e10}, 1e-4, 1e-4)
+    self._test_adamw(1, {'lr': 1e10}, 1e-4, 1e-4)
+    dtypes.default_float = old_default_float
+
+  def test_assert_tensor_train(self):
+    t = Tensor.ones((1,1))
+    optimizer = Adam([t])
+    optimizer.zero_grad()
+    old_state = Tensor.training
+    t.sum().backward()
+    Tensor.training = False
+    self.assertRaises(RuntimeError, optimizer.step)
+    Tensor.training = True
+    optimizer.step()
+    Tensor.training = old_state
+
+  def test_lamb_cpu_offload(self):
+    # test that LAMB works when optimizer params (m, v, b1_t, b2_t) are moved to CPU
+    t = Tensor(x_init.copy())
+    opt = LAMB([t])
+    # move optimizer state to CPU
+    for p in opt.m + opt.v + [opt.b1_t, opt.b2_t]: p.to_("CPU")
+    # run a step
+    t.sum().backward()
+    opt.step()
+    self.assertEqual(t.device, Device.DEFAULT)
+    self.assertEqual(opt.m[0].device, "CPU")
+
+  @needs_second_gpu
+  def test_lamb_cpu_offload_multi(self):
+    ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
+    t = Tensor(x_init.copy()).shard(ds, axis=1)
+    ds = t.device
+    opt = LAMB([t])
+    # move optimizer state to CPU
+    for p in opt.m + opt.v + [opt.b1_t, opt.b2_t]: p.to_("CPU")
+    # run a step
+    t.sum().backward()
+    opt.step()
+    self.assertEqual(t.device, ds)
+    self.assertEqual(opt.m[0].device, "CPU")
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_pickle.py

@@ -0,0 +1,180 @@
+import unittest, pickle, types
+import numpy as np
+from tinygrad import Tensor, TinyJit, Variable, dtypes
+from tinygrad.helpers import GlobalCounters, ContextVar, Context
+from tinygrad.uop.ops import PatternMatcher, UPat, UOp
+
+class TestPickle(unittest.TestCase):
+  def test_pickle_code_object(self):
+    y = lambda x: x*2  # noqa: E731
+    code_str = pickle.dumps(y.__code__)
+    fxn = types.FunctionType(pickle.loads(code_str), globals())
+    self.assertEqual(fxn(2), 4)
+
+  def test_pickle_pattern_matcher(self):
+    pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)])
+    sink = UOp.const(dtypes.int, 2)
+    tt = pm.rewrite(sink)
+    pm_str = pickle.dumps(pm)
+    pm2 = pickle.loads(pm_str)
+    self.assertEqual(pm2.rewrite(sink).key, tt.key)
+
+  def test_pickle_main_pattern_matcher(self):
+    from tinygrad.uop.symbolic import sym
+    ssym = pickle.dumps(sym)
+    dsym = pickle.loads(ssym)
+    self.assertEqual(dsym.patterns[0][0].location, sym.patterns[0][0].location)
+
+  def test_pickle_realized_tensor(self):
+    print("** init")
+    t = Tensor.rand(10, 10).realize()
+    st = pickle.dumps(t)
+    t_values = t.numpy()
+    del t # free buffers
+    print("** post pickle")
+    GlobalCounters.reset()
+    t2:Tensor = pickle.loads(st)
+    np.testing.assert_equal(t_values, t2.numpy())
+    # expect at most one COPY kernel
+    self.assertLessEqual(GlobalCounters.kernel_count, 1)
+
+  def test_pickle_realized_tensor_alt(self):
+    print("** init")
+    t = Tensor.rand(10, 10).to("CPU").realize()
+    st = pickle.dumps(t)
+    t_values = t.numpy()
+    del t # free buffers
+    print("** post pickle")
+    t2:Tensor = pickle.loads(st)
+    assert t2.uop.is_realized
+    np.testing.assert_equal(t_values, t2.numpy())
+
+  def test_pickle_realized_tensor_alt2(self):
+    print("** init")
+    t = Tensor.rand(10, 10).to("CPU").realize()
+    tensor_uop = t.uop
+    assert tensor_uop.is_realized, f"expected {tensor_uop} to be realized"
+    t_values = t.numpy()
+    # pickle
+    st = pickle.dumps(t)
+    # free buffers
+    del t
+    del tensor_uop
+    print("** post pickle")
+    t2:Tensor = pickle.loads(st)
+    assert t2.uop.is_realized, f"expected {t2.uop} to be realized"
+    np.testing.assert_equal(t_values, t2.numpy())
+
+  # NOTE: currently Buffer exists on the uop, not tensor
+  def test_pickle_buffer_uop(self):
+    t = Tensor.arange(4).clone().realize()
+    a = t.uop
+    assert a.is_realized
+    self.assertIsNotNone(buffer:=a.base.realized)
+    s = pickle.dumps(a)
+    # free buffers
+    del a
+    del buffer
+    a2:UOp = pickle.loads(s)
+    self.assertListEqual(a2.base.realized.as_memoryview().cast("I").tolist(), [0, 1, 2, 3])
+
+  def test_pickle_unrealized_tensor(self):
+    t = Tensor.ones(10, 10)
+    st = pickle.dumps(t)
+    t2:Tensor = pickle.loads(st)
+    np.testing.assert_equal(t.numpy(), t2.numpy())
+
+  def test_pickle_variable(self):
+    v = Variable("i", 1, 20).bind(10)
+    t1 = Tensor.ones(10, v).contiguous()
+    t2 = Tensor.ones(10, v).contiguous()
+    ret = (t1+t2).sum(1)
+    st = pickle.dumps(ret)
+    del ret
+    vt2 = pickle.loads(st)
+    np.testing.assert_equal(vt2.numpy(), 20)
+
+  def test_pickle_buffer_view(self):
+    t = Tensor.arange(10).clone(device="CPU").realize()
+    vt = t[3:5].contiguous().realize()
+    assert hasattr(vt.uop.buffer, 'base')
+    ref_value = vt.tolist()
+    st = pickle.dumps(vt)
+    del t, vt
+    vt2 = pickle.loads(st)
+    assert hasattr(vt2.uop.buffer, 'base')
+    assert ref_value == vt2.tolist()
+
+  def test_pickle_numpy(self):
+    t = Tensor(np.array([1,2,3,4.]), dtype=dtypes.float32)
+    st = pickle.dumps(t)
+    t2:Tensor = pickle.loads(st)
+    np.testing.assert_equal(t.numpy(), t2.numpy())
+
+  def test_pickle_jit(self):
+    @TinyJit
+    def add(a, b): return a.sum()+b+1
+    for _ in range(3): add(Tensor.rand(10, 10), Tensor.rand(10, 10))
+    st = pickle.dumps(add)
+    del add
+
+    add_fxn = pickle.loads(st)
+    x = Tensor.ones(10, 10).contiguous().realize()
+    y = Tensor.ones(10, 10).contiguous().realize()
+    print("post jit")
+    out = add_fxn(x, y)
+    np.testing.assert_equal(out.numpy(), 102)
+
+  def test_pickle_jit_no_del(self):
+    @TinyJit
+    def fn(x): return x + 1.0
+    for _ in range(3): fn(Tensor.randn(4))
+    loaded = pickle.loads(pickle.dumps(fn))
+    self.assertEqual(loaded(Tensor([1.0,2.0,3.0,4.0])).tolist(), [2.0,3.0,4.0,5.0])
+
+  def test_pickle_context_var(self):
+    v = ContextVar("test_var", 0)
+    with Context(test_var=1):
+      vs = pickle.dumps(v)
+    v2 = pickle.loads(vs)
+    self.assertEqual(v2.value, 1)
+
+  def test_pickle_schedule(self):
+    a = Tensor([1,2])
+    out = a + 2
+    sched = out.schedule_linear()
+    pk = pickle.dumps(sched)
+    sched_pk = pickle.loads(pk)
+    self.assertEqual(sched_pk.src[-1].src[0], sched.src[-1].src[0])
+
+  def test_pickle_renderer(self):
+    from tinygrad.device import Device
+    pk = pickle.dumps(Device.default.renderer)
+    pickle.loads(pk)
+
+class TestPickleJIT(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    N = 10
+    @TinyJit
+    def add(a, b): return a.sum()+b+1
+    for _ in range(3): add(Tensor.rand(N, N), Tensor.rand(N, N))
+    cls.st = pickle.dumps(add)
+    del add
+
+  def test_inspect(self):
+    import io
+    class FakeClass:
+      def __init__(self, *args, **kwargs):
+        print(self.module, self.name)
+    class InspectUnpickler(pickle.Unpickler):
+      def find_class(self, module, name): return type("SpecializedFakeClass", (FakeClass,), {"name": name, "module": module})
+    InspectUnpickler(io.BytesIO(self.st)).load()
+
+  @unittest.skip("we are still saving intermediate buffers")
+  def test_size(self):
+    # confirm no intermediate buffers are saved
+    self.assertLess(len(self.st), 1_000_000)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_profiler.py

@@ -0,0 +1,227 @@
+import unittest, struct, contextlib, statistics, gc
+from tinygrad import Device, Tensor, dtypes, TinyJit
+from tinygrad.helpers import DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
+from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
+from tinygrad.runtime.support.hcq import HCQCompiled
+from tinygrad.engine.realize import get_runtime
+from tinygrad.codegen import to_program
+
+MOCKGPU = DEV.interface.startswith("MOCK")
+def _dev_base(d):
+  p = d.split(":")
+  return p[0] if len(p) < 2 or not p[1].isdigit() else f"{p[0]}:{p[1]}"
+
+@contextlib.contextmanager
+def helper_collect_profile(*devs):
+  for dev in devs: dev.synchronize()
+  saved = [x for x in Compiled.profile_events if isinstance(x, ProfileDeviceEvent) and x.device.startswith("METAL")]
+  Compiled.profile_events.clear()
+  for x in saved: Compiled.profile_events.append(x)
+
+  cpu_events.clear()
+
+  profile_list = []
+  with Context(PROFILE=1):
+    yield profile_list
+    for dev in devs: dev.synchronize()
+    for dev in devs: dev._at_profile_finalize()
+    for x in Compiled.profile_events: profile_list.append(x)
+    profile_list.extend(cpu_events)
+
+def helper_profile_filter_device(profile, device:str):
+  assert any(getattr(x, "device", None) == device and isinstance(x, ProfileDeviceEvent) for x in profile), f"device {device} is not registred"
+  dev_events = [x for x in profile if getattr(x, "device", None) == device and isinstance(x, ProfileDeviceEvent)]
+  assert len(dev_events) == 1, "only one device registration event is expected"
+  return [x for x in profile if getattr(x, "device", None) == device], dev_events[0]
+
+# TODO: support in HCQCompiled
+is_cpu_hcq = Device.DEFAULT in {"CPU"}
+
+@unittest.skipUnless((issubclass(type(Device[Device.DEFAULT]), HCQCompiled) and not is_cpu_hcq) or Device.DEFAULT in {"METAL"}, "Dev not supported")
+class TestProfiler(unittest.TestCase):
+  @classmethod
+  def setUpClass(self):
+    TestProfiler.d0 = Device[Device.DEFAULT]
+
+    TestProfiler.a = Tensor([0.,1.], device=Device.DEFAULT).realize()
+    TestProfiler.b = self.a + 1
+    si = self.b.schedule_linear().src[-1]
+
+    TestProfiler.prg = to_program(si.src[0], TestProfiler.d0.renderer)
+    TestProfiler.runtime = get_runtime(TestProfiler.d0.device, TestProfiler.prg)
+    TestProfiler.b.uop.buffer.allocate()
+
+  def test_profile_kernel_run(self, wait=False):
+    runner_name = TestProfiler.runtime.name
+    with helper_collect_profile(TestProfiler.d0) as profile:
+      gs, ls = TestProfiler.prg.arg.launch_dims({})
+      TestProfiler.runtime(TestProfiler.b.uop.buffer._buf, TestProfiler.a.uop.buffer._buf, global_size=gs, local_size=ls, wait=wait)
+
+    profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
+    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent)]
+    assert len(kernel_runs) == 1, "one kernel run is expected"
+    assert kernel_runs[0].name == runner_name, "kernel name is not correct"
+    assert _dev_base(kernel_runs[0].device) == kernel_runs[0].device, "kernel should not be on a sub-device"
+
+  def test_profile_kernel_run_wait(self):
+    self.test_profile_kernel_run(wait=True)
+
+  def test_profile_copyin(self):
+    buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
+
+    with helper_collect_profile(TestProfiler.d0) as profile:
+      buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
+
+    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
+    assert len(kernel_runs) == 1, "one kernel run is expected"
+
+  def test_profile_multiops(self):
+    runner_name = TestProfiler.runtime.name
+    buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
+
+    with helper_collect_profile(TestProfiler.d0) as profile:
+      buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
+      gs, ls = TestProfiler.prg.arg.launch_dims({})
+      TestProfiler.runtime(buf1._buf, TestProfiler.a.uop.buffer._buf, global_size=gs, local_size=ls)
+      buf1.copyout(memoryview(bytearray(buf1.nbytes)))
+
+    evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
+
+    assert len(evs) == 3, "3 kernel runs are expected"
+    # NOTE: order of events does not matter, the tool is responsible for sorting them
+    prg_events = [e for e in evs if e.device == TestProfiler.d0.device]
+    assert any(e.name == runner_name for e in prg_events), "kernel name is not correct"
+
+    #for i in range(1, 3):
+    #  assert evs[i].st > evs[i-1].en, "timestamp not aranged"
+
+  def test_profile_multidev(self):
+    try: d1 = Device[f"{Device.DEFAULT}:1"]
+    except Exception as e: self.skipTest(f"second device not available {e}")
+
+    buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
+    buf2 = Buffer(f"{Device.DEFAULT}:1", 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
+
+    with helper_collect_profile(TestProfiler.d0, d1) as profile:
+      buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
+      buf2.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
+
+    for dev in [TestProfiler.d0.device, d1.device]:
+      evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and _dev_base(x.device) == dev]
+      assert len(evs) == 1, "one kernel runs are expected"
+
+  def test_profile_multidev_transfer(self):
+    try: d1 = Device[f"{Device.DEFAULT}:1"]
+    except Exception as e: self.skipTest(f"second device not available {e}")
+
+    buf1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:0").realize()
+    with helper_collect_profile(TestProfiler.d0, d1) as profile:
+      buf1.to(f"{Device.DEFAULT}:1").realize()
+
+    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
+    assert len(kernel_runs) == 1, "one kernel run is expected"
+
+  @unittest.skipIf(Device.DEFAULT in "METAL" or (MOCKGPU and Device.DEFAULT == "AMD"), "AMD mockgpu does not support queue wait interrupts")
+  def test_profile_graph(self):
+    try: d1 = Device[f"{Device.DEFAULT}:1"]
+    except Exception as e: self.skipTest(f"second device not available {e}")
+
+    def f(a):
+      x = (a + 1).realize()
+      return x, x.to(d1.device).realize()
+
+    a = Tensor.randn(10, 10, device=TestProfiler.d0.device).realize()
+    with helper_collect_profile(TestProfiler.d0, d1) as profile:
+      jf = TinyJit(f)
+      for _ in range(3): jf(a)
+      del jf
+
+    graph_evs = [x for x in profile if isinstance(x, ProfileGraphEvent)]
+
+    _, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
+    _, _ = helper_profile_filter_device(profile, d1.device)
+
+    assert len(graph_evs) == 2, "2 graph events are expected"
+    assert len(graph_evs[0].ents) == 2, "two entities are expected"
+
+  @unittest.skipIf(MOCKGPU, "skip MOCKGPU")
+  @unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "must be HCQ")
+  def test_dev_jitter_matrix(self):
+    dev_cnt = 6
+    try: devs = [Device[f"{Device.DEFAULT}:{i}"] for i in range(dev_cnt)]
+    except Exception as e: self.skipTest(f"multiple devices not available {e}")
+
+    for dev in devs: dev.synchronize()
+    for dev in devs: dev._at_profile_finalize()
+
+    def _sync_d2d(d1:HCQCompiled, d2:HCQCompiled):
+      d1.hw_compute_queue_t().signal(d1.timeline_signal, d1.timeline_value).wait(d2.timeline_signal, d2.timeline_value) \
+                             .timestamp(d1.timeline_signal).signal(d1.timeline_signal, d1.timeline_value+1).submit(d1)
+      d2.hw_compute_queue_t().signal(d2.timeline_signal, d2.timeline_value).wait(d1.timeline_signal, d1.timeline_value) \
+                             .timestamp(d2.timeline_signal).signal(d2.timeline_signal, d2.timeline_value+1).submit(d2)
+      d1.timeline_value += 2
+      d2.timeline_value += 2
+      d1.timeline_signal.wait(d1.timeline_value - 1)
+      d2.timeline_signal.wait(d2.timeline_value - 1)
+      return d2.timeline_signal.timestamp - d1.timeline_signal.timestamp
+
+    # then test it by timing the GPU to GPU times
+    dev_evs = {x.device:x for x in Compiled.profile_events if isinstance(x, ProfileDeviceEvent)}
+    jitter_matrix = [[float('nan')] * len(devs) for _ in range(len(devs))]
+    pairs = [(p1, p2) for p1 in enumerate(devs) for p2 in enumerate(devs) if p1 != p2]
+    for (i1, d1), (i2, d2) in pairs:
+      cpu_diff = dev_evs[d1.device].tdiff - dev_evs[d2.device].tdiff
+      jitter_matrix[i1][i2] = statistics.median(_sync_d2d(d1, d2) - _sync_d2d(d2, d1) for _ in range(20)) / 2 - cpu_diff
+
+    print("pairwise clock jitter matrix (us):\n" + '\n'.join([''.join([f'{float(item):8.3f}' for item in row]) for row in jitter_matrix]))
+
+    for (i1, d1), (i2, d2) in pairs:
+      assert abs(jitter_matrix[i1][i2]) < 0.5, "jitter should be less than 0.5us"
+
+  def test_cpu_profile(self):
+    def test_fxn(err=False):
+      if err: raise Exception()
+
+    with helper_collect_profile(dev:=TestProfiler.d0) as profile:
+      with cpu_profile("test_1", dev):
+        test_fxn(err=False)
+      with self.assertRaises(Exception):
+        with cpu_profile("test_2", dev):
+          test_fxn(err=True)
+
+    range_events = [p for p in profile if isinstance(p, ProfileRangeEvent) and p.device == dev]
+    self.assertEqual(len(range_events), 2)
+
+  @unittest.skip("this test is flaky")
+  @unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
+  def test_graph(self):
+    from test.backend.test_graph import helper_alloc_rawbuffer, helper_exec_op, helper_test_graphs
+    device = TestProfiler.d0.device
+    bufs = [helper_alloc_rawbuffer(device, fill=True) for _ in range(5)]
+    graphs = [[helper_exec_op(device, bufs[0], [bufs[1], bufs[2]]), helper_exec_op(device, bufs[0], [bufs[3], bufs[4]]),]]
+    with helper_collect_profile(dev:=TestProfiler.d0) as profile:
+      helper_test_graphs(dev.graph, graphs, runs:=2)
+      # NOTE: explicitly trigger deletion of all graphs
+      graphs.clear()
+      gc.collect()
+    graphs = [e for e in profile if isinstance(e, ProfileGraphEvent)]
+    self.assertEqual(len(graphs), runs)
+    for ge in graphs:
+      self.assertEqual(len(ge.ents), len(graphs))
+
+  @unittest.skip("this test is flaky")
+  def test_trace_metadata(self):
+    with Context(TRACEMETA=1):
+      a = Tensor.empty(1)+2
+      b = Tensor.empty(1)+2
+      with helper_collect_profile(TestProfiler.d0) as profile:
+        Tensor.realize(a, b)
+    profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
+    exec_points = [e for e in profile if isinstance(e, ProfilePointEvent) and e.name == "exec"]
+    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent) and _dev_base(e.device) == e.device]
+    self.assertEqual(len(exec_points), len(range_events), 2)
+    self.assertEqual(len(dedup(e.arg['name'] for e in exec_points)), 1)
+    self.assertEqual(len(dedup(e.arg['metadata'] for e in exec_points)), 1)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/backend/test_quantize_onnx.py

@@ -0,0 +1,237 @@
+# ruff: noqa: E501
+import numpy as np
+import unittest
+from tinygrad import Tensor, Context, Device, dtypes, UOp
+from tinygrad.uop.ops import Ops
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.engine.realize import run_linear
+from tinygrad.codegen import to_program
+from test.helpers import replace_opts
+
+N = 512
+
+def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias=False):
+  import onnx
+  from onnx import helper, numpy_helper, TensorProto
+  # Define input and output
+  input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [batch_size, in_size])
+  output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, [batch_size, out_size])
+
+  # Create random weights and bias
+  W_data = np.random.randn(in_size, out_size).astype(np.float32)
+  W_init = numpy_helper.from_array(W_data, name="W")
+
+  if bias:
+    B_data = np.random.randn(out_size).astype(np.float32)
+    B_init = numpy_helper.from_array(B_data, name="B")
+    gemm_node = helper.make_node("Gemm", inputs=["input", "W", "B"], outputs=["output"], alpha=1.0, beta=1.0, transB=0)
+    graph_def = helper.make_graph([gemm_node], "SingleGemmGraph", [input_tensor], [output_tensor], initializer=[W_init, B_init])
+  else:
+    gemm_node = helper.make_node("Gemm", inputs=["input", "W"], outputs=["output"], alpha=1.0, beta=1.0, transB=0)
+    graph_def = helper.make_graph([gemm_node], "SingleGemmGraph", [input_tensor], [output_tensor], initializer=[W_init])
+
+  # Create and save the model
+  #model_def = helper.make_model(graph_def, producer_name="single_gemm_example")
+  # TODO remove this once ORT supports 1.18.0
+  model_def = helper.make_model(graph_def, producer_name="single_gemm_example", ir_version=10, opset_imports=[helper.make_opsetid("", 22)])
+  onnx.save_model(model_def, model_path)
+  return model_path
+
+def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
+  linear = out.schedule_linear()
+  call = linear.src[-1]
+  prg = to_program(replace_opts(call.src[0], opts), renderer=Device[Device.DEFAULT].renderer)
+  if replace_src is not None:
+    old_name = prg.src[3].arg.split("__attribute__((noinline)) void ")[1].split("(")[0]
+    new_src = replace_src + "/* DSP boilerplate */" + prg.src[3].arg.split("/* DSP boilerplate */")[1].replace(old_name, "fxn")
+    # drop BINARY and replace SOURCE so run_linear recompiles
+    prg = prg.replace(src=prg.src[:3] + (UOp(Ops.SOURCE, arg=new_src),))
+  linear = linear.replace(src=linear.src[:-1] + (call.replace(src=(prg, *call.src[1:])),))
+  for _ in range(run_count): run_linear(linear)
+
+def get_quantized_model(sz):
+  from onnxruntime.quantization import quantize_static, QuantFormat, QuantType, CalibrationDataReader
+  class FakeDataReader(CalibrationDataReader):
+    def __init__(self): self.cnt = 0
+    def get_next(self) -> dict:
+      self.cnt += 1
+      if self.cnt == 100: return None
+      return {"input": np.random.uniform(size=(sz, sz)).astype(np.float32)}
+  out_file = "/tmp/test_out.onnx"
+  quantize_static(create_gemm_model("/tmp/test_in.onnx", sz, sz, sz), out_file,
+                  FakeDataReader(), quant_format=QuantFormat.QDQ, per_channel=False, reduce_range=False,
+                  activation_type=QuantType.QUInt8, weight_type=QuantType.QInt8,
+                  extra_options={"ActivationSymmetric": False})
+  return out_file
+
+@unittest.skip("this is broken")
+@unittest.skipIf(Device.DEFAULT != "CPU", "only tests for CPU")
+class TestQuantizeOnnxCPU(unittest.TestCase):
+  def test_quant_128(self, sz=128):
+    try:
+      import onnx # noqa: F401 # pylint: disable=unused-import
+    except ImportError:
+      raise unittest.SkipTest()
+    from tinygrad.nn.onnx import OnnxRunner
+    out_file = get_quantized_model(sz)
+    run_onnx = OnnxRunner(out_file)
+    inp = Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32))
+    with Context(QUANTIZE=1):
+      linear = run_onnx({"input":inp})["output"].schedule_linear()
+      prg = to_program(linear.src[-2].src[0], renderer=Device[Device.DEFAULT].renderer)
+      daccs = [u for u in tuple(prg.src[2].src) if u.op is Ops.DEFINE_REG]
+      assert all(u.dtype.scalar() is dtypes.int for u in daccs)
+
+@unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP")
+class TestQuantizeOnnx(unittest.TestCase):
+  def test_quant_128(self): self.test_quant(128)
+  def test_quant(self, sz=512):
+    from examples.benchmark_onnx import load_onnx_model
+    # divide is ~1500-2000 without reduce_range, 750-900 with it
+    out_file = get_quantized_model(sz)
+    run_onnx_jit, _ = load_onnx_model(out_file)
+    run_onnx_jit(input=Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32)))
+
+  def test_prequant_conv2d_1x1(self):
+    X = Tensor(np.random.uniform(0, 255, size=(1, 32, 128, 128)).astype(np.uint8))
+    W = Tensor(np.random.uniform(0, 255, size=(64, 32, 1, 1)).astype(np.uint8))
+    out = X.conv2d(W, dtype=X.dtype)
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
+    sexec(out, opts)
+
+  def test_prequant_gemm(self):
+    N = 512
+    X = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8))
+    W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8))
+    out = X.matmul(W, dtype=X.dtype)
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
+    sexec(out, opts)
+
+  # TODO: this has to work
+  def test_prequant_gemm_intacc_early(self, xi=np.int8, wi=np.int8):
+    N = 512
+    X = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(xi))
+    W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(wi))
+    # this divide is interesting and forces the accumulator to actually be an int
+    out = (X.cast("int").matmul(W.cast("int"))//1000).cast("int8")
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
+    sexec(out, opts)
+
+  def test_prequant_gemm_handcode(self):
+    src = """typedef int int128 __attribute__((aligned(512),vector_size(512)));
+    typedef int int32 __attribute__((aligned(128),vector_size(128)));
+    typedef int int64 __attribute__((aligned(256),vector_size(256)));
+    typedef unsigned char unsigned_char4 __attribute__((aligned(4),vector_size(4)));
+    typedef signed char signed_char128 __attribute__((aligned(128),vector_size(128)));
+    typedef unsigned char unsigned_char128 __attribute__((aligned(128),vector_size(128)));
+    typedef unsigned char unsigned_char256 __attribute__((aligned(256),vector_size(256)));
+    union V256 {
+      unsigned_char256 vec256;
+      struct {
+        unsigned_char128 lo128;
+        unsigned_char128 hi128;
+      };
+    };
+    __attribute__((noinline)) void fxn(unsigned char* restrict __attribute__((align_value(128))) data0,
+                                       unsigned char* restrict __attribute__((align_value(128))) data1,
+                                       signed char* restrict __attribute__((align_value(128))) data2) {
+      for (int ridx0 = 0; ridx0 < 512; ridx0++) {
+        int alu0 = (ridx0<<9);
+        for (int ridx1 = 0; ridx1 < 4; ridx1++) {
+          int alu1 = (ridx1<<7);
+          int32 acc0 = __builtin_HEXAGON_V6_vd0_128B();
+          int32 acc1 = __builtin_HEXAGON_V6_vd0_128B();
+          int32 acc2 = __builtin_HEXAGON_V6_vd0_128B();
+          int32 acc3 = __builtin_HEXAGON_V6_vd0_128B();
+
+          for (int ridx2 = 0; ridx2 < 128; ridx2++) {
+            unsigned_char4 val0 = *((unsigned_char4*)((data1+(alu0+(ridx2<<2)))));
+            int alu2 = (alu1+(ridx2<<11));
+            signed_char128 x0 = *((signed_char128*)((data2+alu2)));
+            signed_char128 x1 = *((signed_char128*)((data2+(alu2+512))));
+            signed_char128 x2 = *((signed_char128*)((data2+(alu2+1024))));
+            signed_char128 x3 = *((signed_char128*)((data2+(alu2+1536))));
+
+            union V256 ss01;
+            // ss01.lo128 = (x0[0], x1[0], x0[2], x1[2], x0[4], x1[4], ...)
+            // ss01.hi128 = (x0[1], x1[1], x0[3], x1[3], x0[5], x1[5], ...)
+            ss01.vec256 = __builtin_HEXAGON_V6_vshufoeb_128B(x1, x0);
+
+            union V256 ss23;
+            // ss23.lo128 = (x2[0], x3[0], x2[2], x3[2], x2[4], x3[4], ...)
+            // ss23.hi128 = (x2[1], x3[1], x2[3], x3[3], x2[5], x3[5], ...)
+            ss23.vec256 = __builtin_HEXAGON_V6_vshufoeb_128B(x3, x2);
+
+            union V256 sslo;
+            // sslo.lo128 = (x0[0], x1[0], x2[0], x3[0], x0[4], x1[4], ...)
+            // sslo.hi128 = (x0[2], x1[2], x2[2], x3[2], x0[6], x1[6], ...)
+            sslo.vec256 = __builtin_HEXAGON_V6_vdealvdd_128B(ss23.lo128, ss01.lo128, 2);
+
+            union V256 sshi;
+            // sshi.lo128 = (x0[1], x1[1], x2[1], x3[1], x0[5], x1[5], ...)
+            // sshi.hi128 = (x0[3], x1[3], x2[3], x3[3], x0[7], x1[7], ...)
+            sshi.vec256 = __builtin_HEXAGON_V6_vdealvdd_128B(ss23.hi128, ss01.hi128, 2);
+
+            //unsigned_char128 w0 = (unsigned_char128){val0[0],val0[1],val0[2],val0[3],val0[0],val0[1],val0[2],val0[3],...
+            unsigned_char128 w0 = __builtin_HEXAGON_V6_lvsplatw_128B(*((unsigned int*)&val0));
+
+            acc0 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc0, w0, sslo.lo128);
+            acc1 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc1, w0, sshi.lo128);
+            acc2 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc2, w0, sslo.hi128);
+            acc3 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc3, w0, sshi.hi128);
+          }
+          acc0 /= 1000;
+          acc1 /= 1000;
+          acc2 /= 1000;
+          acc3 /= 1000;
+          // ','.join([f"acc{j}[{i}]" for i in range(32) for j in range(4)])
+          // acc0[0], acc0[1], acc0[2], ..... acc3[30], acc3[31]
+          unsigned_char128 packed = __builtin_HEXAGON_V6_vpackhub_sat_128B(__builtin_HEXAGON_V6_vpackwh_sat_128B(acc3, acc2),
+                                                                           __builtin_HEXAGON_V6_vpackwh_sat_128B(acc1, acc0));
+          packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
+          packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
+          // acc0[0], acc1[0], acc2[0], ..... acc2[31], acc3[31]
+          *((unsigned_char128*)((data0+(alu0+alu1)))) = packed;
+        }
+      }
+    }"""
+    self.test_prequant_gemm_intacc(np.uint8, np.int8, src)
+
+  def test_prequant_gemm_intacc_32(self):
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UNROLL, axis=0, arg=0)]
+    self.test_prequant_gemm_intacc(np.uint8, np.int8, N=32, opts=opts)
+  def test_prequant_gemm_intacc_128(self): self.test_prequant_gemm_intacc(np.uint8, np.int8, N=128)
+  def test_prequant_gemm_intacc_256(self): self.test_prequant_gemm_intacc(np.uint8, np.int8, N=256)
+  def test_prequant_gemm_intacc(self, xi=np.uint8, wi=np.uint8, replace_src=None, N=512, clip=True, opts=None):
+    X = Tensor(m1:=(np.random.uniform(0, 255, size=(N,N)).astype(xi))).realize()
+    W = Tensor(m2:=(np.random.uniform(0, 255, size=(N,N)).astype(wi))).realize()
+    tg_dtype = dtypes.int8 if xi == np.int8 else dtypes.uint8
+    out = (X.int().matmul(W.int())//1000)
+    if clip: out = out.clip(tg_dtype.min, tg_dtype.max)
+    out = out.cast(tg_dtype)
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
+    sexec(out, opts, replace_src, run_count=1)
+    tout = out.numpy()
+    mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000)
+    if clip: mout = mout.clip(tg_dtype.min, tg_dtype.max)
+    mout = mout.astype(xi)
+    print(tout)
+    print(mout)
+    np.testing.assert_equal(tout, mout)
+
+  def test_prequant_gemm_intacc_wi(self): self.test_prequant_gemm_intacc(wi=np.int8)
+  def test_prequant_gemm_intacc_xiwi(self): self.test_prequant_gemm_intacc(xi=np.int8, wi=np.int8)
+  def test_prequant_gemm_intacc_xiwi_noclip(self): self.test_prequant_gemm_intacc(xi=np.int8, wi=np.int8, clip=False)
+
+  def test_prequant_gemv(self):
+    N = 2048
+    X = Tensor(np.random.uniform(0, 255, size=(1,N)).astype(np.uint8)).realize()
+    W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8)).realize()
+    #out = X.cast(dtypes.int) @ W.cast(dtypes.int)
+    #out = X @ W
+    out = X.matmul(W, dtype=X.dtype)
+    opts = [Opt(op=OptOps.UPCAST, axis=0, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
+    sexec(out, opts)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/backend/test_randomness.py

@@ -0,0 +1,262 @@
+import unittest, math
+
+from tinygrad import dtypes, Tensor, Device
+from tinygrad.helpers import getenv, DEV
+from tinygrad.codegen import to_program
+
+from tinygrad.uop.ops import Ops
+from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
+from tinygrad.renderer.isa.x86 import X86Renderer
+from test.helpers import not_support_multi_device, needs_second_gpu
+from test.unit.test_randomness import equal_distribution, normal_test
+
+import numpy as np
+import torch
+from hypothesis import given, settings, strategies as strat
+
+settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+settings.load_profile("my_profile")
+
+class TestRandomness(unittest.TestCase):
+  def test_rand(self):
+    self.assertFalse(normal_test(Tensor.rand))
+    self.assertTrue(equal_distribution(Tensor.rand, torch.rand, lambda x: np.random.rand(*x)))
+
+  def test_rand_is_lazy(self):
+    Tensor.manual_seed(0)
+    r1 = Tensor.rand(10)
+    self.assertFalse(r1.uop.is_realized, "rand should be lazy - tensor should not be realized")
+    counter = Tensor._device_rng_counters[Device.DEFAULT]
+    self.assertFalse(counter.uop.is_realized, "rand should be lazy - counter should not be realized")
+    # second rand triggers assign path
+    r2 = Tensor.rand(10)
+    self.assertFalse(r2.uop.is_realized, "rand should be lazy - tensor should not be realized after second rand")
+    self.assertFalse(counter.uop.is_realized, "rand should be lazy - counter should not be realized after second rand")
+    Tensor.realize(r1, r2)
+    self.assertTrue(r1.uop.is_realized, "tensor should be realized after .realize()")
+    self.assertTrue(r2.uop.is_realized, "tensor should be realized after .realize()")
+
+  @unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need float16 support")
+  def test_rand_float16(self):
+    N = 128
+    x = Tensor.rand((2, N, N), dtype=dtypes.float16)
+    assert x.dtype == dtypes.float16
+    nx = x.numpy()
+    # seed dependant, check output range is [0, 1)
+    assert nx[nx == 1].size == 0
+    assert nx[nx == 0].size > 0
+    equal_distribution(lambda *x: Tensor.rand(*x, dtype=dtypes.float16), torch.rand, lambda x: np.random.rand(*x), shape=(2, N, N))
+
+  @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "gpuocelot doesn't support certain ops needed for threefry")
+  def test_threefry_against_reference(self):
+    Tensor.manual_seed(1337)
+
+    # reference generated using
+    """
+    key0 = 1337
+    key1 = 0
+    values = jax.extend.random.threefry_2x32((np.uint32(key1), np.uint32(key0)), np.arange(20, dtype=np.uint32))
+    print(f"[{', '.join(f'{v}' for v in values)}]")
+    """
+    jr = np.array([2221762175, 1752107825, 653745012, 1967534793, 1395205442, 3840423848, 2159346757,
+                   603508235, 3319473678, 3363866483, 3544324138, 1436466838, 2169858556, 2570072943,
+                   2387150698, 3678370550, 2911697663, 403244401, 2560861638, 1692360114])
+
+    counts = Tensor.arange(20, dtype=dtypes.uint32)
+    counts0, counts1 = counts.chunk(2)
+    r = Tensor._threefry_random_bits(Tensor([0, 1337], dtype='uint32'), counts0, counts1).numpy()
+
+    np.testing.assert_allclose(jr, r)
+
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "PTX and NIR use pointer arithmetic")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "X86 callee saved registers have ulong dtype")
+  def test_threefry_doesnt_use_long(self):
+    linear = Tensor.rand(20).schedule_linear()
+    for call in linear.src:
+      ast = call.src[0]
+      if ast.op is Ops.SINK:
+        prg = to_program(ast, renderer=Device[Device.DEFAULT].renderer)
+        for u in tuple(prg.src[2].src):
+          self.assertNotIn(u.dtype, {dtypes.long, dtypes.ulong}, msg=f"long found in {prg.arg.name}")
+
+  def test_threefry_against_reference_full(self):
+    Tensor.manual_seed(1337)
+
+    # reference generated using
+    """
+    key0 = 1337
+    key1 = int.from_bytes(hashlib.sha256(int(0).to_bytes(4)).digest(), "big") & 0xffffffff
+    # derive new key for the counter offset (c_low=0, c_high=0 for first call)
+    new_key_values = jax.extend.random.threefry_2x32((np.uint32(key1), np.uint32(key0)), np.array([0, 0], dtype=np.uint32))
+    new_key = (np.uint32(new_key_values[0]), np.uint32(new_key_values[1]))
+    values = jax.extend.random.threefry_2x32(new_key, np.arange(20, dtype=np.uint32))
+    values = (values >> (32 - 23)) | np.array(1, dtype=np.float32).view(np.uint32)
+    values = values.view(np.float32) - 1
+    print(f"[{', '.join(f'{v}' for v in values)}]")
+    """
+    jr = np.array([0.45735931396484375, 0.6311527490615845, 0.15571284294128418, 0.8149417638778687, 0.7862188816070557,
+                   0.8008807897567749, 0.568588376045227, 0.9852620363235474, 0.42314577102661133, 0.9811755418777466,
+                   0.38059568405151367, 0.09186363220214844, 0.9497315883636475, 0.5826880931854248, 0.3796330690383911,
+                   0.5610522031784058, 0.16122901439666748, 0.3732343912124634, 0.9795231819152832, 0.3280656337738037], dtype=np.float32)
+    r = Tensor.rand(20).numpy()
+    np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
+
+    # next 20 (c_low=20, c_high=0)
+    jr = np.array([0.09199333190917969, 0.9130761623382568, 0.7048608064651489, 0.22254979610443115, 0.0014830827713012695,
+                   0.37023448944091797, 0.7790107727050781, 0.7484984397888184, 0.7524604797363281, 0.19875383377075195,
+                   0.48537540435791016, 0.10002851486206055, 0.5369305610656738, 0.3294715881347656, 0.5246957540512085,
+                   0.7659651041030884, 0.7949080467224121, 0.34988296031951904, 0.9798505306243896, 0.2599533796310425], dtype=np.float32)
+    r = Tensor.rand(20).numpy()
+    np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
+
+    # next 10 (c_low=40, c_high=0)
+    jr = np.array([0.3198714256286621, 0.7984923124313354, 0.320881724357605, 0.4716068506240845, 0.7323365211486816,
+                   0.9663800001144409, 0.13873648643493652, 0.16062307357788086, 0.49300849437713623, 0.10077548027038574], dtype=np.float32)
+    r = Tensor.rand(10).numpy()
+    np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
+
+  @needs_second_gpu
+  @unittest.skipIf(not_support_multi_device(), "no multi")
+  def test_threefry_tensors_cnt(self):
+    Tensor.manual_seed(1337)
+
+    Tensor.rand(20).realize()
+
+    assert len(Tensor._device_rng_counters) == 1
+    assert len(Tensor._device_seeds) == 1
+
+    Tensor.rand(20, device=f"{Device.DEFAULT}:1").realize()
+
+    assert len(Tensor._device_rng_counters) == 2
+    assert len(Tensor._device_seeds) == 2
+
+    Tensor.manual_seed(2)
+
+    assert len(Tensor._device_rng_counters) == 0
+    assert len(Tensor._device_seeds) == 0
+
+  @needs_second_gpu
+  @unittest.skipIf(not_support_multi_device(), "no multi")
+  def test_threefry_same_kernels(self):
+    Tensor.manual_seed(0)
+
+    Tensor.rand(1).realize()
+
+    s = Tensor.rand(20).schedule_linear().src
+    s2 = Tensor.rand(20).schedule_linear().src
+
+    assert len(s) == len(s2), f"{len(s)} != {len(s2)}"
+    for x,y in zip(s, s2):
+      if not (x.src[0] == y.src[0]):
+        print(f"{x.src[0]} != {y.src[0]}")
+
+    Tensor.rand(1, device=f"{Device.DEFAULT}:1").realize()
+
+    s3 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule_linear().src
+    s4 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule_linear().src
+
+    assert len(s3) == len(s4), f"{len(s3)} != {len(s4)}"
+    assert len(s2) == len(s4), f"{len(s)} != {len(s3)}"
+    for x,y in zip(s3, s4):
+      if not (x.src[0] == y.src[0]):
+        print(f"{x.src[0]} != {y.src[0]}")
+
+  @unittest.skipUnless(dtypes.bfloat16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need bfloat16 support")
+  def test_rand_bfloat16(self):
+    N = 128
+    x = Tensor.rand((2, N, N), dtype=dtypes.bfloat16)
+    assert x.dtype == dtypes.bfloat16
+    nx = x.numpy()
+    assert nx[nx == 1].size == 0
+    assert nx[nx == 0].size > 0
+    equal_distribution(lambda *x: Tensor.rand(*x, dtype=dtypes.bfloat16).float(), torch.rand, lambda x: np.random.rand(*x), shape=(2, N, N))
+
+  def test_rand_like(self):
+    empty = Tensor.empty((80, 44))
+    rand = Tensor.rand_like(empty)
+    assert rand.shape == empty.shape
+    assert rand.dtype == empty.dtype
+    assert rand.device == empty.device
+
+  def test_randn_like(self):
+    empty = Tensor.empty((80, 44))
+    rand = Tensor.randn_like(empty)
+    assert rand.shape == empty.shape
+    assert rand.dtype == empty.dtype
+    assert rand.device == empty.device
+
+  def test_rand_like_zero_shape(self):
+    empty = Tensor.empty(0, 20)
+    rand = Tensor.rand_like(empty)
+    assert rand.shape == empty.shape
+    assert rand.dtype == empty.dtype
+    assert rand.device == empty.device
+
+  def test_rand_like_more_dims(self):
+    empty = Tensor.empty((1, 2, 3, 4, 5, 6))
+    rand = Tensor.rand_like(empty)
+    assert rand.shape == empty.shape
+    assert rand.dtype == empty.dtype
+    assert rand.device == empty.device
+
+  def test_rand_like_dtype(self):
+    empty = Tensor.empty((80, 44), dtype=dtypes.float16)
+    rand = Tensor.rand_like(empty)
+    assert rand.shape == empty.shape
+    assert rand.dtype == empty.dtype
+    assert rand.device == empty.device
+
+    empty = Tensor.empty((80, 44))
+    rand = Tensor.rand_like(empty, dtype=dtypes.float16)
+    assert rand.shape == empty.shape
+    assert rand.dtype == dtypes.float16
+    assert rand.device == empty.device
+
+  def test_randn_like_dtype(self):
+    empty = Tensor.empty((80, 44), dtype=dtypes.float16)
+    rand = Tensor.randn_like(empty)
+    assert rand.shape == empty.shape
+    assert rand.dtype == empty.dtype
+    assert rand.device == empty.device
+
+    empty = Tensor.empty((80, 44))
+    rand = Tensor.randn_like(empty, dtype=dtypes.float16)
+    assert rand.shape == empty.shape
+    assert rand.dtype == dtypes.float16
+    assert rand.device == empty.device
+
+  def test_randn_device(self):
+    self.assertEqual(Tensor.randn(3,3,device="CPU").device, "CPU")
+
+  @given(strat.sampled_from([dtypes.float, dtypes.float16, dtypes.bfloat16]))
+  def test_randn_finite(self, default_float):
+    if default_float not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
+    old_default_float = dtypes.default_float
+    # low precision can result in inf from randn
+    dtypes.default_float = default_float
+    t = Tensor.randn(64, 64)
+    mx = t.max().numpy().item()
+    mn = t.min().numpy().item()
+    print(f"testing with {default_float=}")
+    assert math.isfinite(mx), mx
+    assert math.isfinite(mn), mn
+    dtypes.default_float = old_default_float
+
+  def test_random_counter_overflow(self):
+    device = Device.DEFAULT
+    Tensor.manual_seed(1337)
+    Tensor.rand(1).realize()
+
+    Tensor._device_rng_counters[device].assign(Tensor([dtypes.uint32.max - 5, 0], device=device, dtype=dtypes.uint32)).realize()
+
+    Tensor.rand(10).realize()
+    c = Tensor._device_rng_counters[device].numpy()
+    np.testing.assert_allclose(c, [4, 1])
+
+    Tensor.rand(10).realize()
+    c = Tensor._device_rng_counters[device].numpy()
+    np.testing.assert_allclose(c, [14, 1])
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/backend/test_rangeify.py

@@ -0,0 +1,222 @@
+import unittest
+from tinygrad import Tensor, nn, Device
+from tinygrad.helpers import Context, GlobalCounters, getenv, PCONTIG, DEBUG
+from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops
+from tinygrad.codegen.opt import OptOps, Opt
+from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
+
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
+class TestDoubleMatmul(unittest.TestCase):
+  def setUp(self):
+    with Context(DEBUG=0):
+      self.a, self.b, self.c = [Tensor.randn(16, 16).contiguous().realize() for _ in range(3)]
+      self.ref = (self.a @ self.b @ self.c).realize()
+
+  def _test(self, opts):
+    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
+      out = (self.a @ self.b @ self.c).contiguous(arg=opts).realize()
+
+    with Context(DEBUG=0):
+      err = (out-self.ref).square()
+      self.assertLess(err.max().item(), 1e-4)
+      self.assertLess(err.mean().item(), 1e-6)
+
+  def test_baseline(self): self._test(())
+  def test_upcast_0(self): self._test((Opt(OptOps.UPCAST, 0, 4),))
+  def test_upcast_1(self): self._test((Opt(OptOps.UPCAST, 1, 4),))
+  def test_upcast_2(self): self._test((Opt(OptOps.UPCAST, 2, 4),))
+  def test_upcast_01(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4)))
+  def test_upcast_01_mismatch(self): self._test((Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 4)))
+  def test_upcast_02(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 2, 4)))
+  def test_upcast_12(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 2, 4)))
+
+  def test_unroll_0(self): self._test((Opt(OptOps.UNROLL, 0, 4),))
+  def test_unroll_1(self): self._test((Opt(OptOps.UNROLL, 1, 4),))
+  def test_unroll_01(self): self._test((Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
+
+  def test_upcast_0_unroll_0(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)))
+  def test_upcast_1_unroll_0(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)))
+  def test_upcast_2_unroll_0(self): self._test((Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 0, 4)))
+
+  def test_upcast_0_unroll_1(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
+  def test_upcast_1_unroll_1(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 1, 4)))
+  def test_upcast_2_unroll_1(self): self._test((Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 1, 4)))
+
+  def test_upcast_1_unroll_1_small(self): self._test((Opt(OptOps.UPCAST, 1, 2), Opt(OptOps.UNROLL, 1, 2)))
+  def test_upcast_1_unroll_1_rev(self): self._test((Opt(OptOps.UNROLL, 1, 2), Opt(OptOps.UPCAST, 1, 2)))
+
+  def test_upcast_01_unroll_01(self):
+    self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
+  def test_upcast_12_unroll_01(self):
+    self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
+
+class TestRangeifyAssign(unittest.TestCase):
+  def test_assign_permuted(self):
+    A = Tensor.empty(4, 4, dtype='int')
+    B = Tensor.arange(16).reshape(4,4)
+    ret = A.permute(1,0).assign(B)
+    lst = ret.tolist()
+    lst2 = A.tolist()
+    lst3 = B.tolist()
+    print(lst)
+    print(lst2)
+    print(lst3)
+    self.assertListEqual(lst, lst3)
+    self.assertListEqual(lst2, B.permute(1, 0).tolist())
+
+class TestRangeifyEdgeCase(unittest.TestCase):
+  def test_matmul_relu_cat(self):
+    a = Tensor.ones(100, 512).contiguous().realize()
+    c = Tensor.ones(1, 512).contiguous().realize()
+    cm = Tensor.ones(512, 512)
+    c = c @ cm
+    c = c.relu()
+
+    res = Tensor.cat(a, c, dim=0)
+    self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
+
+  def test_pcontig_multi_gather(self):
+    # regression test: local bufferize must have device set for const_like to work
+    with Context(PCONTIG=2):
+      # NOTE: with uint type, this will become a long and fail on WEBGPU
+      forest = Tensor(list(range(8)), dtype='int')
+      idx = Tensor([0, 0], dtype='int')
+      node_val = forest.gather(0, idx)
+      idx2 = idx * 2 + 1
+      node_val2 = forest.gather(0, idx2)
+      result = (node_val + node_val2).numpy()
+    self.assertEqual(result.tolist(), [1, 1])
+
+if getenv("BIG") > 2:
+  # llama 8B (8192)
+  BS, HEADS, SEQLEN, EMB = 4, 32, 8192, 128
+elif getenv("BIG") > 1:
+  # llama 8B
+  BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
+elif getenv("BIG") > 0:
+  # bigger
+  BS, HEADS, SEQLEN, EMB = 4, 32, 128, 128
+else:
+  BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
+
+def fa():
+  Tensor.manual_seed(1337)
+  with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
+  GlobalCounters.reset()
+  return q.scaled_dot_product_attention(k, v)
+
+def fa_bw():
+  Tensor.manual_seed(1337)
+  with Context(DEBUG=0):
+    q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
+    attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
+    attn_output.weight.realize()
+    target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
+
+  GlobalCounters.reset()
+  attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward()
+  attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1)
+  out = attn_output(attn)
+  loss = (out - target).square().mean()
+  loss.backward()
+  #ret = [out, Tensor.stack(q.grad, k.grad, v.grad, dim=-1)]
+  #ret = [out, Tensor.stack(q.grad, k.grad, dim=-1), v.grad]
+  ret = [out, q.grad, k.grad, v.grad]
+  Tensor.realize(*ret)
+  return ret
+
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
+class TestPcontig(unittest.TestCase):
+  def test_flash_attention_bw(self):
+    with Context(PCONTIG=max(2, PCONTIG.value), DEBUG=2):
+      grads = fa_bw()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+
+    with Context(PCONTIG=0, DEBUG=2):
+      cmp_grads = fa_bw()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+
+    with Context(DEBUG=0):
+      mses = [((x-y)**2).sum().item() for x,y in zip(grads, cmp_grads)]
+    mse = sum(mses)
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-6)
+
+  def test_flash_attention(self, opts=None):
+    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
+      ret = fa().realize() if opts is None else fa().contiguous(arg=opts).realize()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+    with Context(DEBUG=2):
+      cmp = fa().realize()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+    with Context(DEBUG=0):
+      mse = ((cmp-ret)**2).sum().item()
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-6)
+
+  def test_flash_attention_opt(self):
+    opts = ()
+    # columns in top matrix
+    opts += (Opt(OptOps.UPCAST, 0, 4),)
+    # columns in bottom matrix
+    opts += (Opt(OptOps.UPCAST, 3, 4),)
+    # rows in all the matrix
+    opts += (Opt(OptOps.UPCAST, 4, 4),)
+    self.test_flash_attention(opts)
+
+# contiguous + reduce can support ranges?
+
+@unittest.skip("pm_rangeify no longer exists. test this in a different way")
+class TestRangeifyPM(unittest.TestCase):
+  def setUp(self): self.base = Tensor.empty(10*10).reshape(10, 10).contiguous()
+  def assert_same(self, a, b):
+    def run_pm_rangeify(t:Tensor):
+      from tinygrad.schedule.rangeify import pm_rangeify, RangeifyContext
+      sink = t.uop.sink()
+      pm_realize = PatternMatcher([(UPat(Ops.CONTIGUOUS, name="x"), lambda x: x.replace(op=Ops.REALIZE))])
+      sink = graph_rewrite(sink, pm_realize)
+      return graph_rewrite(sink, pm_rangeify, ctx=RangeifyContext())
+    self.assertIs(run_pm_rangeify(a.contiguous()), run_pm_rangeify(b.contiguous()))
+
+  def test_nothing_match(self):
+    a = self.base.pad(((0,0),(0,1)))
+    b = self.base.pad(((0,0),(0,1)))
+    self.assert_same(a, b)
+
+  def test_reshape_match(self):
+    a = self.base
+    b = self.base.reshape(100).reshape(10, 10)
+    self.assert_same(a, b)
+
+  def test_permute_reshape_match(self):
+    a = self.base
+    b = self.base.permute(1,0).reshape(100).reshape(10, 10).permute(1,0)
+    self.assert_same(a, b)
+
+  def test_padded_permute_match(self):
+    a = self.base.pad(((0,0),(0,1)))
+    b = self.base.permute(1,0).pad(((0,1),(0,0))).permute(1,0)
+    self.assert_same(a, b)
+
+  @unittest.expectedFailure
+  def test_padded_reshape_match(self):
+    a = self.base.pad(((0,0),(0,1)))
+    b = self.base.reshape(100).reshape(10, 10).pad(((0,0),(0,1)))
+    self.assert_same(a, b)
+
+  @unittest.expectedFailure
+  def test_padded_permute_reshape_match(self):
+    a = self.base.pad(((0,0),(0,1)))
+    b = self.base.permute(1,0).reshape(100).reshape(10, 10).pad(((0,1),(0,0))).permute(1,0)
+    self.assert_same(a, b)
+
+  # why is this failing?
+  @unittest.expectedFailure
+  def test_cross_pad_match(self):
+    a = self.base.pad(((0,0),(0,1))).pad(((0,1),(0,0)))
+    b = self.base.pad(((0,1),(0,0))).pad(((0,0),(0,1)))
+    self.assert_same(a, b)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_renderer_failures.py

@@ -0,0 +1,108 @@
+import unittest
+import numpy as np
+from tinygrad.device import Device
+from tinygrad.dtype import dtypes, ConstType
+from tinygrad.engine.realize import run_linear
+from tinygrad.codegen import to_program
+from tinygrad.helpers import prod
+from tinygrad.renderer.cstyle import CStyleLanguage
+from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.wgsl import WGSLRenderer
+from tinygrad.runtime.ops_python import PythonRenderer
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu
+from tinygrad.tensor import Tensor, _to_np_dtype
+
+def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):
+  for x in inputs: x.realize()
+  sz = 1 if local_size is None else prod(local_size)
+  outs = [UOp.new_buffer(Device.DEFAULT, sz, u.src[1].dtype) for u in sink.src if u.op is Ops.STORE]
+  for u in outs: u.buffer.allocate().copyin(np.zeros(sz, dtype=_to_np_dtype(u.dtype)).data)
+  run_linear(UOp(Ops.LINEAR, src=(sink.call(*outs, *(x.uop.base for x in inputs)),)))
+  return [u.buffer.numpy() for u in outs]
+
+def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
+  dtype = alu_src_uops[0].dtype
+  a = UOp.param(0, dtype.ptr())
+  b = UOp.param(1, dtype.ptr())
+  idx = UOp.const(dtypes.int, 0)
+  ld = b.index(idx)
+  alu = ld.alu(alu_op, *alu_src_uops)
+  store = UOp.store(a.index(idx), alu)
+  return _test_uop_result([Tensor([input_val])], UOp(Ops.SINK, dtypes.void, (store,), arg=KernelInfo()))[0]
+
+class TestRendererFailures(unittest.TestCase):
+  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
+  def test_gated_store_with_alu(self):
+    a = UOp.param(0, dtypes.int.ptr())
+    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
+    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
+    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
+    ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
+    np.testing.assert_equal(ret, [0, 1, 1, 1])
+
+  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
+  def test_gated_store_with_alu_2d(self):
+    a = UOp.param(0, dtypes.int.ptr())
+    gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
+    gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
+    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
+    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
+    ret = _test_uop_result([], sink, local_size=[4, 2, 1])[0]
+    np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
+
+@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage), "uops are for cstyle")
+class TestCStyleFailures(unittest.TestCase):
+  def test_inline_const_alu(self):
+    # CPU doesn't use the max function
+    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.int.min+1))
+    self.assertEqual(ret[0], 1)
+
+  def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):
+    dtype = "bool" if op in (Ops.OR, Ops.XOR, Ops.AND) else None
+    ret = Tensor.empty(1, dtype=dtype)
+    for _ in range(5): ret = python_alu[op](ret, Tensor.empty(1, dtype=dtype))
+    linear = ret.schedule_linear()
+    assert len(linear.src) == 1
+    src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[3].arg
+    self.assertEqual("("*5 not in src, should_strip_paren)
+
+  def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD)
+  def test_repeat_mul(self): self._test_src_strip_paren(Ops.MUL)
+  def test_repeat_xor(self): self._test_src_strip_paren(Ops.XOR)
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "wgsl ends up with '(' * 5")
+  def test_repeat_or(self): self._test_src_strip_paren(Ops.OR)
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "wgsl ends up with '(' * 5")
+  def test_repeat_and(self): self._test_src_strip_paren(Ops.AND)
+  def test_repeat_sub(self): self._test_src_strip_paren(Ops.SUB, should_strip_paren=False)
+
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer")
+class TestWGSLFailures(unittest.TestCase):
+  def test_multiply_infinity(self):
+    # multiplying a positive constant by infinity should return infinity
+    # WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer
+    ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(dtypes.float32, float("inf")))
+    self.assertEqual(ret[0], float("inf"))
+
+@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "tests for ptx renderer")
+class TestPTXFailures(unittest.TestCase):
+  @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
+  def test_gated_store_with_if(self):
+    a = UOp.param(0, dtypes.int.ptr())
+    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
+    val = UOp.const(dtypes.int, 1)
+    if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,))
+    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, if_uop), val))
+    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
+    ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
+    np.testing.assert_equal(ret, [0, 1, 1, 1])
+
+  @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
+  def test_gated_define_acc_with_half_dtype(self):
+    a = Tensor.randn(32, 32, dtype=dtypes.half).realize()
+    b = Tensor.randn(34, 32, dtype=dtypes.half).realize()
+    result = a.pad((1,1)).matmul(b, dtype=dtypes.half).numpy()
+    reference = a.pad((1,1)).matmul(b, dtype=dtypes.float).numpy()
+    np.testing.assert_allclose(result, reference, atol=1e-2, rtol=1e-2)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_setitem.py

@@ -0,0 +1,378 @@
+import unittest
+from tinygrad import Tensor, TinyJit, Variable, dtypes, Device
+from tinygrad.helpers import Context
+import numpy as np
+
+class TestSetitem(unittest.TestCase):
+  def test_simple_setitem(self):
+    cases = (
+      ((6,6), (slice(2,4), slice(3,5)), Tensor.ones(2,2)),
+      ((6,6), (slice(2,4), slice(3,5)), Tensor([1.,2.])),
+      ((6,6), (slice(2,4), slice(3,5)), 1.0),
+      ((6,6), (3, 4), 1.0),
+      ((6,6), (3, None, 4, None), 1.0),
+      ((4,4,4,4), (Ellipsis, slice(1,3), slice(None)), Tensor(4.0)),
+      ((4,4,4,4), (Ellipsis, slice(1,3)), 4),
+      ((4,4,4,4), (2, slice(1,3), None, 1), 4),
+      ((4,4,4,4), (slice(1,3), slice(None), slice(0,4,2)), 4),
+      ((4,4,4,4), (slice(1,3), slice(None), slice(None), slice(0,3)), 4),
+      ((6,6), (slice(1,5,2), slice(0,5,3)), 1.0),
+      ((6,6), (slice(5,1,-2), slice(5,0,-3)), 1.0),
+    )
+    for shp, slc, val in cases:
+      t = Tensor.zeros(shp).contiguous()
+      t[slc] = val
+      n = np.zeros(shp)
+      n[slc] = val.numpy() if isinstance(val, Tensor) else val
+      np.testing.assert_allclose(t.numpy(), n)
+
+  def test_padded_setitem(self):
+    t = Tensor.arange(10)
+    t[4:1:-2] = 11
+    self.assertListEqual(t.tolist(), [0, 1, 11, 3, 11, 5, 6, 7, 8, 9])
+
+  def test_setitem_inplace_mul(self):
+    t = Tensor.arange(10).clone().realize()
+    t[:3] *= 10
+    self.assertListEqual(t.tolist(), [0, 10, 20, 3, 4, 5, 6, 7, 8, 9])
+
+  @unittest.skip("crashed in LLVM CI")
+  def test_setitem_fancy_on_unrealized_view(self):
+    # fancy indexing setitem on unrealized SHRINK view (triggered infinite loop in graph_rewrite)
+    base = Tensor.arange(20, dtype=dtypes.float).reshape(4, 5).clone().realize()
+    sub = base[1:3]
+    flat = sub.reshape(sub.numel()).contiguous()
+    idx = Tensor([0, 3, 7, 9])
+    flat[idx] = Tensor([99, 98, 97, 96], dtype=dtypes.float)
+    sub.assign(flat.reshape(2, 5))
+    np.testing.assert_allclose(sub.numpy(), [[99, 6, 7, 98, 9], [10, 11, 97, 13, 96]])
+
+  def test_setitem_dtype(self):
+    for dt in (dtypes.int, dtypes.float, dtypes.bool):
+      for v in (5., 5, True):
+        t = Tensor.ones(6,6, dtype=dt).contiguous()
+        t[1] = v
+        self.assertEqual(t.dtype, dt)
+
+  def test_setitem_dtype_mismatch(self):
+    t = Tensor.zeros(6, dtype=dtypes.float).contiguous().realize()
+    with self.assertRaises(RuntimeError): t[2:4] = Tensor([1, 2], dtype=dtypes.int)
+
+  def test_setitem_chained_indexing(self):
+    # N[i][j] must work the same as N[i, j]
+    N1 = Tensor.zeros((3, 3)).contiguous().realize()
+    N1[1, 2] = 5
+    N2 = Tensor.zeros((3, 3)).contiguous().realize()
+    N2[1][2] = 5
+    np.testing.assert_equal(N1.numpy(), N2.numpy())
+
+  def test_setitem_detach(self):
+    # setitem on detached tensor should work
+    t = Tensor.zeros((3, 3)).contiguous().realize()
+    t.detach()[1, 2] = 5
+    self.assertEqual(t[1, 2].item(), 5.0)
+
+  def test_setitem_permute(self):
+    # setitem on permuted tensor should modify original
+    t = Tensor.zeros((2, 3)).contiguous().realize()
+    t.T[1, 0] = 5  # t.T is (3, 2), so [1, 0] maps to t[0, 1]
+    self.assertEqual(t[0, 1].item(), 5.0)
+
+  def test_setitem_flip(self):
+    # setitem on flipped tensor should modify original
+    t = Tensor.zeros((3,)).contiguous().realize()
+    t[::-1][0] = 5  # flip, then set first element (which is last in original)
+    self.assertEqual(t[2].item(), 5.0)
+
+  def test_setitem_inplace_operator(self):
+    t = Tensor.arange(4).reshape(2, 2).contiguous()
+    t[1] += 2
+    np.testing.assert_allclose(t.numpy(), [[0, 1], [4, 5]])
+
+    t = Tensor.arange(4).reshape(2, 2).contiguous()
+    t[1] -= 1
+    np.testing.assert_allclose(t.numpy(), [[0, 1], [1, 2]])
+
+    t = Tensor.arange(4).reshape(2, 2).contiguous()
+    t[1] *= 2
+    np.testing.assert_allclose(t.numpy(), [[0, 1], [4, 6]])
+
+    # NOTE: have to manually cast setitem target to least_upper_float for div
+    t = Tensor.arange(4, dtype=dtypes.float).reshape(2, 2).contiguous()
+    t[1] /= 2
+    np.testing.assert_allclose(t.numpy(), [[0, 1], [1, 1.5]])
+
+    t = Tensor.arange(4).reshape(2, 2).contiguous()
+    t[1] **= 2
+    np.testing.assert_allclose(t.numpy(), [[0, 1], [4, 9]])
+
+    t = Tensor.arange(4).reshape(2, 2).contiguous()
+    t[1] ^= 5
+    np.testing.assert_allclose(t.numpy(), [[0, 1], [7, 6]])
+
+  def test_setitem_consecutive_inplace_operator(self):
+    t = Tensor.arange(4).reshape(2, 2).contiguous()
+    t[1] += 2
+    t[1] -= 1
+    np.testing.assert_allclose(t.numpy(), [[0, 1], [3, 4]])
+
+  def test_setitem_overlapping_indices(self):
+    t = Tensor([1,2,3,4])
+    # regular overlapping indices
+    t[[1,1]] = Tensor([5,6])
+    np.testing.assert_allclose(t.numpy(), [1,6,3,4])
+
+    # overlapping indices with zero value overlapped
+    t[[1,1]] = Tensor([0,1])
+    np.testing.assert_allclose(t.numpy(), [1,1,3,4])
+
+  def test_setitem_overlapping_indices_with_0(self):
+    t = Tensor([1,2,3,4])
+    t[[1,1]] = Tensor([1,0])
+    np.testing.assert_allclose(t.numpy(), [1,0,3,4])
+
+  def test_setitem_with_1_in_shape(self):
+    t = Tensor([[1],[2],[3]])
+    t[[0,0]] = Tensor([[1],[2]])
+    np.testing.assert_allclose(t.numpy(), [[2],[2],[3]])
+
+  def test_fancy_setitem(self):
+    t = Tensor.zeros(6,6).contiguous()
+    t[[1,2], [3,2]] = 3
+    n = np.zeros((6,6))
+    n[[1,2], [3,2]] = 3
+    np.testing.assert_allclose(t.numpy(), n)
+
+  def test_simple_jit_setitem(self):
+    @TinyJit
+    def f(t:Tensor, a:Tensor):
+      t[2:4, 3:5] = a
+      # NOTE: without return t or an explicit realize, it's lazy and not captured
+      return t
+
+    for i in range(1, 6):
+      t = Tensor.zeros(6, 6).contiguous().realize()
+      a = Tensor.full((2, 2), fill_value=i, dtype=dtypes.float).contiguous()
+      f(t, a)
+
+      n = np.zeros((6, 6))
+      n[2:4, 3:5] = np.full((2, 2), i)
+      np.testing.assert_allclose(t.numpy(), n)
+
+  def test_jit_setitem_variable_offset(self):
+    with Context(CHECK_OOB=0):
+      @TinyJit
+      def f(t:Tensor, a:Tensor, v:Variable):
+        t.shrink(((v,v+1), None)).assign(a).realize()
+
+      t = Tensor.zeros(6, 6).contiguous().realize()
+      n = np.zeros((6, 6))
+
+      for i in range(6):
+        v = Variable("v", 0, 6).bind(i)
+        a = Tensor.full((1, 6), fill_value=i+1, dtype=dtypes.float).contiguous()
+        n[i, :] = i+1
+        f(t, a, v)
+        np.testing.assert_allclose(t.numpy(), n)
+      np.testing.assert_allclose(t.numpy(), [[1,1,1,1,1,1],[2,2,2,2,2,2],[3,3,3,3,3,3],[4,4,4,4,4,4],[5,5,5,5,5,5],[6,6,6,6,6,6]])
+
+  def test_setitem_overlapping_inplace1(self):
+    t = Tensor([[3.0], [2.0], [1.0]]).contiguous()
+    t[1:] = t[:-1]
+    self.assertEqual(t.tolist(), [[3.0], [3.0], [2.0]])
+
+  def test_setitem_overlapping_inplace2(self):
+    t = Tensor([[3.0], [2.0], [1.0]]).contiguous()
+    t[:-1] = t[1:]
+    self.assertEqual(t.tolist(), [[2.0], [1.0], [1.0]])
+
+  # TODO: WEBGPU pipeline validation error. this generates (1==gidx0)|(2==gidx0)|(3==gidx0)|(4==gidx0)|(5==gidx0) ...
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU pipeline validation error")
+  def test_setitem_big(self):
+    idx_size, val = 256, 4
+    t = Tensor.arange(0, idx_size+1)
+    idx = Tensor.arange(0, idx_size)
+    t[idx] = val
+    self.assertEqual(t.tolist(), [val]*idx_size+[idx_size])
+
+  def test_setitem_advanced_indexing(self):
+    # Example from https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing
+    t = Tensor.zeros(10,20,30,40,50, dtype=dtypes.int).contiguous()
+    ind_1 = Tensor([5,3,7,8])
+    ind_2 = Tensor([[[0],[1],[2]],[[3],[4],[5]]])
+    v = Tensor.arange(2*3*4*10*30*50).reshape(2,3,4,10,30,50)
+    t[:, ind_1, :, ind_2, :] = v
+    n = np.zeros((10,20,30,40,50), dtype=np.int32)
+    n[:, ind_1.numpy(), :, ind_2.numpy(), :] = v.numpy()
+    np.testing.assert_equal(t.numpy(), n)
+
+  def test_setitem_tensor_int_indexing(self):
+    t = Tensor.zeros(4, 3, dtype=dtypes.int).contiguous()
+    t[Tensor([0, 2]), 0] = Tensor([99, 88], dtype=dtypes.int)
+    n = np.zeros((4, 3), dtype=np.int32)
+    n[[0, 2], 0] = [99, 88]
+    np.testing.assert_equal(t.numpy(), n)
+
+  def test_setitem_tensor_slice_indexing(self):
+    t = Tensor.zeros(4, 3, dtype=dtypes.int).contiguous()
+    t[Tensor([0, 2]), :2] = Tensor([[10, 20], [30, 40]], dtype=dtypes.int)
+    n = np.zeros((4, 3), dtype=np.int32)
+    n[[0, 2], :2] = [[10, 20], [30, 40]]
+    np.testing.assert_equal(t.numpy(), n)
+
+  def test_setitem_2d_tensor_indexing(self):
+    t = Tensor.zeros(2, dtype=dtypes.int).contiguous()
+    index = Tensor([[0, 1], [1,0]])
+    v = Tensor.arange(2*2).reshape(2, 2).contiguous()
+    t[index] = v
+    n = np.zeros((2,), dtype=np.int32)
+    n[index.numpy()] = v.numpy()
+    np.testing.assert_equal(t.numpy(), n)
+
+  def test_setitem_swap_rows(self):
+    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
+    tmp = t[0]
+    t[0] = t[1]
+    t[2] = tmp
+    # NOTE: not [[2, 3], [2, 3], [0, 1]], same with eager
+    np.testing.assert_allclose(t.numpy(), [[2, 3], [2, 3], [2, 3]])
+
+    # eager version
+    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
+    tmp = t[0].realize()
+    t[0] = t[1].realize()
+    t[2] = tmp.realize()
+    np.testing.assert_allclose(t.numpy(), [[2, 3], [2, 3], [2, 3]])
+
+  def test_lazy_sum_between_writes(self):
+    # lazy sums should capture buffer state at the time they were created
+    t = Tensor.zeros(6).contiguous().realize()
+    s0 = t.sum()
+    t[:3].assign(1.0)
+    s1 = t.sum()
+    t[3:].assign(2.0)
+    s2 = t.sum()
+    try:
+      np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [0.0, 3.0, 9.0])
+    except AssertionError:
+      # TODO: broken now, lazy sums all see final buffer state
+      np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [9.0, 9.0, 9.0])
+
+    # eager version
+    t = Tensor.zeros(6).contiguous().realize()
+    s0 = t.sum().realize()
+    t[:3].assign(1.0).realize()
+    s1 = t.sum().realize()
+    t[3:].assign(2.0).realize()
+    s2 = t.sum().realize()
+    np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [0.0, 3.0, 9.0])
+
+  def test_cross_assign_independence(self):
+    # when assigning to two tensors using computations from both,
+    # both assigns should see the OLD values of both tensors
+    a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
+    b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
+    new_a = a + b    # [4, 6, 8, 10]
+    new_b = a * 2    # [0, 2, 4, 6] -- should use OLD a
+    a.assign(new_a)
+    b.assign(new_b)
+    np.testing.assert_allclose(a.numpy(), [4, 6, 8, 10])
+    try:
+      np.testing.assert_allclose(b.numpy(), [0, 2, 4, 6])
+    except AssertionError:
+      # TODO: broken now, new_b sees mutated a
+      np.testing.assert_allclose(b.numpy(), [8, 12, 16, 20])
+
+    # eager version
+    a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
+    b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
+    new_a = (a + b).realize()
+    new_b = (a * 2).realize()
+    a.assign(new_a).realize()
+    b.assign(new_b).realize()
+    np.testing.assert_allclose(a.numpy(), [4, 6, 8, 10])
+    np.testing.assert_allclose(b.numpy(), [0, 2, 4, 6])
+
+  def test_setitem_multiple_disjoint_on_invalid(self):
+    z = Tensor.invalids(10, dtype="int").realize()
+    z[2:5] = 2
+    z[6:7] = 3
+    z.realize()
+    self.assertListEqual(z[2:5].tolist(), [2, 2, 2])
+    self.assertListEqual(z[6:7].tolist(), [3])
+
+class TestWithGrad(unittest.TestCase):
+  def test_basic_setitem_works(self):
+    z = Tensor.rand(8, 8)
+    x = Tensor.rand(8)
+    z[:3] = x
+
+  def test_set_backward(self):
+    z = Tensor.ones(8, 8)
+    x = Tensor.rand(8, 8)
+    z[:] = x
+    z.sum().backward()
+    np.testing.assert_allclose(x.grad.numpy(), np.ones((8, 8)))
+
+  def test_set_nonleaf_backward(self):
+    x = Tensor([1.0, 2.0, 3.0, 4.0])
+    z = x * 2
+    z[:2] = Tensor([10.0, 20.0])
+    z.sum().backward()
+    np.testing.assert_allclose(x.grad.numpy(), [0, 0, 2, 2])
+
+  def test_set_overlapping_backward(self):
+    z = Tensor.zeros(6)
+    x = Tensor.ones(4).contiguous()
+    y = Tensor.ones(4) * 2
+    z[:4] = x
+    z[2:] = y
+    z.sum().backward()
+    np.testing.assert_allclose(x.grad.numpy(), [1, 1, 0, 0])
+    np.testing.assert_allclose(y.grad.numpy(), np.ones(4))
+
+  def test_set_iadd_backward(self):
+    z = Tensor([1.0, 2.0, 3.0, 4.0])
+    x = Tensor([10.0, 20.0])
+    z[:2] += x
+    z.sum().backward()
+    np.testing.assert_allclose(z.grad.numpy(), np.ones(4))
+    np.testing.assert_allclose(x.grad.numpy(), np.ones(2))
+
+  def test_set_used_before_setitem(self):
+    z = Tensor([1.0, 2.0, 3.0, 4.0])
+    _ = z.sum()
+    with self.assertRaises(RuntimeError):
+      z[:2] = Tensor([0.0, 0.0])
+
+  def test_setitem_raises_with_unrealized_downstream(self):
+    x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
+    _y = x * 2.0
+    with self.assertRaises(RuntimeError):
+      x[0] = 99.0
+
+  def test_setitem_raises_on_unrealized_compute_base(self):
+    # y has a compute (unrealized) base; tmp is a view of y. eager: tmp would follow y's mutation. lazy: tmp keeps the old MUL graph.
+    x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
+    y = x * 2.0
+    _tmp = y[:1]
+    with self.assertRaises(RuntimeError):
+      y[0] = 99.0
+
+  def test_setitem_raises_on_aliased_uop(self):
+    # two Tensor objects sharing the exact same unrealized uop. setitem on one updates its uop, the other keeps the stale graph reference.
+    x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
+    y = x * 2.0
+    _z = Tensor(y.uop)
+    with self.assertRaises(RuntimeError):
+      y[0] = 99.0
+
+class TestSetitemLoop(unittest.TestCase):
+  def test_arange(self):
+    N = 10
+    cmp = Tensor.empty(N)
+    for i in range(N): cmp[i] = i
+    self.assertListEqual(Tensor.arange(N).tolist(), cmp.tolist())
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_softmax_fusion.py

@@ -0,0 +1,206 @@
+import unittest
+import numpy as np
+from tinygrad import Tensor, GlobalCounters, Context, Device
+from tinygrad.dtype import DTypeLike, dtypes
+from tinygrad.engine.realize import run_linear
+from tinygrad.helpers import DEBUG, get_single_element
+
+def single_kernel_softmax(x_in:Tensor, axis=-1, dtype:DTypeLike|None=None) -> Tensor:
+  # only support axis =-1
+  x = x_in.reshape(-1, x_in.shape[-1])
+  nr_dim, r_dim = x.shape
+
+  inp = x.reshape(nr_dim, 1, 1, r_dim).expand(nr_dim, r_dim, 1, r_dim)
+  imx = x.reshape(nr_dim, 1, r_dim, 1).expand(nr_dim, r_dim, r_dim, r_dim).max(axis=-2, keepdim=True)
+  m = inp - imx.detach()
+  if dtype is not None: m = m.cast(dtype)
+  e = m.exp()
+  ss = e.sum(axis=-1, keepdim=True)
+
+  inp = x.reshape(nr_dim, r_dim, 1, 1)
+  imx = x.reshape(nr_dim, 1, r_dim, 1).expand(nr_dim, r_dim, r_dim, 1).max(axis=-2, keepdim=True)
+  m = inp - imx.detach()
+  if dtype is not None: m = m.cast(dtype)
+  e = m.exp()
+
+  out = e.div(ss).reshape(x_in.shape)
+  return out
+
+def run_one_schedule_item(out):
+  linear = out.schedule_linear()
+  get_single_element(linear.src)
+  run_linear(linear)
+
+class TestFuse(unittest.TestCase):
+  def _test_fuse(self, fxn, *args, atol=1e-6, allow_multiple=False, **kwargs):
+    GlobalCounters.reset()
+    out_single = fxn(*args, **kwargs)
+    if not allow_multiple: run_one_schedule_item(out_single)
+    np_single = out_single.numpy()
+    GlobalCounters.reset()
+    np_multi = fxn(*args, **kwargs).numpy()
+    np.testing.assert_allclose(np_single, np_multi, atol=atol)
+
+  @unittest.skip("needs RANGEIFY>1")
+  def test_fuse_norm(self):
+    a = Tensor.rand(50,50).realize()
+    self._test_fuse(lambda a: a / a.mean(axis=1), a)
+
+  @unittest.skip("needs RANGEIFY>1")
+  def test_fuse_argmax(self):
+    a = Tensor.rand(50,50).realize()
+    self._test_fuse(lambda a: a.argmax(axis=-1), a)
+
+  @unittest.skip("needs RANGEIFY>1")
+  def test_fuse_softmax(self):
+    a = Tensor.rand(50,50).realize()
+    self._test_fuse(lambda a: a.softmax(axis=-1), a)
+
+  def test_fuse_gemm_softmax(self):
+    a = Tensor.rand(50,50).realize()
+    b = Tensor.rand(50,50).realize()
+    self._test_fuse(lambda a,b: ((a@b).relu()+a).contiguous().softmax(axis=-1), a,b, allow_multiple=True)
+
+  @unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), f"no float16 on {Device.DEFAULT}")
+  @unittest.skip("needs RANGEIFY>1")
+  def test_fuse_softmax_dtype(self):
+    a = Tensor.rand(50,50).realize()
+    self._test_fuse(lambda a: a.softmax(axis=-1, dtype='half'), a, atol=3e-4)
+
+  def test_fuse_arange_eye(self):
+    self._test_fuse(lambda: Tensor.arange(10).reshape(10,1).expand(10,10) == Tensor.arange(10).reshape(1,10).expand(10,10))
+
+  @unittest.skip("needs RANGEIFY>1")
+  def test_double_gemm(self):
+    N = 32
+    with Context(TRACK_MATCH_STATS=0, DEBUG=0):
+      a = (Tensor.rand(N,N)-0.5).realize()
+      b = (Tensor.rand(N,N)-0.5).realize()
+      c = (Tensor.rand(N,N)-0.5).realize()
+    self._test_fuse(lambda a,b,c: a@b@c, a, b, c, atol=1e-5)
+
+  def test_embedding(self):
+    with Context(TRACK_MATCH_STATS=0, DEBUG=0):
+      vocab_sz = 123
+      embed_sz = 16
+      weight = (Tensor.rand(vocab_sz, embed_sz)-0.5).realize()
+      a = Tensor([1, 1, 2, 3]).realize()
+    def embedding(idx:Tensor):
+      arange = Tensor.arange(vocab_sz).unsqueeze(-1)
+      big_shp = idx.shape + (vocab_sz, embed_sz)
+      arange, vals = arange.expand(big_shp), weight.expand(big_shp)
+      idx = idx.reshape(idx.shape+(1, 1)).expand(big_shp)
+      return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
+    self._test_fuse(embedding, a, atol=1e-5)
+
+  @unittest.skip("needs RANGEIFY>1")
+  def test_attention_kernel_count(self):
+    wq = Tensor.empty(32, 32)
+    wk = Tensor.empty(32, 32)
+    wv = Tensor.empty(32, 32)
+    x = Tensor.empty(2, 100, 32)
+    q = (x @ wq).contiguous()
+    k = (x @ wk).contiguous()
+    v = (x @ wv).contiguous()
+    attn = q.scaled_dot_product_attention(k, v)
+    s = attn.schedule_linear()
+    self.assertEqual(len(s.src), 4) # 3 matmul and 1 attention
+
+  @unittest.skip("needs RANGEIFY>1")
+  def test_flash_attention(self):
+    BS = 4
+    HEADS = 2
+    MATDIM = 16
+    EMB = 8
+    with Context(TRACK_MATCH_STATS=0, DEBUG=0):
+      q = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
+      k = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
+      v = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
+    # TODO: OPT is breaking things. NOOPT isn't linearizing
+    with Context(NOOPT=1):
+      self._test_fuse(Tensor.scaled_dot_product_attention, q, k, v, atol=1e-5)
+
+  def test_mismatch_reduce(self):
+    a = Tensor.ones(16, 10).contiguous().realize()
+    b = Tensor.ones(16, 20).contiguous().realize()
+    c = (a.sum(axis=1) + b.sum(axis=1))
+    self.assertListEqual(c.tolist(), [30]*16)
+
+  @unittest.skipUnless(Device.DEFAULT == "METAL", "METAL TC")
+  def test_fuse_and_tc_opt(self):
+    A = Tensor.randn(8, 8).realize()
+    B = Tensor.randn(8, 8).realize()
+    C = Tensor.ones(1, 8, 8).pad(((1,1), None, None),).sum(0)
+    out = (C + (A @ B))
+    out.realize()
+
+class TestSoftmaxFusion(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    with Context(TRACK_MATCH_STATS=0): cls.test = Tensor.rand(32, 10).contiguous().realize()
+
+  def setUp(self):
+    GlobalCounters.reset()
+
+  def test_norm(self):
+    print("*** norm ***")
+    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
+      # NOTE: there's an implied expand on the mean here
+      sout = self.test / self.test.mean(-1, keepdim=True)
+      sout.realize()
+
+    print("*** single kernel norm ***")
+    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
+      inp = self.test.reshape(32, 10, 1)
+      div = self.test.reshape(32, 1, 10).expand(32, 10, 10).mean(axis=-1, keepdim=True)
+      out = (inp / div).reshape(32, 10)
+      out.realize()
+
+    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
+
+  def test_softmax(self):
+    # this is the softmax from scaled_dot_product_attention
+    # it becomes 3 kernels
+    print("*** softmax ***")
+    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
+      sout = self.test.softmax(-1)
+      sout.realize()
+
+    print("*** single kernel softmax ***")
+    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
+      out = single_kernel_softmax(self.test)
+      out.realize()
+
+    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
+
+  @unittest.skip("needs RANGEIFY>1")
+  def test_auto_softmax(self):
+    print("*** softmax ***")
+    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
+      sout = self.test.softmax(-1)
+      sout.realize()
+
+    print("*** auto single kernel softmax ***")
+    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
+      out = self.test.contiguous().softmax(-1)
+      run_one_schedule_item(out)
+
+    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
+
+  def test_softmax_bw(self):
+    print("*** softmax bw ***")
+    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
+      self.test.softmax(-1).sum().backward()
+      sg = self.test.grad.realize()
+
+    self.test.grad = None
+
+    print("*** single kernel softmax bw ***")
+    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
+      single_kernel_softmax(self.test).sum().backward()
+      g = self.test.grad.realize()
+
+    np.testing.assert_allclose(sg.numpy(), g.numpy(), atol=1e-7)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_stunning.py

@@ -0,0 +1,59 @@
+import unittest
+from tinygrad import nn, Tensor, Variable, Context, Device
+from tinygrad.helpers import trange
+
+class Model:
+  def __init__(self): self.layer = nn.Linear(28*28, 10)
+  def __call__(self, x:Tensor) -> Tensor: return self.layer(x.flatten(1))
+
+class TestStunning(unittest.TestCase):
+  def test_indexing_variable(self):
+    a = Tensor.arange(100*10).reshape(100, 10).contiguous()
+
+    # index without variable
+    nv = a[12].tolist()
+
+    # index with variable
+    vi = Variable('i', 0, a.shape[0]-1)
+    wv = a[vi.bind(12)].tolist()
+
+    self.assertListEqual(nv, wv)
+
+  def test_indexing_two_bind(self):
+    a = Tensor.arange(100*10).reshape(100, 10).contiguous()
+
+    nv = a[12].cat(a[76]).tolist()
+
+    vi = Variable('i', 0, a.shape[0]-1)
+    with self.assertRaisesRegex(RuntimeError, "bind mismatch on"):
+      wv = a[vi.bind(12)].cat(a[vi.bind(76)]).tolist()
+      self.assertListEqual(nv, wv)
+
+  @unittest.skipIf(Device.DEFAULT in {"WEBGPU", "NV", "CUDA"}, "Too many buffers / too slow")
+  @unittest.skip("This is binding a Variable to two different values")
+  def test_simple_train(self, steps=6, bs=4, adam=True):
+    X_train, Y_train, _, _ = nn.datasets.mnist()
+    model = Model()
+    if adam: opt = nn.optim.Adam(nn.state.get_parameters(model))
+    else: opt = nn.optim.SGD(nn.state.get_parameters(model), momentum=0.1)
+    samples = Tensor.randint(steps, bs, high=X_train.shape[0])
+    Y_train = Y_train.one_hot(10)
+    X_samp, Y_samp = X_train[samples], Y_train[samples]
+    vi = Variable('i', 0, samples.shape[0]-1)
+    with Context(SPLIT_REDUCEOP=0):
+      with Tensor.train():
+        losses = []
+        for i in range(samples.shape[0]):
+          vib = vi.bind(i)
+          opt.zero_grad()
+          pred = model(X_samp[vib].realize())
+          loss = (pred - Y_samp[vib]).square().mean()
+          losses.append(loss.backward())
+          opt.schedule_step()
+        #losses = Tensor.stack(*losses)
+
+    # run
+    for i in (t:=trange(len(losses))): t.set_description(f"loss: {losses[i].item():6.2f}")
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_subbuffer.py

@@ -0,0 +1,184 @@
+import unittest
+from tinygrad import Device, dtypes, Tensor
+from tinygrad.device import Buffer
+from tinygrad.helpers import Context, DEV
+from test.helpers import needs_second_gpu
+
+@unittest.skipUnless(hasattr(Device[Device.DEFAULT].allocator, "_offset"), "subbuffer not supported")
+class TestSubBuffer(unittest.TestCase):
+  def setUp(self):
+    self.buf = Buffer(Device.DEFAULT, 10, dtypes.uint8).ensure_allocated()
+    self.buf.copyin(memoryview(bytearray(range(10))))
+    self.buf_unalloc = Buffer(Device.DEFAULT, 10, dtypes.uint8)
+
+  def test_subbuffer(self):
+    vbuf = self.buf.view(2, dtypes.uint8, offset=3).ensure_allocated()
+    tst = vbuf.as_memoryview().tolist()
+    assert tst == [3, 4]
+
+  def test_subbuffer_cast(self):
+    # NOTE: bitcast depends on endianness
+    vbuf = self.buf.view(2, dtypes.uint16, offset=3).ensure_allocated()
+    tst = vbuf.as_memoryview().cast("H").tolist()
+    assert tst == [3|(4<<8), 5|(6<<8)]
+
+  def test_subbuffer_double(self):
+    vbuf = self.buf.view(4, dtypes.uint8, offset=3).ensure_allocated()
+    vvbuf = vbuf.view(2, dtypes.uint8, offset=1).ensure_allocated()
+    tst = vvbuf.as_memoryview().tolist()
+    assert tst == [4, 5]
+
+  def test_subbuffer_len(self):
+    vbuf = self.buf.view(5, dtypes.uint8, 2).ensure_allocated()
+    mv = vbuf.as_memoryview()
+    assert len(mv) == 5
+    mv = vbuf.as_memoryview(allow_zero_copy=True)
+    assert len(mv) == 5
+
+  def test_subbuffer_used(self):
+    t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
+    vt = t[2:4].realize()
+    out = (vt + 100).tolist()
+    assert out == [102, 103]
+
+  @needs_second_gpu
+  @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or DEV.interface.startswith("MOCK"), "only NV, AMD, CUDA")
+  def test_subbuffer_transfer(self):
+    t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
+    vt = t[2:5].contiguous().realize()
+    out = vt.to(f"{Device.DEFAULT}:1").realize().tolist()
+    assert out == [2, 3, 4]
+
+  def test_subbuffer_deallocate(self):
+    with Context(LRU=0):
+      vbuf = self.buf.view(2, dtypes.uint8, offset=3).ensure_allocated()
+      self.buf.deallocate()
+      vbuf.deallocate()
+
+      # Allocate a fake one on the same place
+      _ = Buffer(Device.DEFAULT, 10, dtypes.uint8).ensure_allocated()
+
+      self.buf.ensure_allocated()
+      self.buf.copyin(memoryview(bytearray(range(10, 20))))
+
+      vbuf.ensure_allocated()
+
+      tst = vbuf.as_memoryview().tolist()
+      assert tst == [13, 14]
+
+  def test_subbuffer_is_allocated(self):
+    buf = self.buf_unalloc
+    sub_buf = buf.view(3, dtypes.uint8, offset=4)
+    self.assertFalse(buf.is_allocated())
+    self.assertFalse(buf.is_initialized())
+    self.assertFalse(sub_buf.is_allocated())
+    self.assertFalse(sub_buf.is_initialized())
+
+    # base buffer alloc
+    buf.allocate()
+    self.assertTrue(buf.is_allocated())
+    self.assertTrue(buf.is_initialized())
+    self.assertTrue(sub_buf.is_allocated())
+    self.assertFalse(sub_buf.is_initialized())
+
+    # sub buffer alloc
+    sub_buf.allocate()
+    self.assertTrue(sub_buf.is_initialized())
+
+    # sub buffer dealloc
+    sub_buf.deallocate()
+    self.assertTrue(buf.is_allocated())
+    self.assertTrue(buf.is_initialized())
+    self.assertTrue(sub_buf.is_allocated())
+    self.assertFalse(sub_buf.is_initialized())
+
+    # base buffer dealloc
+    buf.deallocate()
+    self.assertFalse(buf.is_allocated())
+    self.assertFalse(buf.is_initialized())
+    self.assertFalse(sub_buf.is_allocated())
+    self.assertFalse(sub_buf.is_initialized())
+
+    # sub buffer alloc
+    sub_buf.ensure_allocated()
+    self.assertTrue(buf.is_allocated())
+    self.assertTrue(buf.is_initialized())
+    self.assertTrue(sub_buf.is_allocated())
+    self.assertTrue(sub_buf.is_initialized())
+
+  def test_subbuffer_copy_in_out(self):
+    sub_buf = self.buf.view(3, dtypes.uint8, offset=3).ensure_allocated() # [3:6]
+    data_out_sub = bytearray([0]*3)
+    sub_buf.copyout(memoryview(data_out_sub))
+    assert data_out_sub == bytearray(range(3, 6))
+    sub_buf.copyin(memoryview(bytearray(range(3))))
+    assert sub_buf.as_memoryview().tolist() == list(range(3))
+    assert self.buf.as_memoryview().tolist()[3:6] == list(range(3))
+    sub_buf.copyout(memoryview(data_out_sub))
+    assert data_out_sub == bytearray(range(3))
+    data_out_base = bytearray([0]*10)
+    self.buf.copyout(memoryview(data_out_base))
+    assert data_out_base[0:3] == bytearray(range(0, 3))
+    assert data_out_base[3:6] == data_out_sub
+    assert data_out_base[6:10] == bytearray(range(6, 10))
+
+  def test_subbuffer_copy_in_out_view_of_view(self):
+    view1 = self.buf.view(7, dtypes.uint8, offset=2).ensure_allocated() # [2:9]
+    view2 = view1.view(3, dtypes.uint8, offset=2).ensure_allocated()   # [4:7]
+    self.assertTrue(view1.is_allocated())
+    self.assertTrue(view2.is_allocated())
+
+    data_in = bytearray([7, 8, 9])
+    view2.copyin(memoryview(data_in))
+    data_out_v2 = bytearray([0]*3)
+    view2.copyout(memoryview(data_out_v2))
+    assert data_in == data_out_v2
+
+    expected_base_data = memoryview(bytearray(range(10)))
+    expected_base_data[4:7] = data_in
+
+    data_out_base = bytearray([0]*10)
+    self.buf.copyout(memoryview(data_out_base))
+    assert expected_base_data == data_out_base
+
+  def test_subbuffer_alloc(self):
+    sub_buf = self.buf.view(4, dtypes.int8, offset=3)
+    sub_buf.allocate()
+    sub_buf.copyin(memoryview(bytearray(range(10, 14))))
+    assert self.buf.as_memoryview().tolist()[3:7] == sub_buf.as_memoryview().tolist()
+
+    sub_buf = self.buf_unalloc.view(4, dtypes.int8, offset=3)
+    sub_buf.allocate()
+    sub_buf.copyin(memoryview(bytearray(range(10, 14))))
+    assert self.buf_unalloc.as_memoryview().tolist()[3:7] == sub_buf.as_memoryview().tolist()
+
+  def test_subbuffer_dealloc(self):
+    sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
+    sub_buf.deallocate()
+    assert self.buf.as_memoryview().tolist() == list(range(10))
+
+  def test_subbuffer_double_dealloc(self):
+    sub_buf = self.buf.view(3, dtypes.uint8, offset=4).ensure_allocated()
+    self.buf.deallocate()
+    with self.assertRaises(AssertionError):
+      self.buf.deallocate()
+    sub_buf.deallocate()
+    with self.assertRaises(AssertionError):
+      sub_buf.deallocate()
+
+  def test_subbuffer_uaf(self):
+    sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
+    assert self.buf.as_memoryview().tolist(), list(range(10))
+    sub_buf.deallocate()
+    with self.assertRaises(AssertionError):
+      sub_buf.as_memoryview().tolist()
+    assert self.buf.as_memoryview().tolist(), list(range(10))
+
+    sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
+    assert sub_buf.as_memoryview().tolist(), list(range(3, 7))
+    self.buf.deallocate()
+    with self.assertRaises(AssertionError):
+      sub_buf.as_memoryview().tolist()
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_symbolic_jit.py

@@ -0,0 +1,332 @@
+import unittest
+
+from test.helpers import assert_jit_cache_len
+from tinygrad import Variable, Tensor, TinyJit
+from tinygrad.engine.jit import JitError
+import numpy as np
+
+class TestSymbolicJit(unittest.TestCase):
+  def test_plus1(self):
+    def f(a): return (a+1).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = jf(a[:, :vi])[:3, :i].numpy()
+      expected = f(a[:, :i]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_plus1_pad(self):
+    # TODO: without contiguous, the pad is not captured in jit
+    def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).contiguous().realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = jf(a[:, :vi]).numpy()
+      expected = f(a[:, :i]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_add(self):
+    def f(a, b): return (a+b).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(3, 10)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = jf(a[:, :vi], b[:, :vi])
+      symbolic = symbolic[:3, :i].numpy()
+      expected = f(a[:, :i], b[:, :i]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_matmul(self):
+    def f(a, b): return (a@b).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(10, 5)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = jf(a[:, :vi], b[:vi, :]).numpy()
+      expected = f(a[:, :i], b[:i, :]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_mixed_with_no_symbol_kernel(self):
+    def f(a, b):
+      s = (a@b).realize()
+      s = (s+s).realize() # this one does not have symbols in input
+      return s
+    jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(10, 5)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = jf(a[:, :vi], b[:vi, :]).numpy()
+      expected = f(a[:, :i], b[:i, :]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 2)
+
+  def test_attention(self):
+    def f(q, k, v): return Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).realize()
+    jf = TinyJit(f)
+    q = Tensor.rand(2, 1, 4, 8)
+    k = Tensor.rand(2, 10, 4, 8)
+    v = Tensor.rand(2, 10, 4, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = jf(q, k[:, :vi], v[:, :vi])[:2, :4, :1, :8].numpy()
+      expected = f(q, k[:, :i], v[:, :i]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 5)
+
+  def test_cat_dim0(self):
+    def f(a, b): return a.cat(b, dim=0).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(2, 3)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = jf(a[:vi], b)[:i+2, :3].numpy()
+      expected = f(a[:i], b).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_cat_dim1(self):
+    def f(a, b): return a.cat(b, dim=1).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(3, 2)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = jf(a[:, :vi], b)[:3, :i+2].numpy()
+      expected = f(a[:, :i], b).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_cat_dim0_two_vars(self):
+    def f(a, b): return a.cat(b, dim=0).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(10, 3)
+    for i in range(2, 5):
+      for j in range(2, 5):
+        vi = Variable("i", 1, 10).bind(i)
+        vj = Variable("j", 1, 10).bind(j)
+        symbolic = jf(a[:vi], b[:vj])[:i+j, :3].numpy()
+        expected = f(a[:i], b[:j]).numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_cat_dim1_two_vars(self):
+    def f(a, b): return a.cat(b, dim=1).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(3, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
+        vi = Variable("i", 1, 10).bind(i)
+        vj = Variable("j", 1, 10).bind(j)
+        symbolic = jf(a[:, :vi], b[:, :vj])[:3, :i+j].numpy()
+        expected = f(a[:, :i], b[:, :j]).numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_two_vars_plus1_ij(self):
+    def f(a, b): return (a@b+1).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(3, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
+        vi = Variable("i", 1, 10).bind(i)
+        vj = Variable("j", 1, 10).bind(j)
+        symbolic = jf(a[:vi, :], b[:, :vj])[:i, :j].numpy()
+        expected = f(a[:i, :], b[:, :j]).numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_two_vars_plus1_ji(self):
+    def f(a, b): return (a@b+1).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(3, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
+        vi = Variable("i", 1, 10).bind(i)
+        vj = Variable("j", 1, 10).bind(j)
+        symbolic = jf(a[:vj, :], b[:, :vi])[:j, :i].numpy()
+        expected = f(a[:j, :], b[:, :i]).numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_jit_symbolic_shape_mismatch(self):
+    @TinyJit
+    def add(a, b): return (a+b).realize()
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(3, 10)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      add(a[:, :vi], b[:, :vi])
+    vi2 = Variable("i", 1, 10).bind(7)
+    a = Tensor.rand(3, 7)[:, :vi2]
+    bad = Tensor.rand(4, 7)[:, :vi2]
+    with self.assertRaises(JitError):
+      add(a, bad)
+
+  def test_shrink(self):
+    # shrink is a movement, so we pair it with a simple function to test the JIT interaction
+    def f(a): return (a+1).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(7, 11)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = a.shrink(((3,5),(vi,vi+2)))
+      symbolic = jf(symbolic).numpy()
+      expected = f(a.shrink(((3,5),(i,i+2)))).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_slice(self):
+    # slice is a movement, so we pair it with a simple function to test the JIT interaction
+    def f(a): return (a+1).realize()
+    jf = TinyJit(f)
+    a = Tensor.rand(7, 11)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = a[3:5, vi:vi+2]
+      symbolic = jf(symbolic).numpy()
+      expected = f(a[3:5, i:i+2]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_slice_var_shape(self):
+    def f(a): return (a+1).realize()
+    jf = TinyJit(f)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      a = Tensor.ones(vi, 11).contiguous()
+      symbolic = a[:, 1:2]
+      symbolic = jf(symbolic)[:i, :1].numpy()
+      expected = f(a[:i, :][:, 1:2]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_ones_sum(self):
+    def f(a): return a.sum().realize()
+    jf = TinyJit(f)
+    t = Tensor.ones(10).contiguous()
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = jf(t[:vi]).item()
+      expected = f(t[:i]).item()
+      np.testing.assert_equal(symbolic, expected)
+
+  def test_mean(self):
+    def f(a): return a.mean().realize()
+    def f0(a): return a.mean(0).realize()
+    def f1(a): return a.mean(1).realize()
+    jf = TinyJit(f)
+    jf0 = TinyJit(f0)
+    jf1 = TinyJit(f1)
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(10, 3)
+    c = Tensor.rand(10, 3)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      # axis = None
+      symbolic = jf(a[:vi]).numpy()
+      expected = a[:i].mean().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+      # axis = 0
+      symbolic = jf0(b[:vi]).numpy()
+      expected = b[:i].mean(0).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+      # axis = 1
+      symbolic = jf1(c[:vi])[:i].numpy()
+      expected = c[:i].mean(1).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_mean_2d(self):
+    def f(a): return a.mean().realize()
+    def f0(a): return a.mean(0).realize()
+    def f1(a): return a.mean(1).realize()
+    jf = TinyJit(f)
+    jf0 = TinyJit(f0)
+    jf1 = TinyJit(f1)
+    a = Tensor.rand(10, 10)
+    b = Tensor.rand(10, 10)
+    c = Tensor.rand(10, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
+        vi = Variable("i", 1, 10).bind(i)
+        vj = Variable("j", 1, 10).bind(j)
+        # axis = None
+        symbolic = jf(a[:vi, :vj]).numpy()
+        expected = a[:i, :j].mean().numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+        # axis = 0
+        symbolic = jf0(b[:vi, :vj])[:j].numpy()
+        expected = b[:i, :j].mean(0).numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+        # axis = 1
+        symbolic = jf1(c[:vi, :vj])[:i].numpy()
+        expected = c[:i, :j].mean(1).numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_var(self):
+    def f(a): return a.var().realize()
+    def f0(a): return a.var(0).realize()
+    def f1(a): return a.var(1).realize()
+    jf = TinyJit(f)
+    jf0 = TinyJit(f0)
+    jf1 = TinyJit(f1)
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(10, 3)
+    c = Tensor.rand(10, 3)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      # axis = None
+      symbolic = jf(a[:vi]).numpy()
+      expected = a[:i].var().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+      # axis = 0
+      symbolic = jf0(b[:vi]).numpy()
+      expected = b[:i].var(0).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+      # axis = 1
+      symbolic = jf1(c[:vi])[:i].numpy()
+      expected = c[:i].var(1).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_var_2d(self):
+    def f(a): return a.var().realize()
+    def f0(a): return a.var(0).realize()
+    def f1(a): return a.var(1).realize()
+    jf = TinyJit(f)
+    jf0 = TinyJit(f0)
+    jf1 = TinyJit(f1)
+    a = Tensor.rand(10, 10)
+    b = Tensor.rand(10, 10)
+    c = Tensor.rand(10, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
+        vi = Variable("i", 1, 10).bind(i)
+        vj = Variable("j", 1, 10).bind(j)
+        # axis = None
+        symbolic = jf(a[:vi, :vj]).numpy()
+        expected = a[:i, :j].var().numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+        # axis = 0
+        symbolic = jf0(b[:vi, :vj])[:j].numpy()
+        expected = b[:i, :j].var(0).numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+        # axis = 1
+        symbolic = jf1(c[:vi, :vj])[:i].numpy()
+        expected = c[:i, :j].var(1).numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_symbolic_ops.py

@@ -0,0 +1,358 @@
+import unittest
+from tinygrad import Tensor, Variable, GlobalCounters
+from tinygrad.uop.ops import sym_infer
+from tinygrad.dtype import dtypes
+from examples.gpt2 import Attention
+import numpy as np
+
+class TestSymbolicOps(unittest.TestCase):
+  def test_plus1(self):
+    def f(a): return (a+1).realize()
+    a = Tensor.rand(3, 10)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = f(a[:, :vi])[:3, :i].numpy()
+      expected = f(a[:, :i]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_plus1_pad(self):
+    def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
+    a = Tensor.rand(3, 10)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = f(a[:, :vi]).numpy()
+      expected = f(a[:, :i]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_add(self):
+    def f(a, b): return (a+b).realize()
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(3, 10)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = f(a[:, :vi], b[:, :vi])[:, :i].numpy()
+      expected = f(a[:, :i], b[:, :i]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_matmul(self):
+    def f(a, b): return (a@b).realize()
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(10, 5)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = f(a[:, :vi], b[:vi, :]).numpy()
+      expected = f(a[:, :i], b[:i, :]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_attention(self, dropout_p=0.0, imin=1, imax=5, use_symbolic=True):
+    def f(q, k, v): return Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=dropout_p).realize()
+    q = Tensor.rand(2, 1, 4, 8)
+    k = Tensor.rand(2, 10, 4, 8)
+    v = Tensor.rand(2, 10, 4, 8)
+    for i in range(imin, imax):
+      vi = Variable("i", 1, 10).bind(i) if use_symbolic else i
+      Tensor.realize(q, k, v)
+      GlobalCounters.reset()
+      symbolic = f(q, k[:, :vi, :, :], v[:, :vi, :, :])[:2, :4, :1, :8].numpy()
+      expected = f(q, k[:, :i, :, :], v[:, :i, :, :]).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_attention_cmp_symbolic(self):
+    # symbolic isn't seeing if i == i, so it's not putting them on the same axis
+    self.test_attention(imin=4, imax=5, use_symbolic=False)
+    self.test_attention(imin=4, imax=5, use_symbolic=True)
+
+  def test_attention_training(self):
+    with Tensor.train():
+      self.test_attention(dropout_p=0.0)
+      with self.assertRaises(ValueError):
+        # symbolic shape dropout is not supported
+        self.test_attention(dropout_p=0.5)
+
+  def test_sdpa_symbolic_seq_len(self):
+    # symbolic seq_len on all of q/k/v (dim -2 after transpose)
+    q = Tensor.rand(2, 10, 4, 8)
+    k = Tensor.rand(2, 10, 4, 8)
+    v = Tensor.rand(2, 10, 4, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      Tensor.realize(q, k, v)
+      symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention(
+        k[:, :vi].transpose(1, 2), v[:, :vi].transpose(1, 2)).realize()[:2, :4, :i, :8].numpy()
+      expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention(
+        k[:, :i].transpose(1, 2), v[:, :i].transpose(1, 2)).realize().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_sdpa_symbolic_seq_len_query_only(self):
+    # symbolic seq_len on query only (dim -2 after transpose)
+    q = Tensor.rand(2, 10, 4, 8)
+    k = Tensor.rand(2, 5, 4, 8)
+    v = Tensor.rand(2, 5, 4, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      Tensor.realize(q, k, v)
+      symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention(
+        k.transpose(1, 2), v.transpose(1, 2)).realize()[:2, :4, :i, :8].numpy()
+      expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention(
+        k.transpose(1, 2), v.transpose(1, 2)).realize().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_sdpa_symbolic_batch(self):
+    # symbolic batch dim (dim 0)
+    q = Tensor.rand(10, 4, 3, 8)
+    k = Tensor.rand(10, 4, 3, 8)
+    v = Tensor.rand(10, 4, 3, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      Tensor.realize(q, k, v)
+      symbolic = q[:vi].scaled_dot_product_attention(k[:vi], v[:vi]).realize()[:i, :4, :3, :8].numpy()
+      expected = q[:i].scaled_dot_product_attention(k[:i], v[:i]).realize().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_sdpa_symbolic_heads(self):
+    # symbolic heads dim (dim -3)
+    q = Tensor.rand(2, 10, 3, 8)
+    k = Tensor.rand(2, 10, 3, 8)
+    v = Tensor.rand(2, 10, 3, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      Tensor.realize(q, k, v)
+      symbolic = q[:, :vi].scaled_dot_product_attention(k[:, :vi], v[:, :vi]).realize()[:2, :i, :3, :8].numpy()
+      expected = q[:, :i].scaled_dot_product_attention(k[:, :i], v[:, :i]).realize().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_attention_pos_0_sz_0(self):
+    Attention(128, 8)(Tensor.ones(1, 0, 128), Variable("start_pos", 0, 128).bind(0), None)
+
+  def test_attention_pos_0_sz_1(self):
+    Attention(128, 8)(Tensor.ones(1, 1, 128), Variable("start_pos", 0, 128).bind(0), None)
+
+  def test_attention_pos_0_sz_2(self):
+    Attention(128, 8)(Tensor.ones(1, 2, 128), Variable("start_pos", 0, 128).bind(0), None)
+
+  def test_cat_dim0(self):
+    def f(a, b): return a.cat(b, dim=0).realize()
+    a = Tensor.rand(10, 3)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      b = Tensor.rand(2, 3)
+      symbolic = f(a[:vi, :], b)[:i+2, :3].numpy()
+      expected = f(a[:i, :], b).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_cat_dim1(self):
+    def f(a, b): return a.cat(b, dim=1).realize()
+    a = Tensor.rand(3, 10)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      b = Tensor.rand(3, 2)
+      symbolic = f(a[:, :vi], b)[:3, :i+2].numpy()
+      expected = f(a[:, :i], b).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_cat_dim0_two_vars(self):
+    def f(a, b): return a.cat(b, dim=0).realize()
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(10, 3)
+    for i in range(2, 5):
+      for j in range(2, 5):
+        vi = Variable("i", 1, 10).bind(i)
+        vj = Variable("j", 1, 10).bind(j)
+        symbolic = f(a[:vi, :], b[:vj, :])[:i+j, :3].numpy()
+        expected = f(a[:i, :], b[:j, :]).numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_cat_dim1_two_vars(self):
+    def f(a, b): return a.cat(b, dim=1).realize()
+    a = Tensor.rand(3, 10)
+    b = Tensor.rand(3, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
+        vi = Variable("i", 1, 10).bind(i)
+        vj = Variable("j", 1, 10).bind(j)
+        symbolic = f(a[:, :vi], b[:, :vj])[:3, :i+j].numpy()
+        expected = f(a[:, :i], b[:, :j]).numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_two_vars_plus1_ij(self):
+    def f(a, b): return (a@b+1).realize()
+    a = Tensor.rand(10, 3).realize()
+    b = Tensor.rand(3, 10).realize()
+    for i in range(2, 5):
+      for j in range(2, 5):
+        vi = Variable("i", 1, 10).bind(i)
+        vj = Variable("j", 1, 10).bind(j)
+        symbolic = f(a[:vi, :], b[:, :vj])[:i, :j].numpy()
+        expected = f(a[:i, :], b[:, :j]).numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_two_vars_plus1_ji(self):
+    # reverse the order of variables
+    def f(a, b): return (a@b+1).realize()
+    a = Tensor.rand(10, 3).realize()
+    b = Tensor.rand(3, 10).realize()
+    for i in range(2, 5):
+      for j in range(2, 5):
+        vi = Variable("i", 1, 10).bind(i)
+        vj = Variable("j", 1, 10).bind(j)
+        symbolic = f(a[:vj, :], b[:, :vi])[:j, :i].numpy()
+        expected = f(a[:j, :], b[:, :i]).numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_invalid_symbolic_reshape(self):
+    a = Tensor.rand(30)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      # Cannot reshape into symbolic from non-symbolic
+      with self.assertRaises(ValueError): a.reshape((3, vi))
+
+  def test_shrink(self):
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      a = Tensor.rand(7, 11)
+      symbolic = a.shrink(((3,5),(vi,vi+2)))
+      symbolic = symbolic.numpy()
+      expected = a.shrink(((3,5),(i,i+2))).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_slice(self):
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      a = Tensor.rand(7, 11)
+      symbolic = a[3:5, vi:vi+2]
+      symbolic = symbolic.numpy()
+      expected = a[3:5, i:i+2].numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_slice_no_start(self):
+    a = Tensor.rand(7, 11)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = a[3:5, :vi:1][:2, :i].numpy()
+      expected = a[3:5, :i:1].numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_expand_padded(self):
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      a = Tensor(1).unsqueeze(0).pad((0, 1)).unsqueeze(0)
+      symbolic = a.expand(vi, 2)[:i, :2].numpy()
+      expected = a.expand(i, 2).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_slice_var_shape(self):
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      a = Tensor.ones(vi, 11).contiguous()
+      symbolic = a[:, 1:2][:i, :1].numpy()
+      expected = Tensor.ones(i, 11)[:, 1:2].numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_ones_sum(self):
+    t = Tensor.ones(10).contiguous()
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = t[:vi].sum().item()
+      expected = t[:i].sum().item()
+      np.testing.assert_equal(symbolic, expected)
+
+  def test_mean(self):
+    a = Tensor.rand(10, 3)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      for axis in [None, 0, 1]:
+        expected = a[:i].mean(axis).numpy()
+        symbolic = a[:vi].mean(axis)
+        if axis is None:
+          symbolic = symbolic.numpy()
+        else:
+          symbolic = symbolic[:expected.shape[0]].numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_mean_2d(self):
+    a = Tensor.rand(10, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
+        vi = Variable("i", 1, 10).bind(i)
+        vj = Variable("j", 1, 10).bind(j)
+        for axis in [None, 0, 1]:
+          expected = a[:i, :j].mean(axis).numpy()
+          symbolic = a[:vi, :vj].mean(axis)
+          if axis is None:
+            symbolic = symbolic.numpy()
+          else:
+            symbolic = symbolic[:expected.shape[0]].numpy()
+          np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_var(self):
+    a = Tensor.rand(10, 3)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      for axis in [None, 0, 1]:
+        expected = a[:i].var(axis).numpy()
+        symbolic = a[:vi].var(axis)
+        if axis is None:
+          symbolic = symbolic.numpy()
+        else:
+          symbolic = symbolic[:expected.shape[0]].numpy()
+        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_var_2d(self):
+    a = Tensor.rand(10, 10)
+    for i in range(2, 5):
+      for j in range(2, 5):
+        vi = Variable("i", 1, 10).bind(i)
+        vj = Variable("j", 1, 10).bind(j)
+        for axis in [None, 0, 1]:
+          expected = a[:i, :j].var(axis).numpy()
+          symbolic_result = a[:vi, :vj].var(axis)
+          if axis is None:
+            symbolic = symbolic_result.numpy()
+          else:
+            symbolic = symbolic_result[:expected.shape[0]].numpy()
+          np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_bitcast_down(self):
+    a = Tensor.rand(10, 3)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      expected = a[:i].bitcast(dtypes.uint8).numpy()
+      symbolic_result = a[:vi].bitcast(dtypes.uint8)
+      if len(expected.shape) == 2:
+        symbolic = symbolic_result[:expected.shape[0], :expected.shape[1]].numpy()
+      else:
+        symbolic = symbolic_result[:].numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)
+
+  def test_bitcast_up(self):
+    a = Tensor.rand(10, 4)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      expected = a[:i].bitcast(dtypes.uint64).numpy()
+      symbolic_result = a[:vi].bitcast(dtypes.uint64)
+      if len(expected.shape) == 2:
+        symbolic = symbolic_result[:expected.shape[0], :expected.shape[1]].numpy()
+      else:
+        symbolic = symbolic_result[:].numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)
+
+  def test_conv2d_ceildiv_edge_case(self):
+    # tests symbolic ceildiv in conv2d output shape calculation
+    # val=79 triggers the edge case where old ceildiv simplifies incorrectly: old gives floor=12, correct ceildiv=13
+    v = Variable('v', 11, 100)
+    val = 79
+    x_full = Tensor.randn(1, 8, 100)
+    weight = Tensor.randn(16, 8, 12)
+
+    # symbolic version
+    result = x_full[:, :, :v.bind(val)].conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))
+    var_val = {v.expr: val}
+    shape = tuple(sym_infer(s, var_val) for s in result.shape)
+    self.assertEqual(shape, (1, 16, 13))
+
+    # concrete version for comparison
+    expected = x_full[:, :, :val].conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))
+    np.testing.assert_allclose(result[:, :, :13].numpy(), expected.numpy(), atol=1e-5, rtol=1e-5)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_tensor.py

@@ -0,0 +1,776 @@
+import numpy as np
+import torch
+import unittest, copy, mmap, random, math, array
+from tinygrad import Tensor, Device, dtypes, nn
+from tinygrad.helpers import getenv, temp, mv_address
+from extra.gradcheck import numerical_jacobian, jacobian, gradcheck
+from hypothesis import given, settings, strategies as strat
+from tinygrad.dtype import DTYPES_DICT
+from tinygrad.uop.ops import UOp
+
+settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+settings.load_profile("my_profile")
+
+x_init = np.random.randn(1,3).astype(np.float32)
+U_init = np.random.randn(3,3).astype(np.float32)
+V_init = np.random.randn(3,3).astype(np.float32)
+W_init = np.random.randn(3,3).astype(np.float32)
+m_init = np.random.randn(1,3).astype(np.float32)
+gradient = np.random.randn(1,3).astype(np.float32)
+
+class TestTinygrad(unittest.TestCase):
+  def test_zerodim_initialization(self):
+    self.assertEqual(Tensor(55).shape, ())
+    self.assertEqual(Tensor(3.14).shape, ())
+
+  def test_deviceless_const_construct_device_repr(self):
+    t = Tensor(UOp.const(dtypes.float, 2.0))
+    self.assertIsNone(t.uop.device)
+    self.assertIsNone(t.device)
+    self.assertIn("<UOp None", repr(t))
+
+  def test_deviceless_const_realize_noop(self):
+    t = Tensor(UOp.const(dtypes.float, 2.0))
+    uop = t.uop
+    t.realize()
+    self.assertIs(t.uop, uop)
+    self.assertIsNone(t.uop.device)
+
+  def test_plus_equals(self):
+    a = Tensor.randn(10,10)
+    b = Tensor.randn(10,10)
+    c = a + b
+    val1 = c.numpy()
+    a += b
+    val2 = a.numpy()
+    np.testing.assert_allclose(val1, val2)
+
+  def test_backward_pass(self):
+    def test_tinygrad():
+      x = Tensor(x_init)
+      W = Tensor(W_init)
+      m = Tensor(m_init)
+      out = x.dot(W).relu()
+      out = out.log_softmax()
+      out = out.mul(m).add(m).sum()
+      out.backward()
+      return out.numpy(), x.grad.numpy(), W.grad.numpy()
+
+    def test_pytorch():
+      x = torch.tensor(x_init, requires_grad=True)
+      W = torch.tensor(W_init, requires_grad=True)
+      m = torch.tensor(m_init)
+      out = x.matmul(W).relu()
+      out = torch.nn.functional.log_softmax(out, dim=1)
+      out = out.mul(m).add(m).sum()
+      out.backward()
+      return out.detach().numpy(), x.grad, W.grad
+
+    for x,y in zip(test_tinygrad(), test_pytorch()):
+      np.testing.assert_allclose(x, y, atol=1e-5)
+
+  # A simple test is to check that we can accumulate gradients (run backward twice or more times)
+  def test_accumulate_gradients(self):
+    x = Tensor(x_init)
+    W = Tensor(W_init)
+    m = Tensor(m_init)
+    out = x.dot(W).relu()
+    out = out.log_softmax()
+    out = out.mul(m).add(m).sum()
+    out.backward()
+    xgrad, wgrad = x.grad.numpy(), W.grad.numpy()
+    out.backward()
+    xgrad2, wgrad2 = x.grad.numpy(), W.grad.numpy()
+    out.backward() # no need to retain again since we will not re-run backward
+    xgrad3, wgrad3 = x.grad.numpy(), W.grad.numpy()
+    np.testing.assert_allclose(xgrad3, xgrad * 3., atol=1e-6)
+    np.testing.assert_allclose(wgrad3, wgrad * 3., atol=1e-6)
+    np.testing.assert_allclose(xgrad2, xgrad * 2., atol=1e-6)
+    np.testing.assert_allclose(wgrad2, wgrad * 2., atol=1e-6)
+
+  def test_second_order_backward_pass(self):
+    def test_pytorch():
+      x_val = torch.tensor([2.0], requires_grad=True)
+      f = x_val**3
+      first_derivative = torch.autograd.grad(outputs=f, inputs=x_val, create_graph=True)[0]
+      second_derivative = torch.autograd.grad(outputs=first_derivative, inputs=x_val)[0]
+      # d^2f/dx^2 = 6x = 6*2 = 12
+      return second_derivative.numpy()
+
+    def test_tinygrad():
+      x_val = Tensor([2.0])
+      f = x_val**3
+      first_derivative = f.sum().gradient(x_val)[0]
+      second_derivative = first_derivative.sum().gradient(x_val)[0]
+      return second_derivative.numpy()
+
+    np.testing.assert_allclose(test_tinygrad(), test_pytorch(), atol=1e-5)
+
+  # passing `gradient` to backward
+  def test_backward_pass_vjp(self):
+    def test_tinygrad():
+      x = Tensor(x_init)
+      W = Tensor(W_init)
+      m = Tensor(m_init)
+      out = x.dot(W).relu()
+      out = out.log_softmax()
+      out = out.mul(m).add(m)
+      out.backward(Tensor(gradient))
+      return out.numpy(), x.grad.numpy(), W.grad.numpy()
+
+    def test_pytorch():
+      x = torch.tensor(x_init, requires_grad=True)
+      W = torch.tensor(W_init, requires_grad=True)
+      m = torch.tensor(m_init)
+      out = x.matmul(W).relu()
+      out = torch.nn.functional.log_softmax(out, dim=1)
+      out = out.mul(m).add(m)
+      out.backward(torch.tensor(gradient))
+      return out.detach().numpy(), x.grad, W.grad
+
+    for x,y in zip(test_tinygrad(), test_pytorch()):
+      np.testing.assert_allclose(x, y, atol=1e-5)
+
+  def test_backward_pass_diamond_model(self):
+    def test_tinygrad():
+      u = Tensor(U_init)
+      v = Tensor(V_init)
+      w = Tensor(W_init)
+      x = u.mul(v).relu()
+      y = u.mul(w).relu()
+      out = x.add(y).mul(y).relu()
+      out = out.log_softmax()
+      out = out.sum()
+      out.backward()
+      return out.numpy(), u.grad.numpy(), v.grad.numpy(), w.grad.numpy()
+
+    def test_pytorch():
+      u = torch.tensor(U_init, requires_grad=True)
+      v = torch.tensor(V_init, requires_grad=True)
+      w = torch.tensor(W_init, requires_grad=True)
+      x = u.mul(v).relu()
+      y = u.mul(w).relu()
+      out = x.add(y).mul(y).relu()
+      out = torch.nn.functional.log_softmax(out, dim=1)
+      out = out.sum()
+      out.backward()
+      return out.detach().numpy(), u.grad, v.grad, w.grad
+
+    for x,y in zip(test_tinygrad(), test_pytorch()):
+      np.testing.assert_allclose(x, y, atol=1e-5, rtol=1e-6)
+
+  def test_const_backward_pass(self):
+    init = 3.5
+
+    def test_pytorch():
+      w1 = torch.tensor(init, requires_grad=True)
+      w2 = torch.tensor(init, requires_grad=True)
+      out = w1.add(w2)
+      out.backward()
+      return w1.grad, w2.grad
+
+    def test_tinygrad():
+      w1 = Tensor(init).clone()
+      w2 = Tensor(init).clone()
+      out = w1.add(w2)
+      out.backward()
+      return w1.grad.numpy(), w2.grad.numpy()
+
+    for x, y in zip(test_tinygrad(), test_pytorch()):
+      np.testing.assert_allclose(x, y, atol=1e-5)
+
+  def test_const_backward_pass_optimizer(self):
+    init = 3.5
+
+    def test_pytorch():
+      w1 = torch.tensor(init, requires_grad=True)
+      w2 = torch.tensor(init, requires_grad=True)
+      out = w1.add(w2)
+      out.backward()
+      return w1.grad.numpy(), w2.grad.numpy()
+
+    def test_tinygrad():
+      w1 = Tensor(init).clone()
+      w2 = Tensor(init).clone()
+      assert w1.is_param is True and w2.is_param is True
+      nn.optim.SGD([w1, w2], lr=0.01)
+      assert w1.is_param is True and w2.is_param is True
+      out = w1.add(w2)
+      out.backward()
+      return w1.grad.numpy(), w2.grad.numpy()
+
+    for x, y in zip(test_tinygrad(), test_pytorch()):
+      np.testing.assert_allclose(x, y, atol=1e-5)
+
+  def test_dropout(self):
+    with Tensor.train():
+      n, rate = 1_000_000, 0.1
+      w = Tensor.ones(n).dropout(rate)
+      non_zeros = np.count_nonzero(w.numpy())
+      expected = n * (1 - rate)
+      np.testing.assert_allclose(non_zeros, expected, rtol=2e-3)
+
+  def test_jacobian(self):
+    W = np.random.RandomState(42069).random((10, 5)).astype(np.float32)
+    x = np.random.RandomState(69420).random((1, 10)).astype(np.float32)
+
+    torch_x = torch.tensor(x, requires_grad=True)
+    torch_W = torch.tensor(W, requires_grad=True)
+    def torch_func(x): return torch.nn.functional.log_softmax(x.matmul(torch_W).relu(), dim=1)
+    PJ = torch.autograd.functional.jacobian(torch_func, torch_x).squeeze().numpy()
+
+    tiny_x = Tensor(x)
+    tiny_W = Tensor(W)
+    def tiny_func(x): return x.dot(tiny_W).relu().log_softmax()
+    J = jacobian(tiny_func, tiny_x)
+    NJ = numerical_jacobian(tiny_func, tiny_x)
+
+    np.testing.assert_allclose(PJ, J, atol = 1e-5)
+    np.testing.assert_allclose(PJ, NJ, atol = 1e-3)
+
+  def test_gradcheck(self):
+    W = np.random.RandomState(1337).random((10, 5)).astype(np.float32)
+    x = np.random.RandomState(7331).random((1, 10)).astype(np.float32)
+
+    tiny_x = Tensor(x)
+    tiny_W = Tensor(W)
+    def tiny_func(x): return x.dot(tiny_W).relu().log_softmax()
+
+    self.assertTrue(gradcheck(tiny_func, tiny_x, eps = 1e-3))
+
+    # coarse approx. since a "big" eps and the non-linearities of the model
+    self.assertFalse(gradcheck(tiny_func, tiny_x, eps = 1e-5))
+
+  def test_random_fns_are_deterministic_with_seed(self):
+    for random_fn in [Tensor.randn, Tensor.normal, Tensor.uniform, Tensor.scaled_uniform, Tensor.glorot_uniform, Tensor.kaiming_normal]:
+      with self.subTest(msg=f"Tensor.{random_fn.__name__}"):
+        Tensor.manual_seed(1337)
+        a = random_fn(10,10).realize()
+        Tensor.manual_seed(1337)
+        b = random_fn(10,10).realize()
+        np.testing.assert_allclose(a.numpy(), b.numpy())
+
+  def test_randperm(self):
+    Tensor.manual_seed(0)
+    a = Tensor.randperm(10).realize()
+    np.testing.assert_equal(a.numpy(), [8, 9, 4, 3, 6, 1, 7, 5, 2, 0])
+    b = Tensor.randperm(1000).realize()
+    np.testing.assert_equal(set(b.numpy()), set(range(1000)))
+
+  def test_rand_rejects_unknown_kwargs(self):
+    with self.assertRaises(TypeError): Tensor.rand(5, generator="foo")
+
+  def test_randn_isnt_inf_on_zero(self):
+    # simulate failure case of rand handing a zero to randn
+    original_rand, Tensor.rand = Tensor.rand, Tensor.zeros
+    try: self.assertNotIn(np.inf, Tensor.randn(16).numpy())
+    except: raise
+    finally: Tensor.rand = original_rand
+
+  def test_zeros_like_has_same_dtype_and_shape(self):
+    for datatype in [dtypes.float16, dtypes.float32, dtypes.int8, dtypes.int32, dtypes.int64, dtypes.uint8]:
+      a = Tensor([1, 2, 3], dtype=datatype)
+      b = Tensor.zeros_like(a)
+      assert a.dtype == b.dtype, f"dtype mismatch {a.dtype=} != {b.dtype}"
+      assert a.shape == b.shape, f"shape mismatch {a.shape} != {b.shape}"
+
+    a = Tensor([1, 2, 3])
+    b = Tensor.zeros_like(a, dtype=dtypes.int8)
+    assert a.dtype == dtypes.default_int and b.dtype == dtypes.int8, "a.dtype should be int and b.dtype should be char"
+    assert a.shape == b.shape, f"shape mismatch {a.shape} != {b.shape}"
+
+  def test_ones_like_has_same_dtype_and_shape(self):
+    for datatype in [dtypes.float16, dtypes.float32, dtypes.int8, dtypes.int32, dtypes.int64, dtypes.uint8]:
+      a = Tensor([1, 2, 3], dtype=datatype)
+      b = Tensor.ones_like(a)
+      assert a.dtype == b.dtype, f"dtype mismatch {a.dtype=} != {b.dtype}"
+      assert a.shape == b.shape, f"shape mismatch {a.shape} != {b.shape}"
+
+    a = Tensor([1, 2, 3])
+    b = Tensor.ones_like(a, dtype=dtypes.int8)
+    assert a.dtype == dtypes.default_int and b.dtype == dtypes.int8, "a.dtype should be int and b.dtype should be char"
+    assert a.shape == b.shape, f"shape mismatch {a.shape} != {b.shape}"
+
+  def test_rand_like_device(self):
+    a = Tensor.ones(3, 3, device="CPU")
+    b = Tensor.rand_like(a)
+    self.assertEqual(b.device, a.device)
+
+  def test_ndim(self):
+    assert Tensor(1).ndim == 0
+    assert Tensor.randn(1).ndim == 1
+    assert Tensor.randn(2,2,2).ndim == 3
+    assert Tensor.randn(1,1,1,1,1,1).ndim == 6
+
+  def test_argfix(self):
+    for f in [Tensor.zeros, Tensor.ones, Tensor.rand, Tensor.randn, Tensor.empty]:
+      self.assertEqual(f().shape, ())
+      self.assertEqual(f(1).shape, (1,))
+      self.assertEqual(f(10,20,40).shape, (10,20,40))
+      self.assertEqual(f([]).shape, ())
+      self.assertEqual(f([1]).shape, (1,))
+      self.assertEqual(f([10,20,40]).shape, (10,20,40))
+      self.assertEqual(f(()).shape, ())
+      self.assertEqual(f((1,)).shape, (1,))
+      self.assertEqual(f((10,20,40)).shape, (10,20,40))
+
+      with self.assertRaises(ValueError): f((2, 2), 2, 2)
+      with self.assertRaises(ValueError): f((2, 2), (2, 2))
+      with self.assertRaises(ValueError): f((128, 128), 0.0, 0.01)
+
+  def test_numel(self):
+    assert Tensor.randn(10, 10).numel() == 100
+    assert Tensor.randn(1,2,5).numel() == 10
+    assert Tensor.randn(1,1,1,1,1,1).numel() == 1
+    assert Tensor([]).numel() == 0
+    assert Tensor.randn(1,0,2,5).numel() == 0
+    assert Tensor(3).numel() == 1
+
+  def test_len(self):
+    assert len(torch.zeros(7)) == len(Tensor.zeros(7))
+    assert len(torch.zeros(10,20)) == len(Tensor.zeros(10,20))
+    assert len(torch.zeros(10,20)) == len(Tensor.zeros(10,20,30))
+    assert len(torch.zeros(1).flatten()) == len(Tensor.zeros(1).flatten())
+    with self.assertRaises(TypeError): len(Tensor(3))
+
+  def test_size(self):
+    t1, t2 = torch.zeros(10,20), Tensor.zeros(10,20)
+    assert t1.size() == t2.size()
+    assert t1.size(0) == t2.size(0)
+    assert t1.size(1) == t2.size(1)
+    assert t1.size(-1) == t2.size(-1)
+    assert t1.size(-2) == t2.size(-2)
+    with self.assertRaises(IndexError): t2.size(2)
+
+  def test_tolist(self):
+    # NOTE: float16 Tensor.tolist() requires python 3.12
+    for arr in [[1,2,3], [1.5,2,3], [[1,2,3], [4,5,6]], 3]:
+      assert Tensor(arr).tolist() == torch.tensor(arr).tolist() == arr
+
+  def test_element_size(self):
+    for _, dtype in DTYPES_DICT.items():
+      assert dtype.itemsize == Tensor.randn(3, dtype=dtype).element_size(), f"Tensor.element_size() not matching Tensor.dtype.itemsize for {dtype}"
+
+  def test_deepwalk_ctx_check(self):
+    layer = Tensor.uniform(1, 1)
+    x = Tensor.randn(1, 1, 1)
+    x.dot(layer).mean().backward()
+    x = Tensor.randn(1, 1, 1)
+    x.dot(layer).mean().backward()
+
+  def test_zerosized_tensors(self):
+    np.testing.assert_equal(Tensor([]).numpy(), np.array([]))
+    np.testing.assert_equal(Tensor(None).numpy(), np.array([]))
+
+  def test_tensor_ndarray_dtype(self):
+    arr = np.array([1]) # where dtype is implicitly int64
+    assert Tensor(arr).dtype == dtypes.int64
+    assert Tensor(arr, dtype=dtypes.float32).dtype == dtypes.float32 # check if ndarray correctly casts to Tensor dtype
+    assert Tensor(arr, dtype=dtypes.float64).dtype == dtypes.float64 # check that it works for something else
+
+  def test_tensor_from_blob(self):
+    x = memoryview(bytearray(16)).cast('I')
+
+    t = Tensor.from_blob(mv_address(x), (4,), dtype=dtypes.int, device="CPU")
+    z = (t+1)
+    np.testing.assert_equal(z.numpy(), [1, 1, 1, 1])
+
+    x[:] = array.array('I', [0, 1, 2, 3])
+    z = (t+1)
+    np.testing.assert_equal(z.numpy(), [1, 2, 3, 4])
+
+  def test_tensor_list_dtype(self):
+    for arr in ([1], [[[1]]], [[1,1],[1,1]], [[[1,1],[1,1]],[[1,1],[1,1]]]):
+      assert Tensor(arr).dtype == dtypes.default_int
+      assert Tensor(arr, dtype=dtypes.float32).dtype == dtypes.float32
+      assert Tensor(arr, dtype=dtypes.float64).dtype == dtypes.float64
+
+    for arr in ([True], [[[False]]], [[True,False],[True,False]], [[[False,True],[False,False]],[[True,True],[False,True]]]):
+      assert Tensor(arr).dtype == dtypes.bool
+      assert Tensor(arr, dtype=dtypes.float32).dtype == dtypes.float32
+      assert Tensor(arr, dtype=dtypes.float64).dtype == dtypes.float64
+
+    # empty tensor defaults
+    for arr in ([], [[[]]], [[],[]]):
+      t = Tensor(arr)
+      assert t.dtype == dtypes.default_float
+      np.testing.assert_allclose(t.numpy(), np.array(arr))
+
+    # mixture of bool and int
+    for arr in ([True, 3], [[True],[3]], [[[True]], [[3]]], [[True, 3], [3, True]]):
+      t = Tensor(arr)
+      assert t.dtype == dtypes.default_int
+      np.testing.assert_allclose(t.numpy(), np.array(arr))
+
+    # mixture of bool, int and float
+    for arr in ([[True,True],[3.,True]], [[0,1],[3.,4]], [[[0],[1]],[[3.],[4]]], [[[True],[1]],[[3.],[4]]]):
+      t = Tensor(arr)
+      assert t.dtype == dtypes.default_float
+      np.testing.assert_allclose(t.numpy(), np.array(arr))
+
+  def test_tensor_list_shapes(self):
+    self.assertEqual(Tensor([[[]]]).shape, (1,1,0))
+    self.assertEqual(Tensor([[],[]]).shape, (2,0))
+    self.assertEqual(Tensor([[[[]],[[]]], [[[]],[[]]], [[[]],[[]]]]).shape, (3,2,1,0))
+
+  def test_tensor_list_errors(self):
+    # inhomogeneous shape
+    with self.assertRaises(ValueError): Tensor([[],[[]]])
+    with self.assertRaises(ValueError): Tensor([[1],[]])
+    with self.assertRaises(ValueError): Tensor([[1],[1],1])
+    with self.assertRaises(ValueError): Tensor([[[1,1,1],[1,1]]])
+    with self.assertRaises(ValueError): Tensor([[1,1,1],[[1,1,1]]])
+
+  def test_tensor_mixed_list_tuple(self):
+    def _list_or_tuple(): return list if random.random() < 0.5 else tuple
+    def _generate_data(depth):
+      if depth == 0: return _list_or_tuple()()
+      if depth == 1: return _list_or_tuple()([random.random(), random.random()])
+      return _list_or_tuple()([_generate_data(depth-1), _generate_data(depth-1)])
+
+    for depth in range(7):
+      for _ in range(20):
+        data = _generate_data(depth)
+        np.testing.assert_allclose(Tensor(data).numpy(), np.array(data))
+
+  def test_tensor_list_implicit_cast(self):
+    data = [True, False]
+    np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
+    np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
+    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
+    data = [-1, 0, 1, 2, 3]
+    np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
+    np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
+    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
+    data = [-3.5, -2.5, -1.5, 0, 1.5, 2.5, 3.5]
+    np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
+    # NOTE: torch and jax raise OverflowError: Python integer -3 out of bounds for uint8
+    # np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
+    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
+
+  def test_tensor_list_special_values(self):
+    if dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes():
+      data = [math.nan, -math.inf, 65504, 65519, 65519.999, 65520, 65520.1]
+      data = data + [-x for x in data]
+      with np.errstate(over='ignore'): np.testing.assert_allclose(Tensor(data, dtype=dtypes.float16).numpy(), np.array(data).astype(np.float16))
+
+    # uint32
+    data = [1 << 33, 1 << 32, 1 << 32 - 1, 1]
+    data = data + [-x for x in data]
+    np.testing.assert_allclose(Tensor(data, dtype=dtypes.uint32).numpy(), np.array(data).astype(np.uint32))
+
+    # int32
+    data = [1 << 33, 1 << 32, 1 << 32 - 1, 1]
+    data = data + [-x for x in data]
+    np.testing.assert_allclose(Tensor(data, dtype=dtypes.int32).numpy(), np.array(data).astype(np.int32))
+
+  def test_tensor_list_ndarray(self):
+    data = [np.array([1, 2, 3]), np.array([1, 2, 3]), np.array([1, 2, 3])]
+    np.testing.assert_equal(Tensor(data).numpy(), np.array(data))
+    data = [np.array([1.0, 2.0, 3.0]), np.array([1, 2, 3]), np.array([1, 2, 3])]
+    np.testing.assert_equal(Tensor(data).numpy(), np.array(data))
+    data = [np.array(1.0), np.array(2.0), np.array(3.0)]
+    np.testing.assert_equal(Tensor(data).numpy(), np.array(data))
+
+  def test_tensor_dtype_errors(self):
+    with self.assertRaises(AttributeError): Tensor([3], dtype="typo")
+    with self.assertRaises(AttributeError): Tensor([3], dtype=(dtypes.int,))
+
+  def test_tensor_bytes(self):
+    data = b"abc123"
+    t = Tensor(data)
+    assert t.dtype == dtypes.uint8
+    assert t.shape == (6,)
+    np.testing.assert_equal(t.numpy(), list(data))
+
+  def test_tensor_copy(self):
+    x = copy.deepcopy(Tensor.ones((3,3,3)))
+    np.testing.assert_allclose(x.numpy(), np.ones((3,3,3)))
+
+  def test_copy_from_disk(self):
+    t = Tensor.randn(30).to(f"disk:{temp('test_copy_from_disk')}")
+    a = t[10:20]
+    dev = a.to(Device.DEFAULT)
+    np.testing.assert_allclose(a.numpy(), dev.numpy())
+
+  def test_copy_from_numpy_dtype(self):
+    data = np.array([1.0, 2, 3], dtype=np.float32)
+    t = Tensor(data, dtype=dtypes.bfloat16)
+    try:
+      # TODO: fix dtype in tinygrad space
+      assert t.dtype == dtypes.bfloat16
+    except AssertionError:
+      assert t.dtype == dtypes.float32
+    np.testing.assert_equal(t.tolist(), data)
+    np.testing.assert_equal((t+1).tolist(), data+1)
+
+  # Regression test for https://github.com/tinygrad/tinygrad/issues/1751
+  def test_copy_from_numpy_unaligned(self):
+    # 2**15 is the minimum for repro
+    arr = np.random.randn(2**15).astype(np.float32)
+    fn = temp('test_copy_from_numpy_unaligned')
+    with open(fn, 'wb') as f: f.write(b't' + arr.tobytes())
+    with open(fn, "a+b") as f: memview = memoryview(mmap.mmap(f.fileno(), arr.nbytes + 1))
+    ua_arr = np.frombuffer(memview[1:], dtype=arr.dtype, count=arr.shape[0])
+    np.testing.assert_allclose(arr, ua_arr)
+    assert not ua_arr.flags.aligned
+    # force device copy - to() is opt'd away - Tensor(dev)/1 is ignored
+    np.testing.assert_allclose(ua_arr, (Tensor(ua_arr)/Tensor(1)).numpy())
+
+  def test_item_to_tensor_to_item(self):
+    for a in [0, 1, 2, 3, -1, -100, 100, -101.1, 2.345, 100.1, True, False]:
+      item = Tensor(a).item()
+      assert type(item) is type(a), a
+      np.testing.assert_allclose(item, a), a
+      buffered_item = Tensor([a]).item()
+      assert type(buffered_item) is type(a), a
+      np.testing.assert_allclose(buffered_item, a), a
+      reshaped_item = Tensor([a]).reshape((1, 1, 1, 1, 1)).item()
+      assert type(reshaped_item) is type(a), a
+      np.testing.assert_allclose(reshaped_item, a), a
+
+  def test_no_bool(self):
+    with self.assertRaises(TypeError):
+      if Tensor(3):
+        print("hi")
+
+    with self.assertRaises(TypeError):
+      _a = Tensor([3]) in [Tensor([3]), Tensor([4]), Tensor([5])]
+
+  def test_repr_with_grad(self):
+    a = Tensor([1.0])
+    b = Tensor([1])
+    c = (a + b).sum().backward()
+    print(a)
+    print(c)
+
+  def test_no_attributeerror_after_apply_uop_exception(self):
+    try:
+      Tensor.arange(4).reshape(3,2)
+    except ValueError:
+      Tensor.zeros(2, 2).realize()
+
+  def test_shrink(self):
+    t = Tensor.arange(32).clone().realize()
+    self.assertListEqual(t[16:20].tolist(), [16,17,18,19])
+    self.assertListEqual(t.shrink_to(16).tolist(), list(range(16)))
+    t = t.reshape(4, 8).contiguous().realize()
+    self.assertListEqual(t.shrink_to(2, 2).tolist(), [[0, 1], [8, 9]])
+    self.assertListEqual(t.shrink_to(None, 2).tolist(), t.shrink_to(4, 2).tolist())
+    with self.assertRaises(ValueError): t.shrink_to(2)
+    with self.assertRaises(ValueError): t.shrink_to(2, 2, 2)
+
+@unittest.skip("this test is just flaky, sync issue")
+class TestMoveTensor(unittest.TestCase):
+  d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
+  @given(strat.sampled_from([d0, d1]), strat.sampled_from([d0, d1]),
+         strat.sampled_from([dtypes.float16, dtypes.float32]), strat.sampled_from([True, False]))
+  def test_to_preserves(self, src, dest, dtype, is_param):
+    if dtype not in Device[Device.DEFAULT].renderer.supported_dtypes():
+      return
+    s = Tensor([1, 2, 3], device=src, dtype=dtype).is_param_(is_param)
+    if is_param: s.sum().backward()
+    t = s.to(dest)
+    np.testing.assert_equal(s.numpy(), t.numpy())
+    assert s.dtype == t.dtype
+    assert s.is_param == t.is_param
+    if is_param:
+      np.testing.assert_equal(s.grad.numpy(), t.grad.numpy())
+
+  @given(strat.sampled_from([dtypes.float16, dtypes.float32]), strat.sampled_from([True, False]))
+  def test_shard_preserves(self, dtype, is_param):
+    s = Tensor([1, 2, 3], dtype=dtype).is_param_(is_param)
+    t = s.shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"))
+    np.testing.assert_equal(s.numpy(), t.numpy())
+    assert s.dtype == t.dtype
+    assert s.is_param == t.is_param
+
+  @given(strat.sampled_from([d0, d1]))
+  def test_same_dev(self, dev):
+    x = Tensor([1,2,3], device=dev)
+    y = x.to(dev)
+    assert x is y
+
+  def test_to_grad(self):
+    x = Tensor.eye(3, device=self.d0)
+    y = Tensor([[2.0,0,-2.0]], device=self.d0)
+    z = y.matmul(x).to(self.d1).sum()
+    z.backward()
+    np.testing.assert_equal(x.grad.numpy(), [[2,2,2],[0,0,0],[-2,-2,-2]])
+
+class TestZeroShapeTensor(unittest.TestCase):
+  def test_rand(self):
+    t = Tensor.rand(3, 2, 0)
+    assert t.shape == (3, 2, 0)
+    np.testing.assert_equal(t.numpy(), np.zeros((3, 2, 0)))
+    t = Tensor.rand(0)
+    assert t.shape == (0,)
+    np.testing.assert_equal(t.numpy(), np.zeros((0,)))
+    t = Tensor.rand(0, 0, 0)
+    assert t.shape == (0, 0, 0)
+    np.testing.assert_equal(t.numpy(), np.zeros((0, 0, 0)))
+
+  def test_full(self):
+    t = Tensor.zeros(3, 2, 0)
+    assert t.shape == (3, 2, 0)
+    np.testing.assert_equal(t.numpy(), np.zeros((3, 2, 0)))
+    t = Tensor.full((3, 2, 0), 12)
+    assert t.shape == (3, 2, 0)
+    np.testing.assert_equal(t.numpy(), np.full((3, 2, 0), 12))
+
+  def test_reshape(self):
+    t = Tensor.zeros(3, 2, 0)
+    a = t.reshape(7, 0)
+    assert a.shape == (7, 0)
+    np.testing.assert_equal(a.numpy(), np.zeros((7, 0)))
+    a = t.reshape(0)
+    assert a.shape == (0,)
+    np.testing.assert_equal(a.numpy(), np.zeros((0,)))
+    with self.assertRaises(ValueError):
+      # cannot reshape from size 0 to size 1
+      a = t.reshape(())
+
+  def test_expand(self):
+    t = Tensor.full((1, 2, 0), 12).expand((6, 2, 0))
+    assert t.shape == (6, 2, 0)
+    np.testing.assert_equal(t.numpy(), np.full((6, 2, 0), 12))
+
+  def test_pad(self):
+    t = Tensor.rand(3, 2, 0).pad((None, None, (1, 1)), value=1)
+    self.assertEqual(t.shape, (3, 2, 2))
+    np.testing.assert_equal(t.numpy(), np.ones((3, 2, 2)))
+
+    t = Tensor.rand(3, 2, 0).pad((None, (1, 1), None), value=1)
+    self.assertEqual(t.shape, (3, 4, 0))
+    np.testing.assert_equal(t.numpy(), np.ones((3, 4, 0)))
+
+    t = Tensor.rand(3, 2, 0).pad(((1, 1), None, None), value=1)
+    self.assertEqual(t.shape, (5, 2, 0))
+    np.testing.assert_equal(t.numpy(), np.ones((5, 2, 0)))
+
+    np.testing.assert_equal(Tensor([1, 2]).pad_to(4).numpy(), [1, 2, 0, 0])
+    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3).numpy(), [[1, 2, 0], [0, 0, 0]])
+    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(1, 3).numpy(), [[1, 2, 0]])
+    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(None, 3).numpy(), [[1, 2, 0]])
+    with self.assertRaises(ValueError): Tensor([1, 2]).pad_to(2, 3)
+    with self.assertRaises(ValueError): Tensor([[1, 2]]).pad_to(3)
+
+  def test_shrink_into_zero(self):
+    t = Tensor.rand(3, 4).realize()
+    assert t.shrink((None, (2, 2))).realize().shape == (3, 0)
+    assert t.shrink(((2, 2), None)).realize().shape == (0, 4)
+    assert t.shrink(((2, 2), (2, 2))).realize().shape == (0, 0)
+
+  def test_cat(self):
+    a = Tensor.rand(3, 2, 2)
+    b = Tensor.rand(3, 2, 0)
+
+    t = a.cat(b, dim=2)
+    assert t.shape == (3, 2, 2)
+    np.testing.assert_equal(t.numpy(), a.numpy())
+
+    t = b.cat(a, dim=2)
+    assert t.shape == (3, 2, 2)
+    np.testing.assert_equal(t.numpy(), a.numpy())
+
+    t = b.cat(b, dim=0)
+    assert t.shape == (6, 2, 0)
+    np.testing.assert_equal(t.numpy(), np.zeros((6, 2, 0)))
+    t = b.cat(b, dim=1)
+    assert t.shape == (3, 4, 0)
+    np.testing.assert_equal(t.numpy(), np.zeros((3, 4, 0)))
+    t = b.cat(b, dim=2)
+    assert t.shape == (3, 2, 0)
+    np.testing.assert_equal(t.numpy(), np.zeros((3, 2, 0)))
+
+  def test_elementwise(self):
+    a = Tensor.rand(3, 2, 0)
+    a_exp = a.exp()
+    assert a_exp.shape == (3, 2, 0)
+    np.testing.assert_equal(a_exp.numpy(), np.exp(a.numpy()))
+
+    b = Tensor.rand(3, 2, 0)
+    assert b.shape == (3, 2, 0)
+    ab = a * b
+    assert ab.shape == (3, 2, 0)
+    np.testing.assert_equal(ab.numpy(), a.numpy() * b.numpy())
+
+    mask = (Tensor.rand(3, 2, 0) > 0.5)
+    assert mask.shape == (3, 2, 0)
+    c = mask.where(a, b)
+    assert c.shape == (3, 2, 0)
+    np.testing.assert_equal(c.numpy(), np.where(mask.numpy(), a.numpy(), b.numpy()))
+
+  def test_reduce_over_non_zero(self):
+    a = Tensor.ones(3, 2, 0).sum(axis=1)
+    assert a.shape == (3, 0)
+    np.testing.assert_equal(a.numpy(), np.sum(np.zeros((3, 2, 0)), axis=1))
+
+  def test_reduce_over_zero(self):
+    a = Tensor.ones(3, 2, 0).sum(axis=2)
+    assert a.shape == (3, 2)
+    np.testing.assert_equal(a.numpy(), np.sum(np.zeros((3, 2, 0)), axis=2))
+
+    a = Tensor.ones(3, 2, 0).sum(axis=2, keepdim=True)
+    assert a.shape == (3, 2, 1)
+    np.testing.assert_equal(a.numpy(), np.sum(np.zeros((3, 2, 0)), axis=2, keepdims=True))
+
+  def test_clone(self):
+    a = Tensor.rand(16, 16).realize()
+    b = a.clone()
+    np.testing.assert_allclose(a.numpy(), b.numpy())
+    self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
+
+    a = Tensor.rand(16, 16).mul(5.0).add(5.0).realize()
+    b = a.clone()
+    np.testing.assert_allclose(a.numpy(), b.numpy())
+    self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
+
+  def test_clone_deviceless_const(self):
+    t = Tensor(UOp.const(dtypes.float, 2.0)).clone()
+    np.testing.assert_equal(t.numpy(), 2.0)
+    self.assertTrue(t.uop.has_buffer_identity())
+
+  def test_numpy_deviceless_const(self):
+    np.testing.assert_equal(Tensor(UOp.const(dtypes.float, 2.0)).numpy(), 2.0)
+
+  def test_clone_with_shrink(self):
+    a = Tensor.rand(16, 16)
+    b = a.shrink(((2, 10), None)).clone()
+    b.realize()
+    self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
+
+  def test_clone_with_shrink_realized(self):
+    a = Tensor.rand(16, 16).realize()
+    b = a.shrink(((2, 10), None)).clone()
+    b.realize()
+    self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
+
+  def test_clone_with_grad(self):
+    a = Tensor.rand(16, 16)
+    a.mul(5.0).add(5.0).mean().backward()
+    b = a.clone()
+    assert a.grad is not None
+    assert b.grad is not None
+    np.testing.assert_allclose(a.grad.numpy(), b.grad.numpy())
+
+  def test_clone_deviceless_const_to_cpu(self):
+    t = Tensor(UOp.const(dtypes.float, 2.0)).clone(device="CPU")
+    self.assertEqual(t.device, "CPU")
+    np.testing.assert_equal(t.numpy(), 2.0)
+
+  def test_reduce_default(self):
+    np.testing.assert_equal(Tensor([]).max().numpy(), -float("inf"))
+    np.testing.assert_equal(Tensor([]).min().numpy(), float("inf"))
+    np.testing.assert_equal(Tensor([]).sum().numpy(), 0)
+    np.testing.assert_equal(Tensor([]).mean().numpy(), float("nan"))
+
+class TestTensorCreationDevice(unittest.TestCase):
+  # test auxiliary tensors are created on the same device
+  def test_one_hot(self):
+    y = Tensor([1, 2, 3]).to("CPU")
+    x = y.one_hot(10)
+    x.realize()
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_tensor_variable.py

@@ -0,0 +1,165 @@
+import unittest
+import numpy as np
+from tinygrad import Tensor, Variable
+
+class TestTensorVariable(unittest.TestCase):
+  def test_add_tvar(self):
+    vv = Variable("a", 0, 10).bind(1)
+    ret = (Tensor(vv) + 3).item()
+    assert ret == 4
+
+  def test_inner_tvar_node(self):
+    vv = Variable("w", 0, 10).bind(2)
+    ret = Tensor(vv * 4).item()
+    assert ret == 8
+
+  def test_inner_tvar_mul(self):
+    vv = Variable("w", 0, 10).bind(2)
+    assert (Tensor(3) * vv).item() == 6
+
+  def test_inner_tvar_mul_node(self):
+    vv = Variable("w", 0, 10).bind(2)
+    assert (Tensor(3) * (vv * 4)).item() == 24
+
+  def test_symbolic_mean(self):
+    vv = Variable("a", 1, 10).bind(2)
+    t = Tensor.ones(2, 10).contiguous()[:, :vv]
+    ret = t.mean().item()
+    assert ret == 1
+
+  def test_symbolic_mean_2d(self):
+    vv = Variable("a", 1, 10).bind(2)
+    vv2 = Variable("b", 1, 10).bind(2)
+    t = Tensor.ones(10, 10).contiguous()[:vv2, :vv]
+    ret = t.mean().item()
+    assert ret == 1
+
+  def test_symbolic_mean_2d_axis_1(self):
+    vv = Variable("a", 1, 10).bind(2)
+    vv2 = Variable("b", 1, 10).bind(2)
+    t = Tensor.ones(10, 10).contiguous()[:vv2, :vv]
+    ret = t.mean(axis=1)[:2].reshape(2, 1).numpy()
+    assert np.all(ret == 1)
+
+  def test_symbolic_mean_2d_add(self):
+    add_term = Variable("c", 0, 10).bind(1)
+    vv = Variable("a", 1, 10).bind(1)
+    vv2 = Variable("b", 1, 10).bind(1)
+    t = Tensor.ones(20, 20).contiguous()[:vv2+add_term, :vv+add_term]
+    ret = t.mean().item()
+    assert ret == 1
+
+  def test_symbolic_var(self):
+    vv = Variable("a", 1, 10).bind(2)
+    t = Tensor.ones(2, 10).contiguous()[:, :vv]
+    ret = t.var().item()
+    assert ret == 0
+
+  def test_symbolic_pad(self):
+    vv = Variable("a", 1, 10).bind(2)
+    t = Tensor.ones(2, 2).contiguous()
+    t = t.pad([vv, vv, vv, vv]).mean()
+    ones = 4
+    zeros = 6+6+4+4+6+6
+    self.assertAlmostEqual(t.item(), ones/(ones+zeros))
+
+  def test_symbolic_arange(self):
+    vv = Variable("a", 1, 10)
+    ret = Tensor.arange(0, vv.bind(4))
+    self.assertListEqual(ret[:4].tolist(), [0,1,2,3])
+
+  def test_symbolic_arange_sym_start(self):
+    vv = Variable("a", 1, 6)
+    ret = Tensor.arange(vv.bind(4), 7)
+    self.assertListEqual(ret[:3].tolist(), [4,5,6])
+
+  def test_symbolic_arange_sym_step(self):
+    vv = Variable("step", 1, 3)
+    ret = Tensor.arange(0, 10, vv.bind(2))
+    self.assertListEqual(ret[:5].tolist(), [0,2,4,6,8])
+
+  def test_symbolic_arange_two_vars(self):
+    begin = Variable("b", 1, 5)
+    end = Variable("e", 6, 10)
+    ret = Tensor.arange(begin.bind(4), end.bind(7))
+    self.assertListEqual(ret[:3].tolist(), [4,5,6])
+
+  def test_symbolic_arange_three_vars(self):
+    begin = Variable("b", 0, 5)
+    end = Variable("e", 10, 20)
+    step = Variable("s", 1, 3)
+    ret = Tensor.arange(begin.bind(2), end.bind(14), step.bind(3))
+    self.assertListEqual(ret[:4].tolist(), [2,5,8,11])
+
+  def test_symbolic_full(self):
+    vv = Variable("x", 1, 10).bind(5)
+    t = Tensor.full((3,), vv)
+    self.assertListEqual(t.tolist(), [5,5,5])
+
+  def test_variable_empty(self):
+    v = Variable("i", 1, 10)
+    # TODO: Tensor creation from unbound variable should assert
+    # with self.assertRaises(AssertionError): t = Tensor.empty(3, v)
+    vb = v.bind(3)
+    t = Tensor.empty(3, vb)
+    assert t.uop.base.buffer.size == 30
+    assert t.uop.shape == (3, vb)
+
+  def test_symbolic_chunk(self):
+    # chunk should work when split dimension is concrete, even if other dims are symbolic
+    vv = Variable("a", 1, 10).bind(4)
+    t = Tensor.ones(10, 8).contiguous()[:vv, :]  # shape (vv, 8)
+    chunks = t.chunk(2, dim=-1)  # split along concrete dim 8
+    assert len(chunks) == 2
+    assert chunks[0].shape[1] == 4
+    assert chunks[1].shape[1] == 4
+    # verify the values by shrinking to concrete shape first
+    np.testing.assert_equal(chunks[0].shrink(((0, 4), (0, 4))).numpy(), np.ones((4, 4)))
+    np.testing.assert_equal(chunks[1].shrink(((0, 4), (0, 4))).numpy(), np.ones((4, 4)))
+
+  def test_symbolic_split(self):
+    # split should work when split dimension is concrete, even if other dims are symbolic
+    vv = Variable("a", 1, 10).bind(3)
+    t = Tensor.arange(30).reshape(10, 3).contiguous()[:, :vv]  # shape (10, vv)
+    splits = t.split(5, dim=0)  # split along concrete dim 10
+    assert len(splits) == 2
+    assert splits[0].shape[0] == 5
+    assert splits[1].shape[0] == 5
+    # verify the values by shrinking to concrete shape first
+    np.testing.assert_equal(splits[0].shrink(((0, 5), (0, 3))).numpy(), np.arange(30).reshape(10, 3)[:5, :3])
+    np.testing.assert_equal(splits[1].shrink(((0, 5), (0, 3))).numpy(), np.arange(30).reshape(10, 3)[5:, :3])
+
+  def test_symbolic_chunk_error_on_symbolic_dim(self):
+    # chunk should fail when trying to split along a symbolic dimension
+    vv = Variable("a", 1, 10).bind(4)
+    t = Tensor.ones(10, 8).contiguous()[:vv, :]  # shape (vv, 8)
+    with self.assertRaises(AssertionError):
+      t.chunk(2, dim=0)  # can't split along symbolic dim
+
+  def test_symbolic_var_sum(self, var_name="u"):
+    t = Variable("t", 1, 10).bind(4)
+    v = Variable(var_name, 1, 5).bind(1)
+    mask = (Tensor.full((1, 1, t, v+t), 1) + 1).contiguous()
+    mask.shrink(((0, 1), (0, 1), (0, 4), (0, 4))).numpy()
+  def test_symbolic_var_sum_alt_name(self): self.test_symbolic_var_sum("s")
+
+  def test_symbolic_triu(self):
+    t = Variable("t", 1, 10).bind(4)
+    for start_pos in (0, 1, 3):
+      var_start_pos = Variable("start_pos", 0, 5).bind(start_pos)
+      mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).triu(var_start_pos+1)
+      out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy()
+      expected = np.triu(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1)
+      np.testing.assert_equal(out, expected)
+
+  def test_symbolic_tril(self):
+    t = Variable("t", 1, 10).bind(4)
+    for start_pos in (0, 1, 3):
+      var_start_pos = Variable("start_pos", 0, 5).bind(start_pos)
+      mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).tril(var_start_pos+1)
+      out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy()
+      expected = np.tril(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1)
+      np.testing.assert_equal(out, expected)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_to_numpy.py

@@ -0,0 +1,17 @@
+from tinygrad.tensor import Tensor
+import numpy as np
+import pickle
+import unittest
+
+class TestToNumpy(unittest.TestCase):
+  def test_numpy_is_numpy(self):
+    output = Tensor.ones((1, 3, 4096)).realize().numpy()
+    new = np.copy(output)
+    print(type(new))
+    serialized = pickle.dumps(new)
+    out = pickle.loads(serialized)
+    assert out.shape == (1,3,4096)
+    assert (out==1).all()
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_transcendental.py

@@ -0,0 +1,205 @@
+import unittest
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.tensor import _to_np_dtype
+from tinygrad.helpers import Context, getenv, DEV, OSX
+from test.backend.test_schedule import check_schedule
+from test.backend.test_dtype_alu import ht, dtypes_float
+import numpy as np
+import math
+from hypothesis import given, settings, strategies as strat
+
+settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+settings.load_profile("my_profile")
+
+supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
+
+class TestTranscendentalMath(unittest.TestCase):
+  @unittest.skipUnless(dtypes.float64 in supported_dtypes, f"no float64 on {Device.DEFAULT}")
+  @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
+  @given(ht.float64, strat.sampled_from([(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin)]))
+  def test_float64(self, x, op):
+    if op[0] == Tensor.sin:
+      # TODO: reduction does not work  # 536870912.125  # 2914593.01171875  # 134217728.03125  # 230581075.65625  # 139216373.71875
+      if abs(x) > 100_000_000: return
+    with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
+      np.testing.assert_allclose(op[0](Tensor([x], dtype=dtypes.float64)).numpy(),
+                                 op[1](np.array([x], dtype=_to_np_dtype(dtypes.float64))),
+                                 atol=3e-2, rtol=1e-5)  # sin can have bigger atol for very big x
+
+  @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
+  @given(ht.float32, strat.sampled_from([(Tensor.exp, np.exp),(Tensor.log, np.log)] +
+    ([(Tensor.sin, np.sin)] if dtypes.ulong in supported_dtypes else [])))
+  def test_float32(self, x, op):
+    # wrong nan behavior on Vulkan
+    if (math.isnan(x) or (x < 0 and op[0] == Tensor.log)) and Device.DEFAULT == "WEBGPU" and not OSX: return
+    with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
+      np.testing.assert_allclose(op[0](Tensor([x], dtype=dtypes.float32)).numpy(),
+                                 op[1](np.array([x], dtype=_to_np_dtype(dtypes.float32))),
+                                 atol=2e-5, rtol=1e-5)
+
+  @unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
+  @given(ht.float16, strat.sampled_from([(Tensor.exp, np.exp),(Tensor.log, np.log)] +
+    ([(Tensor.sin, np.sin)] if dtypes.ulong in supported_dtypes else [])))
+  def test_float16(self, x, op):
+    # wrong nan behavior on Vulkan
+    if (math.isnan(x) or (x < 0 and op[0] == Tensor.log)) and Device.DEFAULT == "WEBGPU" and not OSX: return
+    with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
+      np.testing.assert_allclose(op[0](Tensor([x], dtype=dtypes.float16)).numpy(),
+                                 op[1](np.array([x], dtype=_to_np_dtype(dtypes.float16))),
+                                 atol=1e-2, rtol=5e-3)  # exp can have bigger rtol
+
+  # TODO: WEBGPU produces incorrect values near infinity
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU incorrect values near inf")
+  @given(strat.sampled_from([(dtypes.float64, 709.5), (dtypes.float32, 88.7), (dtypes.float16, 11)]))
+  def test_exp_near_inf(self, dtype_x):
+    # reordering compute might return inf
+    dtype, x = dtype_x
+    if dtype not in supported_dtypes: return
+    with Context(TRANSCENDENTAL=2):
+      y = Tensor([x], dtype=dtype).exp().numpy()
+      expected = np.exp(np.array([x], dtype=_to_np_dtype(dtype)))
+      np.testing.assert_allclose(y, expected, rtol=5e-3)
+
+class TestFromFuzzer(unittest.TestCase):
+  @given(strat.sampled_from(dtypes_float))
+  @unittest.skipUnless(dtypes.ulong in supported_dtypes, "Needs ulong")
+  def test_sin(self, dtype):
+    if dtype not in supported_dtypes: return
+    if dtype == dtypes.float64:
+      # crashes in CI CUDA
+      if DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}: return
+    def _test_value(n: float, unit: float=1.0):
+      next_float = np.nextafter(1.0, 2.0, dtype=_to_np_dtype(dtype))
+      ulp = next_float - 1.0
+      ulp = unit * ulp
+      with Context(TRANSCENDENTAL=2):
+        np.testing.assert_allclose(Tensor([n], dtype=dtype).sin().numpy(), np.sin(np.array([n], dtype=_to_np_dtype(dtype))), atol=ulp, rtol=1e-5)
+    _test_value(-35.0)
+    _test_value(-25.0)
+    _test_value(25.0)
+    _test_value(30.0) # 30.0 == switch_over
+    _test_value(35.0)
+    _test_value(0.0)
+    _test_value(np.pi / 2)
+     # worst case of ulp 1.5
+    _test_value(np.pi * 2, unit=1.5)
+
+  @given(strat.sampled_from(dtypes_float))
+  def test_log2(self, dtype):
+    if dtype not in supported_dtypes: return
+    if dtype == dtypes.float64:
+      # crashes in CI CUDA
+      if DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}: return
+    def _test_value(n: float, unit: float=1.0):
+      next_float = np.nextafter(1.0, 2.0, dtype=_to_np_dtype(dtype))
+      ulp = next_float - 1.0
+      ulp = unit * ulp
+      with Context(TRANSCENDENTAL=2):
+        np.testing.assert_allclose(Tensor([n], dtype=dtype).log2().numpy(), np.log2(np.array([n], dtype=_to_np_dtype(dtype))), atol=ulp, rtol=1e-5)
+    fmin = np.finfo(_to_np_dtype(dtype)).tiny
+    for scale in [1.0, 1e10, 1e20, 1e30]:
+      _test_value(fmin * scale)
+      _test_value(-fmin * scale)
+    _test_value(0)
+    _test_value(0.0000009)
+
+class TestFloat16Log2(unittest.TestCase):
+  """Tests for native float16 log2 implementation (no float32 cast)"""
+  @unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
+  def test_float16_log2_basic(self):
+    # basic values
+    test_values = [1.0, 2.0, 4.0, 0.5, 0.25, 10.0, 100.0, 1000.0]
+    with Context(TRANSCENDENTAL=2):
+      for val in test_values:
+        result = Tensor([val], dtype=dtypes.float16).log2().numpy()[0]
+        expected = np.log2(np.float16(val))
+        np.testing.assert_allclose(result, expected, rtol=1e-3, err_msg=f"log2({val})")
+
+  @unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "Nan handling differs on Vulkan")
+  def test_float16_log2_special(self):
+    # special values: inf, -inf, nan, 0, negative
+    with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
+      # log2(inf) = inf
+      assert np.isinf(Tensor([np.inf], dtype=dtypes.float16).log2().numpy()[0])
+      # log2(0) = -inf
+      assert Tensor([0.0], dtype=dtypes.float16).log2().numpy()[0] == -np.inf
+      # log2(negative) = nan
+      assert np.isnan(Tensor([-1.0], dtype=dtypes.float16).log2().numpy()[0])
+      # log2(nan) = nan
+      assert np.isnan(Tensor([np.nan], dtype=dtypes.float16).log2().numpy()[0])
+
+  @unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
+  def test_float16_log2_denormal(self):
+    # test values near and below float16 min normal (6.1e-5)
+    # these exercise the denormal handling path with 2^10 scaling
+    test_values = [1e-4, 6e-5, 1e-5]
+    with Context(TRANSCENDENTAL=2):
+      for val in test_values:
+        result = Tensor([val], dtype=dtypes.float16).log2().numpy()[0]
+        expected = np.log2(np.float16(val))
+        # denormals have lower precision due to float16 limitations
+        np.testing.assert_allclose(result, expected, rtol=5e-2, err_msg=f"log2({val})")
+
+class TestTranscendentalSchedule(unittest.TestCase):
+  @unittest.skipUnless(dtypes.ulong in supported_dtypes, "Needs ulong")
+  def test_transcendental_sin_fusion(self):
+    with Context(TRANSCENDENTAL=2):
+      a = Tensor.empty(10)
+      b = Tensor.empty(10)
+      c = a.sin() + b.sin()
+      c = c.sin()
+      check_schedule(c, 1)
+
+  def test_transcendental_log2_fusion(self):
+    with Context(TRANSCENDENTAL=2):
+      a = Tensor.empty(10)
+      b = Tensor.empty(10)
+      c = a.log2() + b.log2()
+      c = c.log2()
+      check_schedule(c, 1)
+
+  def test_transcendental_exp2_fusion(self):
+    with Context(TRANSCENDENTAL=2):
+      a = Tensor.empty(10)
+      b = Tensor.empty(10)
+      c = a.exp2() + b.exp2()
+      c = c.exp2()
+      check_schedule(c, 1)
+
+class TestTranscendentalVectorized(unittest.TestCase):
+  def _vectorized_data(self, low, high, vec_size):
+    np_data = np.linspace(low, high, num=(128 // vec_size) * vec_size, dtype=np.float32).reshape(-1, vec_size)
+    data = Tensor(np_data, dtype=dtypes.float32.vec(vec_size))
+    return data, np_data
+
+  def _test_vectorized_op(self, fxn, np_fxn, data_range, vec_size, param_range=None):
+    data, np_data = self._vectorized_data(data_range[0], data_range[1], vec_size)
+    if param_range:
+      param, np_param = self._vectorized_data(param_range[0], param_range[1], vec_size)
+      out, np_out = fxn(data, param), np_fxn(np_data, np_param)
+    else:
+      out, np_out = fxn(data), np_fxn(np_data)
+    np.testing.assert_allclose(out.numpy(), np_out, rtol=1e-4)
+
+  def test_exp2_vectorized(self):
+    for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.exp2, np.exp2, (-100, 100), vec_size)
+
+  def test_log2_vectorized(self):
+    for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.log2, np.log2, (0.001, 200), vec_size)
+
+  @unittest.skipIf(Device.DEFAULT == "DSP", "requires int division")
+  @unittest.skipIf(DEV.renderer == "NAK", "MUFU.SIN is not accurate enough")
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and OSX, "WEBGPU Metal backend is not accurate enough")
+  def test_sin_vectorized(self):
+    for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.sin, np.sin, (-100, 100), vec_size)
+
+  def test_pow_vectorized(self):
+    # np.pow returns nan for negative values raised to a non-integral power
+    for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.pow, np.pow, (0.001, 200), vec_size, param_range=(-10, 10))
+
+  def test_sqrt_vectorized(self):
+    for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.sqrt, np.sqrt, (0, 100), vec_size)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/backend/test_uops.py

@@ -0,0 +1,346 @@
+from typing import Optional, Any
+import unittest, math
+import numpy as np
+from tinygrad.tensor import Tensor, _to_np_dtype
+from tinygrad.helpers import Context
+from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat  # noqa: F401
+from tinygrad.device import Buffer, Device
+from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType, buffers
+from tinygrad.renderer.cstyle import CStyleLanguage
+from tinygrad.engine.realize import run_linear
+from tinygrad.codegen import to_program
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.renderer.ptx import PTXRenderer
+from test.helpers import to_uops_list
+
+def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
+  buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in bufs]
+  for u,b in zip(buf_uops, bufs): buffers[u] = b
+  run_linear(UOp(Ops.LINEAR, src=(UOp.sink(*uops_list, arg=KernelInfo()).call(*buf_uops),)))
+
+def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
+  if op is Ops.CONST: uops.append(UOp.const(dtype, arg))
+  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype).replace(src=()))
+  else: uops.append(UOp(op, dtype, tuple(src), arg))
+  return uops[-1]
+
+def _test_single_value(vals, op, dts):
+  uops = []
+  output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
+  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
+  buf_loads = [uop(uops, Ops.PARAM, dtype.ptr(), (), i+1) for i,dtype in enumerate(dts)]
+  loads = (buf_loads[i].index(uop(uops, Ops.CONST, dtypes.int32, (), 0)) for i, dtype in enumerate(dts))
+  alu = uop(uops, op, output_dtype, loads)
+  out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True), alu))
+  buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
+  buf2 = [Buffer(Device.DEFAULT, 1, dtype).allocate().copyin(np.array([a], dtype=_to_np_dtype(dtype)).data) for a,dtype in zip(vals, dts)]
+  run_uops([out], [buf]+buf2)
+  ret = np.empty(1, _to_np_dtype(output_dtype))
+  buf.copyout(ret.data)
+  return ret[0]
+
+def _test_single_value_const(vals, op, dts):
+  uops = []
+  output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
+  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
+  loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts))
+  alu = uop(uops, op, output_dtype, loads)
+  out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
+  buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
+  run_uops([out], [buf])
+  ret = np.empty(1, _to_np_dtype(output_dtype))
+  buf.copyout(ret.data)
+  return ret[0]
+
+def _test_uops_result(output_dtype, uops, res):
+  # uops = []
+  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
+  # res = output_fn(uops)
+  out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), res))
+  buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
+  run_uops([out], [buf])
+  ret = np.empty(1, _to_np_dtype(output_dtype))
+  buf.copyout(ret.data)
+  return ret[0]
+
+class TestUOps(unittest.TestCase):
+  def _equal(self, v1, v2):
+    assert isinstance(v2, (float, int, bool))
+    if isinstance(v2, float):
+      np.testing.assert_allclose(v1, v2, rtol=2e-7)
+    else:
+      np.testing.assert_equal(v1, v2)
+
+  def _test_uop_fxn(self, op, fxn, dts=(dtypes.float32, )):
+    for f in [_test_single_value, _test_single_value_const]:
+      for a in [-2.0, 0.0, 1.0]:
+        a = dts[0].const(a)
+        self._equal(f([a], op, dts), fxn(a))
+
+  def _test_bop_fxn(self, op, fxn, dts=(dtypes.float32, )*2, no_b_zero=False, no_b_neg=False):
+    for f in [_test_single_value, _test_single_value_const]:
+      for a in [-2.0, 0.0, 1.0]:
+        for b in [-3.0, 1.0] + ([] if no_b_zero else [0.0]):
+          a = dts[0].const(a)
+          b = dts[1].const(abs(b) if no_b_neg else b)
+          self._equal(f([a,b], op, dts), fxn(a,b))
+
+  def _test_top_fxn(self, op, fxn, dts=(dtypes.float32, )*3):
+    for f in [_test_single_value, _test_single_value_const]:
+      for a in [-2.0, 0, 1]:
+        for b in [-3.0, 3.0]:
+          for c in [-4.0, 4.0]:
+            a = dts[0].const(a)
+            b = dts[1].const(b)
+            c = dts[2].const(c)
+            self._equal(f([a,b,c], op, dts), fxn(a,b,c))
+
+class TestFloatUOps(TestUOps):
+  @unittest.skipIf(Device.DEFAULT == "CPU", 'not supported as uop')
+  def test_exp2(self): self._test_uop_fxn(Ops.EXP2, lambda a: np.exp2(a))
+  @unittest.skipIf(Device.DEFAULT == "CPU", 'not supported as uop')
+  def test_log2(self): self._test_uop_fxn(Ops.LOG2, lambda a: math.log2(a) if a > 0 else float('-inf' if a==0 else 'nan'))
+  @unittest.skipIf(Device.DEFAULT == "CPU", 'not supported as uop')
+  def test_sin(self): self._test_uop_fxn(Ops.SIN, lambda a: math.sin(a))
+  def test_recip(self): self._test_uop_fxn(Ops.RECIPROCAL, lambda a: 1/a if a != 0 else float('inf'))
+  def test_sqrt(self): self._test_uop_fxn(Ops.SQRT, lambda a: math.sqrt(a) if a >= 0 else float('nan'))
+
+  def test_add(self): self._test_bop_fxn(Ops.ADD, lambda a,b: a+b)
+  def test_mul(self): self._test_bop_fxn(Ops.MUL, lambda a,b: a*b)
+  def test_max(self): self._test_bop_fxn(Ops.MAX, lambda a,b: max(a,b))
+  def test_cmplt(self): self._test_bop_fxn(Ops.CMPLT, lambda a,b: a<b)
+  def test_cmpne(self): self._test_bop_fxn(Ops.CMPNE, lambda a,b: a!=b)
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support NaN comparison correctly")
+  def test_cmpne_nan(self):  # NaN != x for any x (IEEE 754)
+    for a, b in [(math.nan, 1.0), (1.0, math.nan), (math.nan, math.nan)]:
+      self.assertTrue(_test_single_value(
+        [dtypes.float32.const(a), dtypes.float32.const(b)],
+        Ops.CMPNE, (dtypes.float32, dtypes.float32)))
+  # MOD isn't tested on floats
+
+  def test_where(self):
+    self._test_top_fxn(Ops.WHERE, lambda a,b,c: b if a!=0 else c, (dtypes.bool, dtypes.float, dtypes.float))
+
+  @unittest.skipUnless(Device.DEFAULT == "PYTHON", "only python supports MULACC")
+  def test_mulacc(self):
+    self._test_top_fxn(Ops.MULACC, lambda a,b,c: a*b+c, (dtypes.float, dtypes.float, dtypes.float))
+
+class TestNonFloatUOps(TestUOps):
+  def test_add_int32(self): self._test_bop_fxn(Ops.ADD, lambda a,b: int(a)+int(b), (dtypes.int32, dtypes.int32))
+  def test_mul_int32(self): self._test_bop_fxn(Ops.MUL, lambda a,b: int(a)*int(b), (dtypes.int32, dtypes.int32))
+  @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CStyleLanguage)), "only ptx and cstyle use bitshifts")
+  def test_shr_int32(self): self._test_bop_fxn(Ops.SHR, lambda a,b: int(a)>>int(b), (dtypes.int32, dtypes.int32), no_b_neg=True)
+  @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CStyleLanguage)), "only ptx and cstyle use bitshifts")
+  def test_shl_int32(self): self._test_bop_fxn(Ops.SHL, lambda a,b: int(a)<<int(b), (dtypes.int32, dtypes.int32), no_b_neg=True)
+  def test_div_int32(self):
+    self._test_bop_fxn(Ops.CDIV, lambda a,b: int(a/b), (dtypes.int32, dtypes.int32), no_b_zero=True)
+  def test_and_int32(self): self._test_bop_fxn(Ops.AND, lambda a,b: int(a)&int(b), (dtypes.int32, dtypes.int32))
+  def test_or_int32(self): self._test_bop_fxn(Ops.OR, lambda a,b: int(a)|int(b), (dtypes.int32, dtypes.int32))
+  def test_mod_int32(self):
+    self._test_bop_fxn(Ops.CMOD,
+                       lambda a,b: abs(int(a))%abs(int(b))*(1,-1)[a<0], (dtypes.int32, dtypes.int32), no_b_zero=True)
+  def test_cmplt_int32(self): self._test_bop_fxn(Ops.CMPLT, lambda a,b: int(a)<int(b), (dtypes.int32, dtypes.int32))
+  def test_cmpne_int32(self): self._test_bop_fxn(Ops.CMPNE, lambda a,b: int(a)!=int(b), (dtypes.int32, dtypes.int32))
+  def test_mul_bool(self): self._test_bop_fxn(Ops.MUL, lambda a,b: bool(a) and bool(b), (dtypes.bool, dtypes.bool))
+  def test_where_float16(self):
+    self._test_top_fxn(Ops.WHERE, lambda a,b,c: b if a!=0 else c, (dtypes.bool, dtypes.float16, dtypes.float16))
+
+class TestBoolUOps(TestUOps):
+  def _test_uop_bool_fxn(self, op, fxn):
+    for f in [_test_single_value, _test_single_value_const]:
+      for a in [False, True]:
+        self._equal(f([a], op, (dtypes.bool, )*1), fxn(a))
+
+  def _test_bop_bool_fxn(self, op, fxn):
+    for f in [_test_single_value, _test_single_value_const]:
+      for a in [False, True]:
+        for b in [False, True]:
+          self._equal(f([a,b], op, (dtypes.bool, )*2), fxn(a,b))
+
+  def _test_top_bool_fxn(self, op, fxn):
+    for f in [_test_single_value, _test_single_value_const]:
+      for a in [False, True]:
+        for b in [False, True]:
+          for c in [False, True]:
+            self._equal(f([a,b,c], op, (dtypes.bool, )*3), fxn(a,b,c))
+
+  def test_add_bool(self): self._test_bop_bool_fxn(Ops.ADD, lambda a,b: a or b)
+  def test_mul_bool(self): self._test_bop_bool_fxn(Ops.MUL, lambda a,b: a and b)
+  def test_xor_bool(self): self._test_bop_bool_fxn(Ops.XOR, lambda a,b: a != b)
+  def test_and_bool(self): self._test_bop_bool_fxn(Ops.AND, lambda a,b: a & b)
+  def test_or_bool(self): self._test_bop_bool_fxn(Ops.OR, lambda a,b: a | b)
+  def test_cmpne_bool(self): self._test_bop_bool_fxn(Ops.CMPNE, lambda a,b: a != b)
+  def test_cmplt_bool(self): self._test_bop_bool_fxn(Ops.CMPLT, lambda a,b: a < b)
+  def test_where_bool(self): self._test_top_bool_fxn(Ops.WHERE, lambda a,b,c: b if a else c)
+
+class TestLocalAccess(unittest.TestCase):
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared memory")
+  def test_local_basic(self):
+    uops = []
+    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.float32.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
+    st = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), uop(uops, Ops.CONST, dtypes.float32, (), 42.0)))
+    barr = uop(uops, Ops.BARRIER, dtypes.void, (st,))
+    sres = uop(uops, Ops.LOAD, dtypes.float32, (smem.after(barr).index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True),))
+    self.assertEqual(_test_uops_result(dtypes.float32, uops, sres), 42)
+
+  # NOTE: webgpu specific, since only webgpu performs bitpacking
+  @unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Test local access with packed data type")
+  def test_local_packed(self):
+    uops = []
+    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
+    st = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), uop(uops, Ops.CONST, dtypes.uint8, (), 42)))
+    barr = uop(uops, Ops.BARRIER, dtypes.void, (st,))
+    sres = smem.after(barr).index(uop(uops, Ops.CONST, dtypes.int32, (), 0))
+    self.assertEqual(_test_uops_result(dtypes.uint8, uops, sres), 42)
+
+  # NOTE: webgpu specific, since only webgpu performs bitpacking
+  @unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Test local memory size for packed data types")
+  def test_packed_smem_size(self):
+    _dtypes = [dtypes.char, dtypes.uchar, dtypes.short, dtypes.ushort, dtypes.half]
+    size = 16
+    for dtype in _dtypes:
+      temp = UOp(Ops.DEFINE_LOCAL, dtype.ptr(size=size, addrspace=AddrSpace.LOCAL), (), 'smem')
+      uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
+      out = Device[Device.DEFAULT].renderer.render(uops)
+      # half is supported in wgsl, so it doesn't have to be packed
+      corrected_size = size//(4//dtype.itemsize) if dtype != dtypes.half else size
+      self.assertIn(f"temp0: array<{Device[Device.DEFAULT].renderer.buf_map(dtype)},{corrected_size}>;", out)
+
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared memory")
+  @unittest.skip("tinygrad doesn't support this behavior")
+  def test_local_indirect(self):
+    uops = []
+    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.int32.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
+    st1 = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 1)), uop(uops, Ops.CONST, dtypes.int32, (), 2)))
+    st2 = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 2)), uop(uops, Ops.CONST, dtypes.int32, (), 42)))
+    barr = uop(uops, Ops.BARRIER, dtypes.void, (st1,st2))
+    ofs = uop(uops, Ops.LOAD, dtypes.int32, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 1)), barr))
+    sres = uop(uops, Ops.LOAD, dtypes.int32, (smem.index(ofs),))
+    self.assertEqual(_test_uops_result(dtypes.int32, uops, sres), 42)
+
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends")
+class TestAssembly(unittest.TestCase):
+  def test_bitshift_left(self):
+    g1 = UOp.param(0, dtypes.int32.ptr())
+    out = UOp.param(1, dtypes.int32.ptr())
+    c1 = UOp.const(dtypes.int, 2)
+    c2 = UOp.const(dtypes.int, 3)
+    l1 = g1.index(c1)
+    a1 = UOp(Ops.MUL, dtypes.int, (l1, c1))
+    a2 = UOp(Ops.MUL, dtypes.int, (l1, c2))
+    uops = to_uops_list([out.index(UOp.const(dtypes.int, 0)).store(a1), out.index(UOp.const(dtypes.int, 1)).store(a2)],
+                        ren=Device[Device.DEFAULT].renderer)
+    Device[Device.DEFAULT].renderer.render(uops)
+    ops = [x.op for x in uops]
+    self.assertIn(Ops.SHL, ops)
+    self.assertIn(Ops.MUL, ops)
+
+  def test_mulacc_unrolled(self):
+    # test that     acc = acc + a0*b0 + a1*b1 + a2*b2 + a3*b3
+    # is not        acc = acc + (a0*b0 + a1*b1 + a2*b2 + a3*b3)
+    a = Tensor.empty(1024)
+    b = Tensor.empty(1024)
+    c = (a*b).sum()
+    ast = c.schedule_linear().src[-1].src[0]
+    opts_to_apply = [Opt(OptOps.UNROLL, 0, 4)]
+    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
+    program = to_program(ast, Device[Device.DEFAULT].renderer)
+    uops = tuple(program.src[2].src)
+    self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)
+
+  def test_mulacc_shl(self):
+    g1 = UOp.param(0, dtypes.int32.ptr())
+    c1 = UOp.const(dtypes.int, 0)
+    c2 = UOp.const(dtypes.int, 1)
+    expr = g1.index(c1) * UOp.const(dtypes.int, 4096) + g1.index(c2)
+    uops = to_uops_list([expr], ren=Device[Device.DEFAULT].renderer)
+    Device[Device.DEFAULT].renderer.render(uops)
+    self.assertIn(Ops.MULACC, [x.op for x in uops])
+
+  def test_use_cmpeq(self):
+    g = UOp.param(0, dtypes.uint32.ptr())
+    c = UOp.const(dtypes.uint, 7)
+    comp = g.index(c).ne(c).ne(True)
+    uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
+    Device[Device.DEFAULT].renderer.render(uops)
+    ops = [x.op for x in uops]
+    self.assertIn(Ops.CMPEQ, ops)
+    self.assertNotIn(Ops.CMPNE, ops)
+
+class TestZeroRange(unittest.TestCase):
+  def test_reduce_variable(self):
+    for i in range(3,-1,-1):
+      v = UOp.variable("i", 0, 5).bind(i)
+      out = Tensor.ones(10, dtype=dtypes.int).contiguous().shrink(((0,v),)).sum()
+      self.assertEqual(out.item(), i)
+
+class TestUOpPrograms(unittest.TestCase):
+  def _run(self, prog:UOp, *tensors:Tensor):
+    run_linear(UOp(Ops.LINEAR, src=(prog.call(*[t.uop.buf_uop for t in tensors]),)), update_stats=False)
+
+  def test_simple(self):
+    out = Tensor.empty(10,10,dtype=dtypes.int)
+
+    ptr = UOp.placeholder(out.shape, out.dtype, slot=0)
+    i, j = UOp.range(10, axis_id=0), UOp.range(10, axis_id=1)
+    prog = ptr[i,j].set(42).end(i,j)
+    self._run(prog.sink(arg=KernelInfo()), out)
+
+    with Context(DEBUG=0): self.assertTrue((out == 42).all().item())
+
+  def test_matmul(self):
+    a = Tensor.randn(10,10)
+    b = Tensor.randn(10,10)
+    c = Tensor.empty(10,10)
+    ref = (a@b)
+    with Context(DEBUG=0): Tensor.realize(a, b, c, ref)
+
+    # C[i,j] = sum_k A[i,k] * B[k,j]
+    # Shapes: A[M,K], B[K,N], C[M,N]
+    M = N = K = 10
+    DT = dtypes.float32
+
+    # Placeholders (bind slots explicitly)
+    A = UOp.placeholder((M, K), DT, slot=0)
+    B = UOp.placeholder((K, N), DT, slot=1)
+    C = UOp.placeholder((M, N), DT, slot=2)
+
+    # Axes: i,j are spatial; k is a reduction axis over the shared dim K
+    i = UOp.range(M, axis_id=0)                             # rows of A/C
+    j = UOp.range(N, axis_id=1)                             # cols of B/C
+    k = UOp.range(K, axis_id=2, axis_type=AxisType.REDUCE)  # reduction over K
+
+    # Zero-init: write a scalar 0 to each (i,j).
+    C = C[i, j].set(0.0)
+
+    # Accumulate: C_after(k) enforces the dependency along the reduction axis
+    C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j])
+
+    # Finalize the loop nest / schedule in (i, j, k) order
+    prog = C.end(i, j, k)
+
+    # run program
+    # TODO: make this work with opts_to_apply
+    self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
+
+    with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)
+
+  def test_matmul_relu(self):
+    a, b, c = Tensor.randn(10,10), Tensor.randn(10,10), Tensor.empty(10,10)
+    ref = (a@b).relu()
+    with Context(DEBUG=0): Tensor.realize(a, b, c, ref)
+
+    A, B, C = a.uop.placeholder_like(0), b.uop.placeholder_like(1), c.uop.placeholder_like(2)
+    i, j, k = UOp.range(10, 0), UOp.range(10, 1), UOp.range(10, 2, axis_type=AxisType.REDUCE)
+
+    C = C[i, j].set(0.0)
+    C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
+    C = C[i, j].set(C[i, j].maximum(0.0))
+
+    prog = C.end(i, j)
+
+    self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
+    with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)

tinygrad_repo/test/backend/test_zero_copy.py

@@ -0,0 +1,27 @@
+import unittest
+from tinygrad import Tensor, Device
+import time
+
+def time_tensor_numpy(out:Tensor):
+  times = []
+  for _ in range(5):
+    st = time.perf_counter()
+    out.uop.base.realized.as_memoryview(allow_zero_copy=True)
+    et = time.perf_counter() - st
+    times.append(et)
+  return min(times)
+
+N = 4096
+class TestZeroCopy(unittest.TestCase):
+  @unittest.skipIf(Device.DEFAULT not in {"CPU", "METAL"}, "device isn't zero copy")
+  def test_zero_copy_from_default_to_cpu(self):
+    demo = Tensor.rand(1).realize()
+    t1 = time_tensor_numpy(demo)
+    out = Tensor.rand(N, N).realize()
+    t2 = time_tensor_numpy(out)
+    gbps = out.nbytes()*1e-9/max(t2-t1, 1e-10)
+    print(f"time(base): {t1*1e3:.2f} ms, time(copy): {t2*1e3:.2f} ms :  copy speed {gbps:.2f} GB/s")
+    self.assertGreater(gbps, 600)  # more than 600 GB/s = no copy
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)