IQ.Pilot Release Commit @ 5bc9cd3

tinygrad_repo/test/null/test_attention.py

@@ -0,0 +1,36 @@
+import unittest
+from tinygrad import Tensor, dtypes, TinyJit, UOp
+from tinygrad.llm.model import apply_rope as apply_rope_new, precompute_freqs_cis
+from test.helpers import assert_jit_cache_len
+
+def apply_rope(x:Tensor, start_pos:int):
+  B, H, T, Hd = x.shape
+  precompute_freqs_cis.cache_clear()
+  freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
+  return apply_rope_new(x, freqs_cis)
+
+class TestAttention(unittest.TestCase):
+  def test_half_qkv_buffers(self):
+    BS, seqlen, dim = 10, 4, 100
+    q = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
+    k = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
+    v = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
+    attn = q.scaled_dot_product_attention(k, v)
+    sched = attn.schedule_linear()
+    # attention has 4 kernels now
+    self.assertEqual(len(sched.src), 4)
+
+  def test_apply_rope_jit_prune(self):
+    def rope_fn(x_in, pos): return apply_rope(x_in, pos)
+    rope_noprune = TinyJit(rope_fn)
+    rope_prune = TinyJit(rope_fn, prune=True)
+
+    v_pos = UOp.variable("start_pos", 0, 100)
+    for _ in range(3):
+      rope_noprune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1))
+      rope_prune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1))
+    assert_jit_cache_len(rope_prune, 1)
+    assert_jit_cache_len(rope_noprune, 3)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_autogen.py

@@ -0,0 +1,535 @@
+import ctypes, struct, subprocess, tempfile, unittest
+from tinygrad.helpers import OSX, WIN
+from tinygrad.runtime.support.c import DLL, record, Field
+from tinygrad.runtime.support import c
+from tinygrad.runtime.support.autogen import gen
+
+@unittest.skipIf(WIN, "doesn't compile on windows")
+class TestC(unittest.TestCase):
+  def compile(self, src):
+    with tempfile.NamedTemporaryFile(suffix=".so") as f:
+      subprocess.check_output(('clang', '-x', 'c', '-fPIC', '-shared', '-', '-o', f.name), input=src.encode())
+      return DLL("test", f.name)
+
+  def test_struct_array_init(self):
+    @record
+    class Foo(c.Struct):
+      SIZE = 12
+      a = Field(ctypes.c_int * 3, 0)
+
+    f = Foo((1,2,3))
+    assert f.a[0] == 1
+    assert f.a[1] == 2
+    assert f.a[2] == 3
+    f = Foo((ctypes.c_int * 3)(1,2,3))
+    assert f.a[0] == 1
+    assert f.a[1] == 2
+    assert f.a[2] == 3
+
+  def test_field_ranges(self):
+    @record
+    class Foo(c.Struct):
+      SIZE = 2
+      s = Field(ctypes.c_int8, 0)
+      u = Field(ctypes.c_uint8, 1)
+
+    f = Foo()
+    f.s = -1
+    f.u = -1
+    assert f.s == -1
+    assert f.u == 255
+
+  # this syntax is inherited from ctypes, but it seems a bit nonsensical?
+  def test_voidp_none(self):
+    @record
+    class Foo(c.Struct):
+      SIZE = 8
+      p = Field(ctypes.c_void_p, 0)
+
+    f = Foo(None)
+    assert f.p is None
+    f.p = ctypes.c_void_p(0xDEADBEEF)
+    assert f.p == 0xDEADBEEF
+    f.p = None
+    assert f.p is None
+
+  def test_packed_struct(self):
+    @record
+    class Baz(c.Struct):
+      SIZE = 8
+      a = Field(ctypes.c_uint, 0, 30)
+      b = Field(ctypes.c_uint, 3, 30, 6)
+      c = Field(ctypes.c_uint, 7, 2, 4)
+      d = Field(ctypes.c_uint, 7, 2, 6)
+
+    b = Baz(0x3AAADEAD, 0xBEEF, 1, 0)
+    assert b.a == 0x3AAADEAD
+    assert b.b == 0xBEEF
+    assert b.c == 1
+    assert b.d == 0
+
+    b.a = 0xCAFE
+    assert b.a == 0xCAFE
+    assert b.b == 0xBEEF
+    assert b.c == 1
+    assert b.d == 0
+
+  def test_packed_struct_interop(self):
+    @record
+    class Baz(c.Struct):
+      SIZE = 8
+      a = Field(ctypes.c_int, 0, 30)
+      b = Field(ctypes.c_int, 3, 30, 6)
+      c = Field(ctypes.c_int, 7, 2, 4)
+      d = Field(ctypes.c_int, 7, 2, 6)
+
+    src = '''
+      struct __attribute__((packed)) baz {
+        int a:30;
+        int b:30;
+        int c:2;
+        int d:2;
+      };
+
+      int test(struct baz x) {
+        return x.a + x.b + x.c + x.d;
+      }
+    '''
+    dll = self.compile(src)
+    b = Baz(0xAA000, 0x00BB0, 0, 1)
+    @dll.bind(ctypes.c_int, Baz)
+    def test(x:Baz) -> ctypes.c_int: ...
+    self.assertEqual(test(b), b.a + b.b + b.c + b.d)
+
+  # https://github.com/python/cpython/issues/90914
+  def test_bitfield_interop(self):
+    @record
+    class Baz(c.Struct):
+      SIZE = 1
+      a = Field(ctypes.c_bool, 0, 1, 0)
+      b = Field(ctypes.c_bool, 0, 1, 1)
+      c = Field(ctypes.c_bool, 0, 1, 2)
+      d = Field(ctypes.c_bool, 0, 1, 3)
+      e = Field(ctypes.c_bool, 0, 1, 4)
+      f = Field(ctypes.c_bool, 0, 1, 5)
+      g = Field(ctypes.c_bool, 0, 1, 6)
+      h = Field(ctypes.c_bool, 0, 1, 7)
+    src = '''#include <stdbool.h>
+      struct baz {
+        bool a:1, b:1, c:1, d:1, e:1, f:1, g:1, h:1;
+      };
+
+      int test(struct baz x) {
+        return x.c;
+      }
+    '''
+    dll = self.compile(src)
+    @dll.bind(ctypes.c_int, Baz)
+    def test(x:Baz) -> ctypes.c_int: ...
+    for i in range(8): self.assertEqual(test(Baz(*(j==i for j in range(8)))), i==2)
+
+  def test_struct_interop(self):
+    @record
+    class Baz(c.Struct):
+      SIZE = 32
+      a = Field(ctypes.c_int, 0)
+      b = Field(ctypes.c_int, 4)
+      c = Field(ctypes.c_int, 8)
+      d = Field(ctypes.c_int, 12)
+      e = Field(ctypes.c_int, 16)
+      f = Field(ctypes.c_int, 20)
+      g = Field(ctypes.c_int, 24)
+      h = Field(ctypes.c_int, 28)
+    src = '''#include <stdio.h>
+      struct baz {
+        int a, b, c, d, e, f, g, h;
+      };
+
+      struct baz test(struct baz x) {
+        return (struct baz){x.h, x.g, x.f, x.e, x.d, x.c, x.b, x.a};
+      }
+    '''
+    dll = self.compile(src)
+    @dll.bind(Baz, Baz)
+    def test(x:Baz) -> Baz: ...
+    self.assertEqual(bytes(test(Baz(*range(8)))), struct.pack("8i", *range(7, -1, -1)))
+
+  def test_aos_interop(self):
+    @record
+    class Item(c.Struct):
+      SIZE = 4
+      val = Field(ctypes.c_int, 0)
+    src = """
+    struct item { int val; };
+      int test(struct item arr[3]) {
+        int ret = 0;
+        for (int i = 0; i < 3; i++) ret += arr[i].val;
+        return ret;
+      }
+    """
+    dll = self.compile(src)
+    @dll.bind(ctypes.c_int, Item * 3)
+    def test(arr:(Item * 3)) -> ctypes.c_int: ...
+    self.assertEqual(test((Item * 3)(Item(10), Item(20), Item(30))), 60)
+
+  def test_soa_interop(self):
+    @record
+    class Row(c.Struct):
+      SIZE = 16
+      data = Field(ctypes.c_int * 3, 0)
+    src = """
+      struct row { int data[3]; };
+      struct row test(struct row x) {
+        return (struct row){{ x.data[2], x.data[1], x.data[0] }};
+      }
+    """
+    dll = self.compile(src)
+    @dll.bind(Row, Row)
+    def test(x:Row) -> Row: ...
+    r = test(Row((ctypes.c_int * 3)(10, 20, 30)))
+    self.assertIsInstance(r, Row)
+    self.assertEqual(r.data[0], 30)
+    self.assertEqual(r.data[1], 20)
+    self.assertEqual(r.data[2], 10)
+
+  def test_soa_ptr_interop(self):
+    @record
+    class Row(c.Struct):
+      SIZE = 8
+      data = Field(c.POINTER[ctypes.c_int], 0)
+    src = """
+      struct row { int *data; };
+      int test(struct row x) {
+        return x.data[2] + x.data[1] + x.data[0];
+      }
+    """
+    dll = self.compile(src)
+    @dll.bind(ctypes.c_int, Row)
+    def test(x:Row) -> ctypes.c_int: ...
+    assert test(Row((ctypes.c_int * 3)(10, 20, 30))) == 60
+
+  def test_nested_struct_interop(self):
+    @record
+    class Inner(c.Struct):
+      SIZE = 4
+      a = Field(ctypes.c_int, 0)
+    @record
+    class Outer(c.Struct):
+      SIZE = 8
+      inner = Field(Inner, 0)
+      b = Field(ctypes.c_int, 4)
+    src = """
+      struct i { int a; };
+      struct o { struct i i; int b; };
+      struct o test(struct o x) {
+        return (struct o){(struct i){ x.b }, x.i.a };
+      }
+    """
+    dll = self.compile(src)
+    @dll.bind(Outer, Outer)
+    def test(x:Outer) -> Outer: ...
+    o = test(Outer(Inner(10), 20))
+    self.assertEqual(o.inner.a, 20)
+    self.assertEqual(o.b, 10)
+
+  def test_struct_pointer_interop(self):
+    @record
+    class Foo(c.Struct):
+      SIZE = 8
+      a = Field(ctypes.c_int, 0)
+      b = Field(ctypes.c_int, 4)
+    src = """
+      struct foo { int a, b; };
+      struct foo *test(struct foo *f) {
+        int x = f->a;
+        f->a = f->b;
+        f->b = x;
+        return f;
+      }
+    """
+    dll = self.compile(src)
+    @dll.bind(ctypes.POINTER(Foo), ctypes.POINTER(Foo))
+    def test(f:ctypes.POINTER(Foo)) -> ctypes.POINTER(Foo): ...
+    inp = ctypes.pointer(Foo(10, 20))
+    out = test(inp)
+    self.assertEqual(out.contents.a, 20)
+    self.assertEqual(out.contents.b, 10)
+
+  def test_pointer_field_roundtrip(self):
+    # This tests storing a pointer in a record struct field and passing it to C
+    # Mimics how mesa.struct_lp_build_tgsi_params.mask is used
+    from tinygrad.runtime.support.c import POINTER
+    @record
+    class Inner(c.Struct):
+      SIZE = 8
+      value = Field(ctypes.c_int, 0)
+      flag = Field(ctypes.c_int, 4)
+    @record
+    class Outer(c.Struct):
+      SIZE = 16
+      x = Field(ctypes.c_int, 0)
+      inner_ptr = Field(POINTER[Inner], 8)
+
+    src = """
+      struct inner { int value; int flag; };
+      struct outer { int x; struct inner *inner_ptr; };
+      int test(struct inner *p) {
+        return p->value + p->flag;
+      }
+    """
+    dll = self.compile(src)
+    @dll.bind(ctypes.c_int, ctypes.POINTER(Inner))
+    def test(p:POINTER[Inner]) -> ctypes.c_int: ...
+
+    inner = Inner(value=42, flag=10)
+    outer = Outer(x=1, inner_ptr=ctypes.pointer(inner))
+    # Retrieve pointer from struct field and pass to C
+    self.assertEqual(test(outer.inner_ptr), 52)
+
+  def test_pointer_field_loses_reference(self):
+    # BUG: When a pointer is stored in a record struct field, only the address bytes are saved.
+    # The pointer's _objects dict (which prevents GC of the pointed-to object) is lost.
+    # This causes the pointed-to object to be garbage collected, leading to use-after-free.
+    from tinygrad.runtime.support.c import POINTER
+    @record
+    class MaskContext(c.Struct):
+      SIZE = 16
+      value = Field(ctypes.c_int, 0)
+      initialized = Field(ctypes.c_int, 4)
+      ptr = Field(ctypes.c_void_p, 8)
+    @record
+    class Params(c.Struct):
+      SIZE = 16
+      x = Field(ctypes.c_int, 0)
+      mask = Field(POINTER[MaskContext], 8)
+
+    src = """
+      struct mask_ctx { int value; int initialized; void *ptr; };
+      void mask_begin(struct mask_ctx *m, int val) { m->value = val; m->initialized = 1; }
+      int mask_end(struct mask_ctx *m) { return m->value + m->initialized; }
+    """
+    dll = self.compile(src)
+    @dll.bind(None, ctypes.POINTER(MaskContext), ctypes.c_int)
+    def mask_begin(m:POINTER[MaskContext], val:ctypes.c_int) -> None: ...
+    @dll.bind(ctypes.c_int, ctypes.POINTER(MaskContext))
+    def mask_end(m:POINTER[MaskContext]) -> ctypes.c_int: ...
+
+    # When MaskContext() is created inline, it gets garbage collected after the pointer
+    # is stored because only the address bytes are saved, not the _objects reference.
+    params = Params(x=1, mask=ctypes.pointer(MaskContext()))
+    mask_begin(params.mask, 42)
+    result = mask_end(params.mask)
+    self.assertEqual(result, 43)  # 42 + 1
+
+@unittest.skipIf(OSX and ('MTLCompiler' in DLL._loaded_ or 'llvm' in DLL._loaded_), "libclang can't be loaded after MTLCompiler or llvm on OSX")
+@unittest.skipIf(WIN, "doesn't compile on windows")
+class TestAutogen(unittest.TestCase):
+  def run_gen(self, contents):
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.h') as f:
+      f.write(contents)
+      f.flush()
+
+      generated_code = gen(name="test_header", dll=None, files=[f.name])
+
+      namespace = {}
+      exec(generated_code, namespace)
+      return namespace
+
+  def test_packed_structs(self):
+    ns = self.run_gen("""
+typedef unsigned NvU32;
+typedef unsigned long NvU64;
+
+typedef struct
+{
+    NvU32 version;
+    NvU32 size;
+    NvU64 gfwImageOffset;
+    NvU32 gfwImageSize;
+    NvU32 flags;
+} __attribute__((packed)) FWSECLIC_READ_VBIOS_DESC;
+
+#define FWSECLIC_READ_VBIOS_STRUCT_FLAGS (2)
+
+typedef struct
+{
+    NvU32 version;
+    NvU32 size;
+    NvU32 frtsRegionOffset4K;
+    NvU32 frtsRegionSize;
+    NvU32 frtsRegionMediaType;
+} __attribute__((packed)) FWSECLIC_FRTS_REGION_DESC;
+
+#define FWSECLIC_FRTS_REGION_MEDIA_FB (2)
+#define FWSECLIC_FRTS_REGION_SIZE_1MB_IN_4K (0x100)
+
+typedef struct
+{
+    FWSECLIC_READ_VBIOS_DESC readVbiosDesc;
+    FWSECLIC_FRTS_REGION_DESC frtsRegionDesc;
+} __attribute__((packed)) FWSECLIC_FRTS_CMD;
+""")
+
+    FWSECLIC_READ_VBIOS_DESC = ns['FWSECLIC_READ_VBIOS_DESC']
+    FWSECLIC_FRTS_REGION_DESC = ns['FWSECLIC_FRTS_REGION_DESC']
+    FWSECLIC_FRTS_CMD = ns['FWSECLIC_FRTS_CMD']
+
+    read_vbios_desc = FWSECLIC_READ_VBIOS_DESC(version=0x1, size=ctypes.sizeof(FWSECLIC_READ_VBIOS_DESC), flags=2)
+    frst_reg_desc = FWSECLIC_FRTS_REGION_DESC(version=0x1, size=ctypes.sizeof(FWSECLIC_FRTS_REGION_DESC),
+      frtsRegionOffset4K=0xdead, frtsRegionSize=0x100, frtsRegionMediaType=2)
+    frts_cmd = FWSECLIC_FRTS_CMD(readVbiosDesc=read_vbios_desc, frtsRegionDesc=frst_reg_desc)
+    assert int.from_bytes(frts_cmd, 'little') == 0x2000001000000dead0000001400000001000000020000000000000000000000000000001800000001
+    assert int.from_bytes(frts_cmd.readVbiosDesc, 'little') == int.from_bytes(read_vbios_desc, 'little')
+    assert int.from_bytes(frts_cmd.frtsRegionDesc, 'little') == int.from_bytes(frst_reg_desc, 'little')
+    assert frts_cmd.readVbiosDesc.__class__ is FWSECLIC_READ_VBIOS_DESC
+    assert frts_cmd.frtsRegionDesc.__class__ is FWSECLIC_FRTS_REGION_DESC
+
+  def test_gen_from_header(self):
+    namespace = self.run_gen("""
+    typedef struct {
+      int x;
+      int y;
+    } Point;
+
+    typedef enum {
+      RED = 0,
+      GREEN = 1,
+      BLUE = 2
+    } Color;
+
+    typedef struct {
+      Point origin;
+      int width;
+      int height;
+      Color color;
+    } Rectangle;
+
+    int add_points(Point a, Point b);""")
+
+    self.assertIn('Point', namespace)
+    self.assertIn('Color', namespace)
+    self.assertIn('Rectangle', namespace)
+    self.assertIn('RED', namespace)
+    self.assertIn('GREEN', namespace)
+    self.assertIn('BLUE', namespace)
+
+    self.assertEqual(namespace['RED'], 0)
+    self.assertEqual(namespace['GREEN'], 1)
+    self.assertEqual(namespace['BLUE'], 2)
+
+    Point = namespace['Point']
+    p = Point()
+    self.assertTrue(hasattr(p, 'x'))
+    self.assertTrue(hasattr(p, 'y'))
+
+    Rectangle = namespace['Rectangle']
+    rect = Rectangle()
+    self.assertTrue(hasattr(rect, 'origin'))
+    self.assertTrue(hasattr(rect, 'width'))
+    self.assertTrue(hasattr(rect, 'height'))
+    self.assertTrue(hasattr(rect, 'color'))
+
+    p2 = Point(10, 20)
+    self.assertEqual(p2.x, 10)
+    self.assertEqual(p2.y, 20)
+
+  def test_struct_ordering(self):
+    namespace = self.run_gen("""
+    struct A;
+    struct C;
+    typedef struct A A;
+
+    struct B {
+      struct C *c_ptr;
+    };
+
+    struct C {
+      struct A *a_ptr;
+    };
+
+    struct A {
+      int x;
+      struct B *b_ptr;
+    };""")
+
+    self.assertIn('struct_A', namespace)
+    self.assertIn('struct_B', namespace)
+    self.assertIn('struct_C', namespace)
+    A, B, C = namespace['A'], namespace['struct_B'], namespace['struct_C']
+    a, b, c = A(), B(), C()
+    self.assertTrue(hasattr(a, 'x'))
+    self.assertTrue(hasattr(a, 'b_ptr'))
+    self.assertTrue(hasattr(b, 'c_ptr'))
+    self.assertTrue(hasattr(c, 'a_ptr'))
+
+  def test_anonymous_children(self):
+    namespace = self.run_gen("""
+      struct foo {
+        struct {
+          int a,b;
+        } bar;
+      };
+    """)
+
+    self.assertIn('struct_foo', namespace)
+    self.assertIn('struct_foo_bar', namespace)
+
+  def test_enums(self):
+    namespace = self.run_gen("""
+      enum Foo { A, B, C };
+      enum Bar { X, Y, Z };
+    """)
+
+    assert namespace["A"] == 0
+    assert namespace["B"] == 1
+    assert namespace["C"] == 2
+    assert namespace["X"] == 0
+    assert namespace["Y"] == 1
+    assert namespace["Z"] == 2
+    assert namespace["enum_Foo"].get(0) == "A"
+    assert namespace["enum_Foo"].get(1) == "B"
+    assert namespace["enum_Foo"].get(2) == "C"
+    assert namespace["enum_Bar"].get(0) == "X"
+    assert namespace["enum_Bar"].get(1) == "Y"
+    assert namespace["enum_Bar"].get(2) == "Z"
+
+  @unittest.skipIf(OSX, "can't find stdint?")
+  def test_packed_fields(self):
+    ns = self.run_gen("""#include <stdint.h>
+typedef struct die_info
+ {
+	 uint16_t die_id;
+	 uint16_t die_offset; /* Points to the corresponding die_header structure */
+ } die_info;
+
+typedef struct ip_discovery_header
+ {
+	 uint32_t signature;    /* Table Signature */
+	 uint16_t version;      /* Table Version */
+	 uint16_t size;         /* Table Size */
+	 uint32_t id;           /* Table ID */
+	 uint16_t num_dies;     /* Number of Dies */
+	 die_info die_info[16]; /* list die information for up to 16 dies */
+	 union {
+		 uint16_t padding[1];	/* version <= 3 */
+		 struct {		/* version == 4 */
+			 uint8_t base_addr_64_bit : 1; /* ip structures are using 64 bit base address */
+			 uint8_t reserved : 7;
+			 uint8_t reserved2;
+		 };
+	 };
+ } ip_discovery_header;
+""")
+
+    ip_discovery_header = ns['ip_discovery_header']
+
+    hdr = b'IPDS\x04\x00|\x1d\x80\x1a\xffd\x01\x00\x00\x00\x8c\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00' # noqa: E501
+    ihdr = ip_discovery_header.from_buffer_copy(hdr)
+
+    assert ctypes.sizeof(ihdr) == 80
+    assert ihdr.signature == 0x53445049
+    assert ihdr.version == 0x0004
+    assert ihdr.num_dies == 1
+    assert ihdr.base_addr_64_bit == 1
+
+if __name__ == "__main__": unittest.main()

tinygrad_repo/test/null/test_compile_failures.py

@@ -0,0 +1,30 @@
+import unittest, io
+from contextlib import redirect_stdout
+from tinygrad import Tensor, dtypes, Device
+from tinygrad.helpers import OSX
+from tinygrad.engine.realize import compile_linear
+from tinygrad.codegen import to_program
+
+class TestCompileFailures(unittest.TestCase):
+  def compile(self, out:Tensor):
+    compile_linear(out.schedule_linear())
+
+  def test_interpolate_atari(self):
+    self.compile(Tensor.empty(210, 160, dtype='uint8').interpolate((64, 64)))
+
+  def test_add_max_uchar(self):
+    self.compile((Tensor.empty(1024, dtype='uint8') + Tensor.empty(1024, dtype='uint8')).max())
+
+class TestDisassembly(unittest.TestCase):
+  @unittest.skipUnless(Device.DEFAULT == "CPU" and OSX, "m series cpus support fp16 arithmetic")
+  def test_float16_alu(self):
+    c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
+    s = c.schedule_linear().src[-1]
+    p = to_program(s.src[0], Device[Device.DEFAULT].renderer)
+    lib = Device[Device.DEFAULT].compiler.compile(p.src[3].arg)
+    out = io.StringIO()
+    with redirect_stdout(out): Device[Device.DEFAULT].compiler.disassemble(lib)
+    assert "fcvt" not in out.getvalue()
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_const_folding.py

@@ -0,0 +1,150 @@
+import unittest, itertools, math
+from tinygrad import Tensor, dtypes, Context
+from tinygrad.dtype import DType, ConstType
+from tinygrad.uop.ops import Ops, UOp
+from test.helpers import full_rewrite
+import numpy as np
+
+def _check_ast_count(desired_count:int, t:Tensor):
+  # NOTE: this has side effect because everything can be scheduled only once
+  linear = t.schedule_linear()
+  asts = [s for s in linear.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 TestUnaryOpsConstFolding(unittest.TestCase):
+  def test_all_consts_ops(self):
+    _check_ast_count(0, Tensor.ones(4).exp())
+    _check_ast_count(0, Tensor.ones(4).sqrt())
+    _check_ast_count(0, Tensor.ones(4) + Tensor.ones(4))
+    _check_ast_count(0, Tensor.ones(4) / Tensor.ones(4))
+
+  def test_cast(self):
+    _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
+    _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))
+
+  def test_neg_folding(self):
+    _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
+    _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
+    _check_ast_count(0, Tensor([1, 2, 3]).neg().neg())
+
+  def test_neg_realized_no_fold(self):
+    x = Tensor.randn(32, 32)
+    x = x.clip(0, 1).realize()
+    _check_ast_count(1, x.neg())
+
+class TestBinaryOpsConstFolding(unittest.TestCase):
+  def test_add_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)
+  def test_add_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(4))
+  def test_literal_zero_add(self):
+    _check_ast_count(0, 0 + Tensor([1.0, 2, 3, 4]))
+  def test_tensor_zero_add(self):
+    _check_ast_count(0, Tensor.zeros(4) + Tensor([1.0, 2, 3, 4]))
+
+  def test_sub_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) - 0)
+  def test_sub_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) - Tensor.zeros(4))
+
+  def test_mul_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 0)
+  def test_mul_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.zeros(4))
+  def test_literal_zero_mul(self):
+    _check_ast_count(0, 0 * Tensor([1.0, 2, 3, 4]) * 0)
+  def test_tensor_zero_mul(self):
+    _check_ast_count(0, Tensor.zeros(4) * Tensor([1.0, 2, 3, 4]))
+
+  def test_mul_literal_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 1)
+  def test_mul_tensor_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(4))
+  def test_literal_one_mul(self):
+    _check_ast_count(0, 1 * Tensor([1.0, 2, 3, 4]))
+  def test_tensor_one_mul(self):
+    _check_ast_count(0, Tensor.ones(4) * Tensor([1.0, 2, 3, 4]))
+
+  def test_bool_tensor_mul_bool(self):
+    _check_ast_count(0, Tensor([True, False]) * True)
+    _check_ast_count(0, Tensor([True, False]) * False)
+  def test_bool_mul_bool_tensor(self):
+    _check_ast_count(0, True * Tensor([True, False]))
+    _check_ast_count(0, False * Tensor([True, False]))
+
+  def test_div_literal_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / 1)
+  def test_div_tensor_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / Tensor.ones(4))
+
+  def test_floordiv_literal_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
+  def test_floordiv_tensor_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) // Tensor.ones(4, dtype=dtypes.int32))
+
+  def test_pow_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 0)
+  def test_pow_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.zeros(4))
+
+  def test_pow_literal_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 1)
+  def test_pow_tensor_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.ones(4))
+  def test_literal_one_pow(self):
+    _check_ast_count(0, 1 ** Tensor([1.0, 2, 3, 4]))
+  def test_tensor_one_pow(self):
+    _check_ast_count(0, Tensor.ones(4) ** Tensor([1.0, 2, 3, 4]))
+
+class TestBitcastConstFolding(unittest.TestCase):
+  def test_scalar_bitcast(self):
+    def t(cases: dict[DType, ConstType]):
+      for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
+        if not math.isnan(from_v):
+          r = full_rewrite(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0]
+          self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
+          self.assertEqual(r.dtype, to_dt, msg)
+          np.testing.assert_equal(r.arg, to_v, msg)
+
+    t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False})
+    t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True})
+
+    t({dtypes.int8:  -1, dtypes.uint8:  2**8-1})
+    t({dtypes.int16: -1, dtypes.uint16: 2**16-1, dtypes.float16: float('nan')})
+    t({dtypes.int32: -1, dtypes.uint32: 2**32-1, dtypes.float32: float('nan')})
+    t({dtypes.int64: -1, dtypes.uint64: 2**64-1, dtypes.float64: float('nan')})
+
+    t({dtypes.int8:  -2**7,  dtypes.uint8:  2**7})
+    t({dtypes.int16: -2**15, dtypes.uint16: 2**15})
+    t({dtypes.int32: -2**31, dtypes.uint32: 2**31})
+    t({dtypes.int64: -2**63, dtypes.uint64: 2**63})
+
+    t({dtypes.int16: 13496, dtypes.uint16: 13496, dtypes.float16: 0.294921875})
+    t({dtypes.int32: 1050081145, dtypes.uint32: 1050081145, dtypes.float32: 0.29485681653022766})
+    t({dtypes.int64: 4598983288165178391, dtypes.uint64: 4598983288165178391, dtypes.float64: 0.29485681936461233})
+
+  def test_vec_bitcast(self):
+    with Context(SPEC=0):
+      srcs = full_rewrite(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src
+    self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs))
+    self.assertEqual(tuple(x.arg for x in srcs), (2**32-1, 2**31, 75))
+
+# folds advance indexing into basic indexing
+class TestIndexingConstFolding(unittest.TestCase):
+  def test_scalar_index(self):
+    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
+    _check_ast_count(1, t[:,:,Tensor(1),:])
+    _check_ast_count(1, t[:,:,Tensor(1)+2,:])
+    _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
+
+  def test_const_tensor_index(self):
+    # TODO: these can be 0, implement const tensor folded indexing
+    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
+    _check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
+    _check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
+    _check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_device.py

@@ -0,0 +1,192 @@
+#!/usr/bin/env python
+import unittest, os, subprocess
+from unittest.mock import patch
+from tinygrad import Tensor
+from tinygrad.device import Device, Compiler, enumerate_devices_str
+from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, Target, WIN, OSX, DEV
+from tinygrad.runtime.support.c import DLL
+
+class TestDevice(unittest.TestCase):
+  def test_canonicalize(self):
+    self.assertEqual(Device.canonicalize(None), Device.DEFAULT)
+    self.assertEqual(Device.canonicalize("CPU"), "CPU")
+    self.assertEqual(Device.canonicalize("cpu"), "CPU")
+    self.assertEqual(Device.canonicalize("CL"), "CL")
+    self.assertEqual(Device.canonicalize("CL:0"), "CL")
+    self.assertEqual(Device.canonicalize("cl:0"), "CL")
+    self.assertEqual(Device.canonicalize("CL:1"), "CL:1")
+    self.assertEqual(Device.canonicalize("cl:1"), "CL:1")
+    self.assertEqual(Device.canonicalize("CL:2"), "CL:2")
+    self.assertEqual(Device.canonicalize("disk:/dev/shm/test"), "DISK:/dev/shm/test")
+    self.assertEqual(Device.canonicalize("disk:000.txt"), "DISK:000.txt")
+
+  def test_getitem_not_exist(self):
+    with self.assertRaises(ModuleNotFoundError):
+      Device["TYPO"]
+
+  @unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")
+  def test_nonexistent_renderer(self):
+    with self.assertRaisesRegex(RuntimeError, "has no renderer"):
+      with Context(DEV="CPU:TYPO"): Device[Device.DEFAULT].renderer
+    with self.assertRaisesRegex(RuntimeError, "did you mean: 'CLANG'"):
+      with Context(DEV="CPU:CLANGJIT"): Device[Device.DEFAULT].renderer
+
+  @unittest.skipIf(Device.DEFAULT != "AMD", "only run on AMD")
+  def test_nonexistent_iface(self):
+    result = subprocess.run(['python3', '-c', 'from tinygrad import Device; Device[Device.DEFAULT].iface'],
+                            env={**os.environ, "DEV":"USA+AMD"}, capture_output=True)
+    self.assertNotEqual(result.returncode, 0)
+    self.assertIn(b"did you mean: 'USB'", result.stderr)
+
+  @unittest.skipIf(Device.DEFAULT != "AMD", "only run on AMD")
+  def test_dev_id_out_of_range(self):
+    result = subprocess.run(['python3', '-c', 'from tinygrad import Device; Device[Device.DEFAULT]'],
+                            env={**os.environ, "DEV":":99+AMD"}, capture_output=True)
+    self.assertNotEqual(result.returncode, 0)
+    self.assertIn(b"invalid visibility filter", result.stderr)
+
+  def test_lowercase_canonicalizes(self):
+    device = Device.DEFAULT
+    with Context(DEV=device.lower()):
+      self.assertEqual(Device.canonicalize(None), device)
+
+  def test_set_device_default_raises(self):
+    with self.assertRaisesRegex(AttributeError, "setting Device.DEFAULT is deprecated"):
+      Device.DEFAULT = "CPU"
+
+  def test_old_device_env_raises(self):
+    result = subprocess.run(['python3', '-c', 'from tinygrad import Device; Device.DEFAULT'],
+                            env={**os.environ, "CPU": "1", "DEV": ""}, capture_output=True)
+    self.assertNotEqual(result.returncode, 0)
+    self.assertIn(b"deprecated", result.stderr)
+
+  def test_old_renderer_env_raises(self):
+    result = subprocess.run(['python3', '-c', 'from tinygrad import Device; Device[Device.DEFAULT].renderer'],
+                            env={**os.environ, "DEV": "CPU", "CPU_LLVM": "1"}, capture_output=True)
+    self.assertNotEqual(result.returncode, 0)
+    self.assertIn(b"deprecated", result.stderr)
+
+  @unittest.skipIf(WIN, "skipping windows test") # TODO: subprocess causes memory violation?
+  def test_env_overwrite_default_compiler(self):
+    if Device.DEFAULT == "CPU":
+      from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangCompiler
+      try: _, _ = CPULLVMCompiler(), ClangCompiler()
+      except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
+
+      imports = "from tinygrad import Device; from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangCompiler"
+      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, CPULLVMCompiler)"'],
+                        shell=True, check=True, env={**os.environ, "DEV": "CPU:LLVM"})
+      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangCompiler)"'],
+                        shell=True, check=True, env={**os.environ, "DEV": "CPU"})
+      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangCompiler)"'],
+                        shell=True, check=True, env={**os.environ, "DEV": "CPU:CLANG"})
+    elif Device.DEFAULT == "AMD":
+      from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
+      try: _, _ = HIPCompiler(Device[Device.DEFAULT].arch), AMDLLVMCompiler(Device[Device.DEFAULT].arch)
+      except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
+
+      imports = "from tinygrad import Device; from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler"
+      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, AMDLLVMCompiler)"'],
+                        shell=True, check=True, env={**os.environ, "DEV": "AMD:LLVM"})
+      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, HIPCompiler)"'],
+                        shell=True, check=True, env={**os.environ, "DEV": "AMD"})
+      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, HIPCompiler)"'],
+                        shell=True, check=True, env={**os.environ, "DEV": "AMD:HIP"})
+    else: self.skipTest("only run on CPU/AMD")
+
+  @unittest.skipIf(WIN, "skipping windows test")
+  def test_env_online(self):
+    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangCompiler
+    try: _, _ = CPULLVMCompiler(), ClangCompiler()
+    except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
+
+    with Context(DEV="CPU:LLVM"):
+      inst = Device["CPU"].compiler
+      self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler)
+    with Context(DEV="CPU"):
+      self.assertIsInstance(Device["CPU"].compiler, ClangCompiler)
+    with Context(DEV="CPU:LLVM"):
+      self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler)
+      assert inst is Device["CPU"].compiler  # cached
+
+  @unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")
+  def test_compiler_autodetect_fallback(self):
+    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler
+
+    try: CPULLVMCompiler()
+    except Exception as e: self.skipTest(f"skipping: LLVM not available: {e}")
+
+    dev = Device["CPU"]
+    dev.cached_renderer.clear()
+    with patch("tinygrad.renderer.cstyle.ClangRenderer.__init__", side_effect=RuntimeError("broken")):
+      self.assertIsInstance(dev.renderer.compiler, CPULLVMCompiler)
+
+  def test_dev_contextvar(self):
+    orig_dev = Device.DEFAULT
+    with Context(DEV="CPU"): self.assertEqual(Tensor.empty(1).device, "CPU")
+    with Context(DEV="NULL"): self.assertEqual(Tensor.empty(1).device, "NULL")
+    self.assertEqual(Tensor.empty(1).device, orig_dev)
+
+class TestDevVar(unittest.TestCase):
+  def test_parse(self):
+    for d, t in [("AMD", Target(device="AMD", renderer="")), ("AMD:LLVM", Target(device="AMD", renderer="LLVM")),
+                 (":LLVM", Target(device="", renderer="LLVM")), ("AMD::gfx1100", Target(device="AMD", arch="gfx1100")),
+                 ("AMD:LLVM:gfx1100", Target(device="AMD", renderer="LLVM", arch="gfx1100")), ("::gfx1100", Target(arch="gfx1100")),
+                 ("CPU:LLVM:arm64,native,AMX", Target(device="CPU", renderer="LLVM", arch="arm64,native,AMX")),
+                 ("USB+", Target(interface="USB")), ("USB+AMD", Target(device="AMD", interface="USB")),
+                 ("PCI:0+AMD", Target(device="AMD", interface="PCI", indices="0")), (":0+AMD", Target(device="AMD", indices="0")),
+                 ("PCI:0,1+AMD", Target(device="AMD", interface="PCI", indices="0,1")),
+                 ("QCOM;USB+AMD", [Target(device="QCOM"), Target(device="AMD", interface="USB")])]:
+      with Context(DEV=d):
+        self.assertEqual(DEV.value, t if isinstance(t, list) else [t])
+        self.assertEqual(str(DEV), d)
+
+  def test_target(self):
+    with Context(DEV="CPU"): self.assertEqual(DEV.target("CPU"), Target("CPU"))
+    with Context(DEV="CPU:LLVM"): self.assertEqual(DEV.target("CPU"), Target("CPU", "LLVM"))
+    with Context(DEV=":LLVM"): self.assertEqual(DEV.target("CPU"), Target("CPU", "LLVM"))
+    with Context(DEV="AMD:LLVM"): self.assertEqual(DEV.target("CPU"), Target("CPU"))
+    with Context(DEV=""): self.assertEqual(DEV.target("CPU"), Target("CPU"))
+    with Context(DEV="QCOM:IR3;AMD:LLVM"):
+      self.assertEqual(DEV.target("QCOM"), Target("QCOM", "IR3"))
+      self.assertEqual(DEV.target("AMD"), Target("AMD", "LLVM"))
+      self.assertEqual(DEV.target("CPU"), Target("CPU"))
+
+  def test_dev_arch_override(self):
+    with Context(DEV="NULL::gfx1100"):
+      self.assertEqual(Device["NULL"].renderer.target.arch, "gfx1100")
+
+class MockCompiler(Compiler):
+  def __init__(self, key): super().__init__(key)
+  def compile(self, src) -> bytes: return src.encode()
+
+class TestCompiler(unittest.TestCase):
+  def test_compile_cached(self):
+    diskcache_put("key", "123", None) # clear cache
+    getenv.cache_clear()
+    with Context(CCACHE=1):
+      self.assertEqual(MockCompiler("key").compile_cached("123"), str.encode("123"))
+      self.assertEqual(diskcache_get("key", "123"), str.encode("123"))
+
+  def test_compile_cached_disabled(self):
+    diskcache_put("disabled_key", "123", None) # clear cache
+    getenv.cache_clear()
+    with Context(CCACHE=0):
+      self.assertEqual(MockCompiler("disabled_key").compile_cached("123"), str.encode("123"))
+      self.assertIsNone(diskcache_get("disabled_key", "123"))
+
+  def test_device_compile(self):
+    getenv.cache_clear()
+    with Context(CCACHE=0):
+      a = Tensor([0.,1.], device=Device.DEFAULT).realize()
+      (a + 1).realize()
+
+@unittest.skip("this test is broken if you have tinymesa installed")
+@unittest.skipIf(OSX and 'libclang' in DLL._loaded_, "MTLCompiler can't be loaded after libclang on OSX")
+class TestRunAsModule(unittest.TestCase):
+  def test_module_runs(self):
+    cpu_line = [l for l in enumerate_devices_str() if "CPU" in l][0]
+    self.assertIn("PASS", cpu_line, f"expected CPU to PASS, got: {cpu_line}")
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/null/test_disk_cache.py

@@ -0,0 +1,108 @@
+import unittest
+import pickle
+from tinygrad.helpers import diskcache_get, diskcache_put, diskcache, diskcache_clear
+
+def remote_get(table,q,k): q.put(diskcache_get(table, k))
+def remote_put(table,k,v): diskcache_put(table, k, v)
+
+class DiskCache(unittest.TestCase):
+  def test_putget(self):
+    table = "test_putget"
+    diskcache_put(table, "hello", "world")
+    self.assertEqual(diskcache_get(table, "hello"), "world")
+    diskcache_put(table, "hello", "world2")
+    self.assertEqual(diskcache_get(table, "hello"), "world2")
+
+  def test_putcomplex(self):
+    table = "test_putcomplex"
+    diskcache_put(table, "k", ("complex", 123, "object"))
+    ret = diskcache_get(table, "k")
+    self.assertEqual(ret, ("complex", 123, "object"))
+
+  def test_getotherprocess(self):
+    table = "test_getotherprocess"
+    from multiprocessing import Process, Queue
+    diskcache_put(table, "k", "getme")
+    q = Queue()
+    p = Process(target=remote_get, args=(table,q,"k"))
+    p.start()
+    p.join()
+    self.assertEqual(q.get(), "getme")
+
+  def test_putotherprocess(self):
+    table = "test_putotherprocess"
+    from multiprocessing import Process
+    p = Process(target=remote_put, args=(table,"k", "remote"))
+    p.start()
+    p.join()
+    self.assertEqual(diskcache_get(table, "k"), "remote")
+
+  def test_no_table(self):
+    self.assertIsNone(diskcache_get("faketable", "k"))
+
+  def test_ret(self):
+    table = "test_ret"
+    self.assertEqual(diskcache_put(table, "key", ("vvs",)), ("vvs",))
+
+  def test_non_str_key(self):
+    table = "test_non_str_key"
+    diskcache_put(table, 4, 5)
+    self.assertEqual(diskcache_get(table, 4), 5)
+    self.assertEqual(diskcache_get(table, "4"), 5)
+
+  def test_decorator(self):
+    calls = 0
+    @diskcache
+    def hello(x):
+      nonlocal calls
+      calls += 1
+      return "world"+x
+    self.assertEqual(hello("bob"), "worldbob")
+    self.assertEqual(hello("billy"), "worldbilly")
+    kcalls = calls
+    self.assertEqual(hello("bob"), "worldbob")
+    self.assertEqual(hello("billy"), "worldbilly")
+    self.assertEqual(kcalls, calls)
+
+  def test_dict_key(self):
+    table = "test_dict_key"
+    fancy_key = {"hello": "world", "goodbye": 7, "good": True, "pkl": pickle.dumps("cat")}
+    fancy_key2 = {"hello": "world", "goodbye": 8, "good": True, "pkl": pickle.dumps("cat")}
+    fancy_key3 = {"hello": "world", "goodbye": 8, "good": True, "pkl": pickle.dumps("dog")}
+    diskcache_put(table, fancy_key, 5)
+    self.assertEqual(diskcache_get(table, fancy_key), 5)
+    diskcache_put(table, fancy_key2, 8)
+    self.assertEqual(diskcache_get(table, fancy_key2), 8)
+    self.assertEqual(diskcache_get(table, fancy_key), 5)
+    self.assertEqual(diskcache_get(table, fancy_key3), None)
+
+  def test_table_name(self):
+    table = "test_gfx1010:xnack-"
+    diskcache_put(table, "key", "test")
+    self.assertEqual(diskcache_get(table, "key"), "test")
+
+  @unittest.skip("disabled by default because this drops cache table")
+  def test_clear_cache(self):
+    # clear cache to start
+    diskcache_clear()
+    tables = [f"test_clear_cache:{i}" for i in range(3)]
+    for table in tables:
+      # check no entries
+      self.assertIsNone(diskcache_get(table, "k"))
+    for table in tables:
+      diskcache_put(table, "k", "test")
+      # check insertion
+      self.assertEqual(diskcache_get(table, "k"), "test")
+
+    diskcache_clear()
+    for table in tables:
+      # check no entries again
+      self.assertIsNone(diskcache_get(table, "k"))
+
+    # calling multiple times is fine
+    diskcache_clear()
+    diskcache_clear()
+    diskcache_clear()
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/null/test_dtype.py

@@ -0,0 +1,135 @@
+import unittest, pickle
+from tinygrad.tensor import Tensor
+from tinygrad.dtype import dtypes, DType, ImageDType, PtrDType, to_dtype, Invalid, InvalidType
+
+class TestImageDType(unittest.TestCase):
+  def test_image_scalar(self):
+    assert dtypes.imagef((10,10)).base.scalar() == dtypes.float32
+    assert dtypes.imageh((10,10)).base.scalar() == dtypes.float32
+  def test_image_vec(self):
+    assert dtypes.imagef((10,10)).base.vec(4) == dtypes.float32.vec(4)
+    assert dtypes.imageh((10,10)).base.vec(4) == dtypes.float32.vec(4)
+
+class TestPtrDType(unittest.TestCase):
+  def test_vec_double(self):
+    dt1 = dtypes.float.vec(4).ptr().vec(4)
+    dt2 = dtypes.float.vec(4).ptr().vec(4)
+    self.assertEqual(dt1, dt2)
+    self.assertEqual(str(dt1), str(dt2))
+
+  def test_scalar(self):
+    dt = dtypes.float.vec(4).ptr().scalar()
+    self.assertEqual(dt.base, dtypes.float.vec(4))
+
+    dt = dtypes.float.vec(4).ptr().vec(4).scalar()
+    self.assertEqual(dt.base, dtypes.float.vec(4))
+
+    dt = dtypes.float.vec(4).scalar()
+    self.assertEqual(dt, dtypes.float)
+
+  def test_serialize(self):
+    dt = dtypes.float.vec(4).ptr().vec(4)
+    self.assertEqual(dt, eval(str(dt)))
+
+  def test_vec_ptr_sz(self):
+    dt = dtypes.float.ptr(1024).vec(4)
+    self.assertEqual(dt, eval(str(dt)))
+    self.assertEqual(str(dt), "dtypes.float.ptr(1024).vec(4)")
+
+  def test_vcount(self):
+    dt = dtypes.float.ptr().vec(4)
+    self.assertEqual(dt.vcount, 4)
+    self.assertEqual(dt.v, 4)
+    self.assertEqual(dt.count, 1)
+
+    dt = dtypes.float.vec(4).ptr()
+    self.assertEqual(dt.vcount, 1)
+    self.assertEqual(dt.v, 1)
+    self.assertEqual(dt.count, 4)
+
+    dt = dtypes.float.vec(4).ptr().vec(4)
+    self.assertEqual(dt.vcount, 4)
+    self.assertEqual(dt.v, 4)
+    self.assertEqual(dt.count, 4)
+
+class TestEqStrDType(unittest.TestCase):
+  def test_image_ne(self):
+    if ImageDType is None: raise unittest.SkipTest("no ImageDType support")
+    assert dtypes.float == dtypes.float32, "float doesn't match?"
+    assert dtypes.imagef((1,2,4)) != dtypes.imageh((1,2,4)), "different image dtype doesn't match"
+    assert dtypes.imageh((1,2,4)) != dtypes.imageh((1,4,2)), "different shape doesn't match"
+    assert dtypes.imageh((1,2,4)) == dtypes.imageh((1,2,4)), "same shape matches"
+    assert isinstance(dtypes.imageh((1,2,4)), ImageDType)
+  def test_ptr_eq(self):
+    assert dtypes.float32.ptr() == dtypes.float32.ptr()
+    assert not (dtypes.float32.ptr() != dtypes.float32.ptr())
+  def test_ptr_nbytes(self):
+    assert dtypes.float16.ptr(32).nbytes() == 32 * dtypes.float16.itemsize
+  def test_ptr_nbytes_unlimited(self):
+    self.assertRaises(RuntimeError, lambda: dtypes.float32.ptr().nbytes())
+  def test_strs(self):
+    if PtrDType is None: raise unittest.SkipTest("no PtrDType support")
+    self.assertEqual(str(dtypes.imagef((1,2,4))), "dtypes.imagef((1, 2, 4))")
+    self.assertEqual(str(dtypes.float32.ptr(16)), "dtypes.float.ptr(16)")
+
+class TestToDtype(unittest.TestCase):
+  def test_dtype_to_dtype(self):
+    dtype = dtypes.int32
+    res = to_dtype(dtype)
+    self.assertIsInstance(res, DType)
+    self.assertEqual(res, dtypes.int32)
+
+  def test_str_to_dtype(self):
+    dtype = "int32"
+    res = to_dtype(dtype)
+    self.assertIsInstance(res, DType)
+    self.assertEqual(res, dtypes.int32)
+
+class TestCastConvenienceMethod(unittest.TestCase):
+  def test_method(self):
+    for input_dtype in (dtypes.float, dtypes.int):
+      t = Tensor([1, 2], dtype=input_dtype)
+      self.assertEqual(t.dtype, input_dtype)
+      self.assertEqual(t.bool().dtype, dtypes.bool)
+      self.assertEqual(t.short().dtype, dtypes.short)
+      self.assertEqual(t.int().dtype, dtypes.int)
+      self.assertEqual(t.long().dtype, dtypes.long)
+      self.assertEqual(t.half().dtype, dtypes.half)
+      self.assertEqual(t.bfloat16().dtype, dtypes.bfloat16)
+      self.assertEqual(t.float().dtype, dtypes.float)
+      self.assertEqual(t.double().dtype, dtypes.double)
+
+class TestDtypeTolist(unittest.TestCase):
+  def test_bfloat16(self):
+    self.assertEqual(Tensor([-60000, 1.5, 3.1, 60000], device="PYTHON", dtype=dtypes.bfloat16).tolist(), [-59904.0, 1.5, 3.09375, 59904.0])
+  def test_fp8(self):
+    # 448
+    self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e4m3).tolist(), [-448.0, 1.5, 3.0, 448.0])
+    # 57344
+    self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e5m2).tolist(), [-28672.0, 1.5, 3.0, 28672.0])
+
+class TestCanLosslessCast(unittest.TestCase):
+  def test_can_lossless_cast(self):
+    from tinygrad.dtype import can_lossless_cast
+    # signed -> unsigned is NOT lossless (negative values wrap)
+    self.assertFalse(can_lossless_cast(dtypes.int8, dtypes.uint64))
+    self.assertFalse(can_lossless_cast(dtypes.int32, dtypes.uint32))
+    # unsigned -> larger signed is lossless
+    self.assertTrue(can_lossless_cast(dtypes.uint8, dtypes.int16))
+    self.assertTrue(can_lossless_cast(dtypes.uint32, dtypes.int64))
+    # large ints don't fit in floats
+    self.assertFalse(can_lossless_cast(dtypes.int32, dtypes.float))
+    self.assertFalse(can_lossless_cast(dtypes.int64, dtypes.double))
+    # half has more mantissa bits
+    self.assertTrue(can_lossless_cast(dtypes.int8, dtypes.half))
+    self.assertFalse(can_lossless_cast(dtypes.int8, dtypes.bfloat16))
+
+class TestInvalidSingleton(unittest.TestCase):
+  def test_singleton(self):
+    self.assertIs(InvalidType(), InvalidType())
+    self.assertIs(InvalidType(), Invalid)
+  def test_pickle(self):
+    self.assertIs(pickle.loads(pickle.dumps(Invalid)), Invalid)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/null/test_dtype_spec.py

@@ -0,0 +1,440 @@
+import unittest, math, struct, operator
+from tinygrad import Tensor, Device
+from tinygrad.dtype import DTYPES_DICT, dtypes, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
+
+from tinygrad.helpers import getenv
+from hypothesis import given, settings, strategies as strat
+import numpy as np
+import torch
+
+settings.register_profile("my_profile", max_examples=50, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+settings.load_profile("my_profile")
+
+core_dtypes = list(DTYPES_DICT.values())
+dtype_ints = [dt for dt in core_dtypes if dtypes.is_int(dt) and dt in Device[Device.DEFAULT].renderer.supported_dtypes()]
+dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and dt in Device[Device.DEFAULT].renderer.supported_dtypes()]
+
+FP8E4M3_MAX = 448.0
+FP8E5M2_MAX = 57344.0
+
+def u32_to_f32(u): return struct.unpack('f', struct.pack('I', u))[0]
+def f32_to_u32(f): return struct.unpack('I', struct.pack('f', f))[0]
+
+class TestHelpers(unittest.TestCase):
+  signed_ints = (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64)
+  uints = (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
+  floats = (dtypes.float16, dtypes.float32, dtypes.float64)
+
+  @given(strat.sampled_from(signed_ints+uints), strat.integers(min_value=1, max_value=8))
+  def test_is_int(self, dtype, amt):
+    assert dtypes.is_int(dtype.vec(amt) if amt > 1 else dtype)
+    assert not dtypes.is_float(dtype.vec(amt) if amt > 1 else dtype)
+
+  @given(strat.sampled_from(uints), strat.integers(min_value=1, max_value=8))
+  def test_is_unsigned_uints(self, dtype, amt):
+    assert dtypes.is_unsigned(dtype.vec(amt) if amt > 1 else dtype)
+
+  @given(strat.sampled_from(signed_ints), strat.integers(min_value=1, max_value=8))
+  def test_is_unsigned_signed_ints(self, dtype, amt):
+    assert not dtypes.is_unsigned(dtype.vec(amt) if amt > 1 else dtype)
+
+  @given(strat.sampled_from(floats), strat.integers(min_value=1, max_value=8))
+  def test_is_float(self, dtype, amt):
+    assert dtypes.is_float(dtype.vec(amt) if amt > 1 else dtype)
+    assert not dtypes.is_int(dtype.vec(amt) if amt > 1 else dtype)
+    assert not dtypes.is_unsigned(dtype.vec(amt) if amt > 1 else dtype)
+
+  def test_bf16_is_float(self):
+    assert dtypes.is_float(dtypes.bfloat16)
+
+  def test_fp8s_are_float(self):
+    assert dtypes.is_float(dtypes.fp8e4m3)
+    assert dtypes.is_float(dtypes.fp8e5m2)
+
+  @given(strat.sampled_from([d for d in DTYPES_DICT.values() if dtypes.is_float(d) or dtypes.is_int(d)]), strat.integers(min_value=2, max_value=8))
+  def test_scalar(self, dtype, amt):
+    assert dtype.vec(amt).scalar() == dtype
+
+  def test_from_py(self):
+    assert dtypes.from_py(True) == dtypes.bool
+    assert dtypes.from_py(2) == dtypes.default_int
+    assert dtypes.from_py(3.0) == dtypes.default_float
+    assert dtypes.from_py([]) == dtypes.default_float
+    assert dtypes.from_py(()) == dtypes.default_float
+    assert dtypes.from_py([True]) == dtypes.bool
+    assert dtypes.from_py([True, 2]) == dtypes.default_int
+    assert dtypes.from_py([True, 3.0]) == dtypes.default_float
+    assert dtypes.from_py([2, 3.0]) == dtypes.default_float
+    assert dtypes.from_py([True, 2, 3.0]) == dtypes.default_float
+    with self.assertRaises(RuntimeError): dtypes.from_py(None)
+    with self.assertRaises(RuntimeError): dtypes.from_py([None])
+    with self.assertRaises(RuntimeError): dtypes.from_py({})
+    with self.assertRaises(RuntimeError): dtypes.from_py(set())
+
+  def test_dtype_range(self):
+    for dt in core_dtypes:
+      if dtypes.is_float(dt):
+        np.testing.assert_equal(dt.min, -math.inf)
+        np.testing.assert_equal(dt.max, math.inf)
+        np.testing.assert_equal(dt.min, -math.inf)
+        np.testing.assert_equal(dt.max, math.inf)
+      elif dtypes.is_int(dt):
+        info = np.iinfo(_to_np_dtype(dt))
+        np.testing.assert_equal(dt.min, info.min)
+        np.testing.assert_equal(dt.max, info.max)
+        np.testing.assert_equal(dt.min, info.min)
+        np.testing.assert_equal(dt.max, info.max)
+      else:
+        assert dt == dtypes.bool, dt
+        np.testing.assert_equal(dt.min, False)
+        np.testing.assert_equal(dt.max, True)
+        np.testing.assert_equal(dt.min, False)
+        np.testing.assert_equal(dt.max, True)
+
+  def test_dtype_range_vec(self):
+    for dt in core_dtypes:
+      self.assertEqual(dt.min, dt.vec(4).min)
+      self.assertEqual(dt.max, dt.vec(4).max)
+
+  def test_float_to_fp16(self):
+    self.assertEqual(float_to_fp16(1), 1)
+    self.assertEqual(float_to_fp16(65504), 65504)
+    self.assertEqual(float_to_fp16(65519.999), 65504)
+    self.assertEqual(float_to_fp16(65520), math.inf)
+    self.assertEqual(float_to_fp16(1e-8), 0.0)
+    self.assertEqual(float_to_fp16(-65504), -65504)
+    self.assertEqual(float_to_fp16(-65519.999), -65504)
+    self.assertEqual(float_to_fp16(-65520), -math.inf)
+    self.assertTrue(math.isnan(float_to_fp16(math.nan)))
+
+  def test_float_to_bf16(self):
+    max_bf16 = torch.finfo(torch.bfloat16).max
+    for a in [1, 1.1, 1234, 23456, -777.777, max_bf16, max_bf16 * 1.00001, -max_bf16, -max_bf16 * 1.00001, math.inf, -math.inf]:
+      self.assertEqual(float_to_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
+    self.assertTrue(math.isnan(float_to_bf16(math.nan)))
+
+  def test_float_to_bf16_nan(self):
+    patterns = [0x7FC00001, 0xFFC00001, 0x7F800001, 0xFF800001, 0x7FFFFFFF, 0xFFFFFFFF]
+    for u in patterns:
+      x = u32_to_f32(u)
+      y = float_to_bf16(x)
+      t = torch.tensor([x], dtype=torch.bfloat16).item()
+      self.assertTrue(math.isnan(y))
+      self.assertTrue(math.isnan(t))
+
+  def test_float_to_bf16_round(self):
+    uppers = [0x3f800000, 0x41230000, 0xC1460000]
+    for upper in uppers:
+      base = upper & 0xFFFF0000
+      base_f32 = u32_to_f32(base)
+      base_f32_round_up = u32_to_f32(base + 0x00010000)
+
+      x = u32_to_f32(base | 0x00007000)
+      self.assertEqual(float_to_bf16(x), base_f32)
+      self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32)
+
+      x = u32_to_f32(base | 0x0000C000)
+      self.assertEqual(float_to_bf16(x), base_f32_round_up)
+      self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32_round_up)
+
+      if ((upper >> 16) & 1) == 0:
+        x = u32_to_f32(base | 0x00008000)
+        self.assertEqual(float_to_bf16(x), base_f32)
+        self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32)
+      else:
+        x = u32_to_f32(base | 0x00008000)
+        self.assertEqual(float_to_bf16(x), base_f32_round_up)
+        self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32_round_up)
+
+  def test_float_to_bf16_boundary(self):
+    base = 0x7F7F0000
+    inf_u32 = 0x7F800000
+
+    x = u32_to_f32(base | 0x00007FFF)
+    self.assertEqual(f32_to_u32(float_to_bf16(x)), base)
+    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), base)
+
+    x = u32_to_f32(base | 0x0000C000)
+    self.assertEqual(f32_to_u32(float_to_bf16(x)), inf_u32)
+    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), inf_u32)
+
+    x = u32_to_f32(base | 0x00008000)
+    self.assertEqual(f32_to_u32(float_to_bf16(x)), inf_u32)
+    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), inf_u32)
+
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
+  def test_truncate_fp8e4m3(self, x):
+    if math.isnan(x): np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), x)
+    elif math.isinf(x): np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), math.copysign(math.nan, x))
+    elif x > FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), FP8E4M3_MAX)
+    elif x < -FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), -FP8E4M3_MAX)
+    else: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), torch.tensor(x, dtype=torch.float8_e4m3fn).float().item())
+
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
+  def test_truncate_fp8e5m2(self, x):
+    if math.isnan(x): np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), x)
+    elif math.isinf(x): np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), x)
+    elif x > FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), FP8E5M2_MAX)
+    elif x < -FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), -FP8E5M2_MAX)
+    else: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), torch.tensor(x, dtype=torch.float8_e5m2).float().item())
+
+  def test_finfo(self):
+    for dt in [dtypes.float16, dtypes.float32, dtypes.float64]:
+      info = np.finfo(_to_np_dtype(dt))
+      self.assertEqual(info.bits, dt.bitsize)
+      self.assertEqual((info.nexp, info.nmant), dtypes.finfo(dt))
+
+class TestTypePromotion(unittest.TestCase):
+  @given(strat.sampled_from(core_dtypes))
+  def test_self_promo_to_self(self, dtype):
+    assert least_upper_dtype(dtype) == dtype
+    assert least_upper_dtype(dtype, dtype) == dtype
+    assert least_upper_dtype(dtype, dtype, dtype) == dtype
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_promo_resulted_higher_than_inputs(self, dtype1, dtype2):
+    result = least_upper_dtype(dtype1, dtype2)
+    assert not (result < dtype1) and not (result < dtype2)
+
+  def test_dtype_promo(self):
+    assert least_upper_dtype(dtypes.bool, dtypes.int8) == dtypes.int8
+    assert least_upper_dtype(dtypes.int8, dtypes.uint8) == dtypes.int16
+    assert least_upper_dtype(dtypes.uint8, dtypes.int16) == dtypes.int16
+    assert least_upper_dtype(dtypes.int16, dtypes.uint16) == dtypes.int32
+    assert least_upper_dtype(dtypes.uint16, dtypes.int32) == dtypes.int32
+    assert least_upper_dtype(dtypes.int32, dtypes.uint32) == dtypes.int64
+    assert least_upper_dtype(dtypes.uint32, dtypes.int64) == dtypes.int64
+    assert least_upper_dtype(dtypes.int64, dtypes.uint64) == dtypes.uint64
+    assert least_upper_dtype(dtypes.float16, dtypes.float32) == dtypes.float32
+    assert least_upper_dtype(dtypes.float32, dtypes.float64) == dtypes.float64
+
+    assert least_upper_dtype(dtypes.bool, dtypes.float32) == dtypes.float32
+    assert least_upper_dtype(dtypes.bool, dtypes.float64) == dtypes.float64
+    assert least_upper_dtype(dtypes.float16, dtypes.int64) == dtypes.float16
+    assert least_upper_dtype(dtypes.float16, dtypes.uint64) == dtypes.float16
+    assert least_upper_dtype(dtypes.fp8e4m3, dtypes.fp8e5m2) == dtypes.half
+    assert least_upper_dtype(dtypes.fp8e4m3, dtypes.bfloat16) == dtypes.bfloat16
+    assert least_upper_dtype(dtypes.fp8e5m2, dtypes.bfloat16) == dtypes.bfloat16
+    assert least_upper_dtype(dtypes.fp8e4m3, dtypes.float16) == dtypes.float16
+    assert least_upper_dtype(dtypes.fp8e5m2, dtypes.float16) == dtypes.float16
+    assert least_upper_dtype(dtypes.fp8e4m3, dtypes.int64) == dtypes.fp8e4m3
+    assert least_upper_dtype(dtypes.fp8e4m3, dtypes.uint64) == dtypes.fp8e4m3
+    assert least_upper_dtype(dtypes.fp8e5m2, dtypes.int64) == dtypes.fp8e5m2
+    assert least_upper_dtype(dtypes.fp8e5m2, dtypes.uint64) == dtypes.fp8e5m2
+
+  def test_weakint_promo(self):
+    # weakint with itself is weakint
+    assert least_upper_dtype(dtypes.weakint, dtypes.weakint) == dtypes.weakint
+    # weakint is above bool
+    assert least_upper_dtype(dtypes.weakint, dtypes.bool) == dtypes.weakint
+    # weakint defers to any concrete int type
+    assert least_upper_dtype(dtypes.weakint, dtypes.int8) == dtypes.int8
+    assert least_upper_dtype(dtypes.weakint, dtypes.uint8) == dtypes.uint8
+    assert least_upper_dtype(dtypes.weakint, dtypes.int16) == dtypes.int16
+    assert least_upper_dtype(dtypes.weakint, dtypes.int32) == dtypes.int32
+    assert least_upper_dtype(dtypes.weakint, dtypes.int64) == dtypes.int64
+    assert least_upper_dtype(dtypes.weakint, dtypes.uint64) == dtypes.uint64
+    # weakint defers to any float type
+    assert least_upper_dtype(dtypes.weakint, dtypes.float16) == dtypes.float16
+    assert least_upper_dtype(dtypes.weakint, dtypes.float32) == dtypes.float32
+    assert least_upper_dtype(dtypes.weakint, dtypes.float64) == dtypes.float64
+
+class TestTypeSpec(unittest.TestCase):
+  def setUp(self):
+    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
+  def tearDown(self):
+    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
+
+  def test_set_dtype_default(self):
+    for default_int in [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64]:
+      dtypes.default_int = default_int
+      assert dtypes.default_int == default_int
+
+    for default_float in [*dtypes.fp8s, dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
+      dtypes.default_float = default_float
+      assert dtypes.default_float == default_float
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from([operator.gt, operator.ge, operator.le, operator.lt, operator.eq, operator.ne]))
+  def test_bool_ops(self, dtype, op):
+    assert op(Tensor.ones(4, 4, dtype=dtype), Tensor.ones(4, 4, dtype=dtype)).dtype == dtypes.bool
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
+  def test_functions_return_index(self, dtype, default_int, default_float):
+    dtypes.default_int, dtypes.default_float = default_int, default_float
+    assert Tensor([0, 1], dtype=dtype).argmax().dtype == dtypes.int32
+    assert Tensor([0, 1], dtype=dtype).argmin().dtype == dtypes.int32
+    assert Tensor([0, 1], dtype=dtype).multinomial().dtype == dtypes.int32
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints))
+  def test_tensor_indexing_returns_same_dtype(self, data_dtype, indices_dtype):
+    X_data =  Tensor.ones(60000, 1, 28, 28, dtype=data_dtype)
+    indices =  Tensor.randint(512, high=X_data.shape[0]).cast(indices_dtype)
+    assert X_data[indices].dtype == X_data.dtype
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints))
+  def test_gather_returns_same_dtype(self, data_dtype, indices_dtype):
+    X_data = Tensor([[1, 0], [0, 1]], dtype=data_dtype)
+    indices = Tensor([[0, 0], [1, 0]], dtype=indices_dtype)
+    assert X_data.gather(0, indices).dtype == X_data.dtype
+    assert X_data.gather(1, indices).dtype == X_data.dtype
+
+  @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
+  def test_attention_returns_same_dtype(self, data_dtype, default_float):
+    dtypes.default_float = default_float
+    query = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
+    key = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
+    value = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
+    mask = (Tensor.rand(32, 8, 128, 128) < 0.5)
+    assert query.scaled_dot_product_attention(key, value, is_causal=True).dtype == data_dtype
+    assert query.scaled_dot_product_attention(key, value, is_causal=True, dropout_p=0.3).dtype == data_dtype
+    assert query.scaled_dot_product_attention(key, value, is_causal=False).dtype == data_dtype
+    assert query.scaled_dot_product_attention(key, value, attn_mask=mask).dtype == data_dtype
+
+class TestAutoCastType(unittest.TestCase):
+  def setUp(self):
+    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
+  def tearDown(self):
+    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
+
+  @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
+  def test_least_upper_float_input_is_float(self, input_dtype, default_float):
+    dtypes.default_float = default_float
+    self.assertEqual(least_upper_float(input_dtype), input_dtype)
+
+  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
+  def test_least_upper_float_input_is_int(self, input_dtype, default_float):
+    dtypes.default_float = default_float
+    self.assertEqual(least_upper_float(input_dtype), default_float)
+
+  @given(strat.sampled_from(core_dtypes))
+  def test_broadcast_scalar(self, dt):
+    assert (Tensor.ones(4, 4, dtype=dt) + 2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
+    assert (Tensor.ones(4, 4, dtype=dt) + 2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
+    assert (Tensor.ones(4, 4, dtype=dt) + True).dtype == dt
+
+  @given(strat.sampled_from(dtype_floats))
+  def test_int_div_int(self, default_float):
+    dtypes.default_float = default_float
+    self.assertEqual(Tensor([1]).div(Tensor([2])).dtype, default_float)
+
+  def test_sum(self):
+    assert (Tensor([0, 1], dtype=dtypes.bool)).sum().dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int8)).sum().dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int16)).sum().dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int32)).sum().dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int64)).sum().dtype == dtypes.int64
+    assert (Tensor([0, 1], dtype=dtypes.uint8)).sum().dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint16)).sum().dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint32)).sum().dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint64)).sum().dtype == dtypes.uint64
+    assert (Tensor([0, 1], dtype=dtypes.fp8e4m3)).sum().dtype == dtypes.fp8e4m3
+    assert (Tensor([0, 1], dtype=dtypes.fp8e5m2)).sum().dtype == dtypes.fp8e5m2
+    assert (Tensor([0, 1], dtype=dtypes.float16)).sum().dtype == dtypes.float16
+    assert (Tensor([0, 1], dtype=dtypes.bfloat16)).sum().dtype == dtypes.bfloat16
+    assert (Tensor([0, 1], dtype=dtypes.float32)).sum().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.float64)).sum().dtype == dtypes.float64
+
+  def test_mean(self):
+    assert (Tensor([0, 1], dtype=dtypes.bool)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.int8)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.int16)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.int32)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.int64)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.uint8)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.uint16)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.uint32)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.uint64)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.fp8e4m3)).mean().dtype == dtypes.fp8e4m3
+    assert (Tensor([0, 1], dtype=dtypes.fp8e5m2)).mean().dtype == dtypes.fp8e5m2
+    assert (Tensor([0, 1], dtype=dtypes.float16)).mean().dtype == dtypes.float16
+    assert (Tensor([0, 1], dtype=dtypes.bfloat16)).mean().dtype == dtypes.bfloat16
+    assert (Tensor([0, 1], dtype=dtypes.float32)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.float64)).mean().dtype == dtypes.float64
+
+  def test_cumsum(self):
+    assert (Tensor([0, 1], dtype=dtypes.bool)).cumsum(0).dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int8)).cumsum(0).dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int16)).cumsum(0).dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int32)).cumsum(0).dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int64)).cumsum(0).dtype == dtypes.int64
+    assert (Tensor([0, 1], dtype=dtypes.uint8)).cumsum(0).dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint16)).cumsum(0).dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint32)).cumsum(0).dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint64)).cumsum(0).dtype == dtypes.uint64
+    assert (Tensor([0, 1], dtype=dtypes.fp8e4m3)).cumsum(0).dtype == dtypes.fp8e4m3
+    assert (Tensor([0, 1], dtype=dtypes.fp8e5m2)).cumsum(0).dtype == dtypes.fp8e5m2
+    assert (Tensor([0, 1], dtype=dtypes.float16)).cumsum(0).dtype == dtypes.float16
+    assert (Tensor([0, 1], dtype=dtypes.bfloat16)).cumsum(0).dtype == dtypes.bfloat16
+    assert (Tensor([0, 1], dtype=dtypes.float32)).cumsum(0).dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.float64)).cumsum(0).dtype == dtypes.float64
+
+  def test_cumsum_empty(self):
+    # empty cumsum dtype must match non-empty
+    for d in (dtypes.bool, dtypes.int8, dtypes.uint8, dtypes.float16, dtypes.float32):
+      self.assertEqual(Tensor([], dtype=d).cumsum(0).dtype, Tensor([0, 1], dtype=d).cumsum(0).dtype)
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_matmul(self, dt1, dt2, acc_dt):
+    t1 = Tensor([0, 1], dtype=dt1)
+    t2 = Tensor([0, 1], dtype=dt2)
+    self.assertEqual(t1.matmul(t2).dtype, least_upper_dtype(t1.dtype, t2.dtype))
+    self.assertEqual(t1.matmul(t2, dtype=acc_dt).dtype, acc_dt)
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_linear(self, dt1, dt2, dt3, acc_dt):
+    x = Tensor([0, 1], dtype=dt1)
+    w = Tensor([0, 1], dtype=dt2)
+    b = Tensor([0, 1], dtype=dt3)
+    self.assertEqual(x.linear(w).dtype, least_upper_dtype(x.dtype, w.dtype))
+    self.assertEqual(x.linear(w, b).dtype, least_upper_dtype(least_upper_dtype(x.dtype, w.dtype), b.dtype))
+    self.assertEqual(x.linear(w, dtype=acc_dt).dtype, acc_dt)
+    self.assertEqual(x.linear(w, b, dtype=acc_dt).dtype, acc_dt)
+
+  @staticmethod
+  def check_where_alternate_input_other(input_, other, data_type):
+    assert (Tensor([True, False]).where(input_, other)).dtype == data_type
+    assert (Tensor([True, False]).where(other, input_)).dtype == data_type
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_where_no_scalar(self, dt1, dt2):
+    self.check_where_alternate_input_other(Tensor(2, dtype=dt1), Tensor(3, dtype=dt2), least_upper_dtype(dt1, dt2))
+
+  @given(strat.sampled_from(core_dtypes))
+  def test_where_one_scalar(self, dt):
+    t = Tensor(2, dtype=dt)
+    self.check_where_alternate_input_other(t, 3.2, (dt if dtypes.is_float(dt) else dtypes.default_float))
+    self.check_where_alternate_input_other(t, 3, (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int))
+    self.check_where_alternate_input_other(t, True, dt)
+
+  def test_where_two_scalars(self):
+    self.check_where_alternate_input_other(3.1, 3.2, dtypes.default_float)
+    self.check_where_alternate_input_other(3.1, 3, dtypes.default_float)
+    self.check_where_alternate_input_other(3.1, True, dtypes.default_float)
+    self.check_where_alternate_input_other(3, 2, dtypes.default_int)
+    self.check_where_alternate_input_other(3, True, dtypes.default_int)
+    self.check_where_alternate_input_other(False, True, dtypes.bool)
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_maximum(self, dt1, dt2):
+    assert Tensor([0, 1, 2], dtype=dt1).maximum(Tensor([2, 0, 5], dtype=dt2)).dtype == least_upper_dtype(dt1, dt2)
+
+  @given(strat.sampled_from(core_dtypes))
+  def test_maximum_const(self, dt):
+    assert Tensor([1, 2], dtype=dt).maximum(3.1).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
+    assert Tensor([1, 2], dtype=dt).maximum(3).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
+    assert Tensor([1, 2], dtype=dt).maximum(True).dtype == dt
+
+  def test_div(self):
+    assert (Tensor([1, 2], dtype=dtypes.int32) / Tensor([2, 2], dtype=dtypes.int32)).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.int16) / Tensor([2, 2], dtype=dtypes.int32)).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.float32) / Tensor([2, 2], dtype=dtypes.float16)).dtype == dtypes.float32
+    assert (Tensor([1, 2], dtype=dtypes.int32) / Tensor([2, 2], dtype=dtypes.float16)).dtype == dtypes.float16
+
+  def test_div_const(self):
+    assert (Tensor([1, 2], dtype=dtypes.int32) / 2).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.int32) / 2.0).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.float16) / 2).dtype == dtypes.float16
+    assert (Tensor([1, 2], dtype=dtypes.float16) / 2.0).dtype == dtypes.float16
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_elf.py

@@ -0,0 +1,38 @@
+import unittest, subprocess, platform
+from tinygrad.runtime.support.compiler_cpu import ClangCompiler
+from tinygrad.runtime.support.elf import elf_loader
+
+class TestElfLoader(unittest.TestCase):
+  def test_load_clang_jit_strtab(self):
+    src = '''
+      int something; // will be a load from a relocation (needed for .rela.text to exist)
+      int test(int x) {
+        return something + x;
+      }
+    '''
+    args = ('-x', 'c', '-c', '-target', f'{platform.machine()}-none-unknown-elf', '-march=native', '-fPIC', '-O2', '-ffreestanding', '-nostdlib')
+    obj = subprocess.check_output(('clang',) + args + ('-', '-o', '-'), input=src.encode('utf-8'))
+    _, sections, _ = elf_loader(obj)
+    section_names = [sh.name for sh in sections]
+    assert '.text' in section_names and '.rela.text' in section_names, str(section_names)
+  def test_clang_jit_compiler_external_raise(self):
+    src = '''
+      int evil_external_function(int);
+      int test(int x) {
+        return evil_external_function(x+2)*2;
+      }
+    '''
+    with self.assertRaisesRegex(RuntimeError, 'evil_external_function'):
+      ClangCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src)
+  def test_link(self):
+    src = '''
+      float powf(float, float); // from libm
+      float test(float x, float y) { return powf(x, y); }
+    '''
+    args = ('-x', 'c', '-c', '-target', f'{platform.machine()}-none-unknown-elf', '-march=native', '-fPIC', '-O2', '-ffreestanding', '-nostdlib')
+    obj = subprocess.check_output(('clang',) + args + ('-', '-o', '-'), input=src.encode())
+    with self.assertRaisesRegex(RuntimeError, 'powf'): elf_loader(obj)
+    elf_loader(obj, link_libs=['m'])
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_gc.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python
+import gc, inspect
+import unittest
+import numpy as np
+from tinygrad.device import Buffer
+from tinygrad.engine.realize import run_linear
+from tinygrad.uop.ops import UOp
+from tinygrad.tensor import Tensor
+
+def _allocations_of_type(t):
+  ret = 0
+  for x in gc.get_objects():
+    try:
+      if isinstance(x, t): ret += 1
+    except ReferenceError:
+      pass
+  return ret
+
+def tensors_allocated():
+  gc.collect()
+  return _allocations_of_type(Tensor)
+
+def bufs_allocated():
+  gc.collect()
+  return _allocations_of_type(Buffer)
+
+class TestGC(unittest.TestCase):
+
+  def test_gc(self):
+    Tensor.manual_seed(0)
+    base = tensors_allocated()
+    a = Tensor.rand(4, 4)
+    b = Tensor.zeros(4, 4)
+    (a*b).mean().backward()
+    assert (tensors_allocated()-base > 0)
+    del a,b
+    assert (tensors_allocated()-base == 2) # one for Tensor._device_rng_counters, and one for Tensor._device_seeds
+    Tensor.manual_seed(0)
+
+  def test_gc_complex(self):
+    Tensor.manual_seed(0)
+    base = tensors_allocated()
+    a = Tensor(np.zeros((4, 4), dtype=np.float32))
+    b = Tensor.rand(4, 4)
+    assert (tensors_allocated()-base == 4)
+    (a*b).mean().backward()
+    assert (tensors_allocated()-base == 6)
+    del b
+    assert (tensors_allocated()-base == 4)
+    b = Tensor(np.zeros((4, 4), dtype=np.float32))
+    print(tensors_allocated())
+    (a*b).mean().backward()
+    print(tensors_allocated())
+    assert (tensors_allocated()-base == 6)
+    del b
+    assert (tensors_allocated()-base == 4)
+    Tensor.manual_seed(0)
+
+  def test_schedule_gc(self):
+    init = bufs_allocated()
+    x = Tensor.ones(256).contiguous().realize()
+    y = Tensor.ones(5, 5).contiguous()
+    y.schedule_linear()
+    del x
+    del y
+    self.assertEqual(bufs_allocated()-init, 0)
+
+  def test_schedule_gc_with_inputs(self):
+    init = bufs_allocated()
+    x = Tensor.ones(256).contiguous().realize()
+    y = x+Tensor.ones(256).contiguous()
+    del x
+    run_linear(*y.linear_with_vars())
+    self.assertEqual(bufs_allocated()-init, 1)
+    del y
+    self.assertEqual(bufs_allocated()-init, 0)
+
+  def test_toposort_blocks_gc(self):
+    init = bufs_allocated()
+    x = Tensor.ones(4,4).contiguous().realize()+1
+    self.assertEqual(bufs_allocated()-init, 1)
+    # try commenting this part out, it's green!
+    x.uop.toposort()
+    del x
+    if bufs_allocated()-init != 0:
+      print(inspect.getclosurevars(UOp.toposort().fget))
+      raise AssertionError(f"never gced {[x for x in gc.get_objects() if isinstance(x, Buffer)]}")
+
+  def test_buffer_refcount(self):
+    init = bufs_allocated()
+    a = Tensor.empty(10)
+    self.assertEqual(bufs_allocated()-init, 0)
+    a.realize()
+    real_buf = a.uop.buffer
+    # after the Tensor UOp is deleted there shouldn't be any references on the Buffer
+    self.assertEqual(real_buf.uop_refcount, 1)
+    self.assertEqual(bufs_allocated()-init, 1)
+    del a.uop
+    self.assertEqual(real_buf.uop_refcount, 0)
+    self.assertEqual(bufs_allocated()-init, 1) # keep the buffer alive
+    del real_buf
+    self.assertEqual(bufs_allocated()-init, 0)
+
+  def test_assign_refcount(self):
+    init = bufs_allocated()
+    a = Tensor.full((4,), 1.).contiguous()
+    a.realize()
+    real_buf = a.uop.buffer
+    self.assertEqual(real_buf.uop_refcount, 1)
+    a.assign(Tensor.full((4,), 2.))
+    self.assertIs(a.uop.src[0].buffer, real_buf)
+    # NOTE: this is still 1, we don't count the ASSIGN
+    self.assertEqual(real_buf.uop_refcount, 1)
+    a.realize()
+    del a
+    self.assertEqual(real_buf.uop_refcount, 0) # no UOps for this Buffer
+    self.assertEqual(bufs_allocated()-init, 1) # Buffer is alive
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_gpudims.py

@@ -0,0 +1,110 @@
+import unittest, math
+import z3
+from tinygrad.codegen.gpudims import get_grouped_dims, add_gpudims
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.uop.validate import uops_to_z3
+from tinygrad.dtype import dtypes
+from tinygrad.renderer import Renderer
+from tinygrad.helpers import flatten, dedup, Target
+
+class TestGroupedDims(unittest.TestCase):
+  def _check_grouped_dims(self, prefix, dims, max_sizes, reverse, expected_sizes, assert_same_length=True):
+    idxs = get_grouped_dims(prefix, dims, max_sizes, reverse)
+    loop_idxs = dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs]))
+    loop_idxs = sorted(loop_idxs, key=lambda uop: uop.arg)
+    sizes = [x.src[0].arg for x in loop_idxs]
+    assert len(idxs) == len(dims), f"expected idxs to have same length as dims {len(dims)}, got {len(idxs)}"
+    if assert_same_length:
+      assert len(loop_idxs) == min(len(sizes), len(dims)), f"expected idxs to have length {min(len(sizes), len(dims))}, got {len(loop_idxs)}"
+    assert sizes == expected_sizes, f"expected sizes={expected_sizes}, got {sizes=}"
+    self._verify_indices_z3(idxs, dims)
+
+  def _verify_indices_z3(self, idxs, dims):
+    """Use z3 to prove bijectivity: bounds (0 <= flat < total) + injectivity (different inputs => different flat)."""
+    total = math.prod(dims)
+    specials = sorted(dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs])), key=lambda u: u.arg)
+    # build flat index and primed flat (same expression with renamed SPECIALs)
+    flat = UOp.const(dtypes.weakint, 0)
+    for i, idx in enumerate(idxs):
+      flat = flat + idx * int(math.prod(dims[i+1:]))
+    flat_p = flat.substitute({s: UOp(Ops.SPECIAL, s.dtype, s.src, s.arg+"_p") for s in specials})
+    solver = z3.Solver()
+    [z3_flat, z3_flat_p] = uops_to_z3(solver, flat, flat_p)
+    # bounds
+    self.assertEqual(solver.check(z3_flat < 0), z3.unsat, f"flat can be negative: {dims=}")
+    self.assertEqual(solver.check(z3_flat >= total), z3.unsat, f"flat can be >= {total}: {dims=}")
+    # injectivity: flat == flat' but inputs differ => unsat
+    inputs_differ = z3.Or(*[z3.Int(s.arg) != z3.Int(s.arg+"_p") for s in specials])
+    self.assertEqual(solver.check(z3.And(z3_flat == z3_flat_p, inputs_differ)), z3.unsat, f"not injective: {dims=}")
+
+  def test_grouped_dims(self):
+    # no-op
+    self._check_grouped_dims("gidx", (2,), (16,16,16), False, [2])
+    self._check_grouped_dims("gidx", (2,3), (16,16,16), False, [2,3])
+
+    # check reverse dims
+    self._check_grouped_dims("gidx", (2,3), (16,16,16), True, [3,2])
+    self._check_grouped_dims("gidx", (2,3,4), (16,16,16), False, [2,3,4])
+
+    # test splitting globals:    len(dims) == len(max)
+    self._check_grouped_dims("gidx", (64,3,4), (16,16,16), False, [16,12,4])
+    self._check_grouped_dims("gidx", (64,3,4), (16,4,16), False, [16,3,16])
+    self._check_grouped_dims("gidx", (64,3,4), (16,16,16), True, [16,3,16])
+    self._check_grouped_dims("gidx", (128,3,4), (16,4,256), False, [16,3,32])
+    self._check_grouped_dims("gidx", (4,4,512), (16,4,256), False, [8,4,256])
+    self._check_grouped_dims("gidx", (5,12,7), (8,4,16), False, [10,3,14])
+
+    # prefer group_dim strategy when possible
+    self._check_grouped_dims("gidx", (512,4,2), (8192,2,2), False, [2048,2])
+
+    # test splitting globals:    len(dims) < len(max)
+    #                            len(dim)        ->          len(limited)
+    #                              1             ->             2
+    self._check_grouped_dims("gidx", (128,), (16,16,256), False, [16,8], False)
+    #                              1             ->             3
+    self._check_grouped_dims("gidx", (65536,), (16,16,256), False, [16,16,256], False)
+    #                              2             ->             2
+    self._check_grouped_dims("gidx", (65536,2), (65535,65535,65535), False, [32768,4], False)
+    # test when the only divisor is the square root of dim
+    self._check_grouped_dims("gidx", (121,), (12,12,12), False, [11,11], False)
+    #                              2             ->             3
+    self._check_grouped_dims("gidx", (128,128), (16,16,256), False, [16,16,64], False)
+
+    # collapse on onto the left most axis
+    self._check_grouped_dims("gidx", (2,3,4,5), (16,16,16), False, [6,4,5])
+    self._check_grouped_dims("gidx", (2,3,4,5), (32,16,16), True, [20,3,2])
+
+    # collapse on left-most available axis (the left most is too small)
+    self._check_grouped_dims("gidx", (2,3,4,5), (4,16,16), False, [2,12,5])
+    self._check_grouped_dims("gidx", (2,3,4,5), (16,16,16), True, [5,12,2])
+
+    # dim too large and not factorable
+    with self.assertRaises(RuntimeError):
+      get_grouped_dims("gidx", (23,), (16,16,16), False,)
+    with self.assertRaises(RuntimeError):
+      get_grouped_dims("gidx", (128,3,4), (16,2,2), False,)
+
+    # too large for sizes
+    with self.assertRaises(RuntimeError):
+      get_grouped_dims("gidx", (2,3,4,5,6), (16,16,16))
+
+  def test_grouped_direct_dims_are_special(self):
+    # when (2,3) are merged into 6, the unmerged dims (4,5) should map directly to SPECIAL ops (no div/mod)
+    idxs = get_grouped_dims("gidx", (2,3,4,5), (16,16,16), False)
+    assert idxs[2].op is Ops.SPECIAL, f"expected SPECIAL for direct-mapped dim, got {idxs[2].op}"
+    assert idxs[3].op is Ops.SPECIAL, f"expected SPECIAL for direct-mapped dim, got {idxs[3].op}"
+
+  def test_global_prod_max(self):
+    g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL)
+    sink = UOp.param(0, dtypes.float.ptr()).index(g + l).store(UOp.const(dtypes.float, 1.0)).end(g, l).sink(arg=KernelInfo())
+    class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128)
+    specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL]
+    self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1)
+    self.assertGreater(len([s for s in specials if "gidx" in s.arg]), 1)
+
+  def test_max_sizes_none(self):
+    self._check_grouped_dims("gidx", (2,3,4), None, False, [2,3,4])
+    self._check_grouped_dims("gidx", (100,), None, False, [100])
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_gradient.py

@@ -0,0 +1,83 @@
+from typing import Callable
+import unittest, math
+import torch
+from tinygrad import Tensor
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import UOp
+from tinygrad.gradient import compute_gradient
+
+class TestGradient(unittest.TestCase):
+  def _cmp_nan_okay(self, x, y):
+    if math.isnan(x) and math.isnan(y): return
+    self.assertAlmostEqual(x, y, places=5)
+
+  def _test_one_input_function(self, f:Callable, jf:Callable|None=None):
+    if jf is None: jf = f
+    x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
+    gx = compute_gradient(f(x), UOp.const(dtypes.float, 1.0), set([x]))[x]
+
+    for val in [-5., -2.0, 0.0, 2.0, 5.]:
+      tg_out = gx.substitute({x: x.const_like(val)}).ssimplify()
+      tx = torch.tensor([val], dtype=torch.float, requires_grad=True)
+      torch_out = torch.autograd.grad(jf(tx), tx)[0].item()
+      self._cmp_nan_okay(tg_out, torch_out)
+
+  def _test_two_input_function(self, f:Callable, jf:Callable|None=None):
+    if jf is None: jf = f
+    x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
+    y = UOp.variable('y', -math.inf, math.inf, dtype=dtypes.float)
+    grads = compute_gradient(f(x, y), UOp.const(dtypes.float, 1.0), set([x, y]))
+    gx, gy = grads[x], grads[y]
+
+    for valx in [-5., -2.0, 0.0, 2.0, 5.]:
+      for valy in [-5., -2.0, 0.0, 2.0, 5.]:
+        # Substitute the values into the gradient expressions
+        substitutions = {x: x.const_like(valx), y: y.const_like(valy)}
+        tg_out_x = gx.substitute(substitutions).ssimplify()
+        tg_out_y = gy.substitute(substitutions).ssimplify()
+
+        tx = torch.tensor([valx], dtype=torch.float, requires_grad=True)
+        ty = torch.tensor([valy], dtype=torch.float, requires_grad=True)
+        torch_grad = torch.autograd.grad(jf(tx, ty), [tx, ty])
+        torch_out_x, torch_out_y = [x.item() for x in torch_grad]
+
+        self._cmp_nan_okay(tg_out_x, torch_out_x)
+        self._cmp_nan_okay(tg_out_y, torch_out_y)
+
+  # unary ops unit
+  def test_recip(self): self._test_one_input_function(lambda x: 1.0/x)
+  def test_sin(self): self._test_one_input_function(lambda x: x.sin())
+  def test_sqrt(self): self._test_one_input_function(lambda x: x.sqrt())
+  def test_log2(self): self._test_one_input_function(lambda x: x.log2())
+  def test_exp2(self): self._test_one_input_function(lambda x: x.exp2())
+
+  # binary ops unit
+  def test_add(self): self._test_two_input_function(lambda x,y: x+y)
+  def test_mul(self): self._test_two_input_function(lambda x,y: x*y)
+
+  # chain rule
+  def test_chain(self): self._test_one_input_function(lambda x: x.sin().sqrt())
+  def test_chain_binop(self): self._test_two_input_function(lambda x,y: (x*y)+x*y)
+  def test_big_add_sin(self): self._test_two_input_function(lambda x,y: x.sin()+3.0/y)
+  def test_big_chain(self): self._test_two_input_function(lambda x,y: (1.0/x*y)+x*y)
+  def test_where(self): self._test_two_input_function(lambda x,y: (x<y).where(x,y), lambda x,y: torch.where(x<y,x,y))
+
+class TestRealizeMeansRealize(unittest.TestCase):
+  def test_randn_realizes(self):
+    x = Tensor.randn(2, 3, 64, 64).realize()
+    assert x.uop is not x.uop.base
+    assert x.uop.is_realized
+
+  def test_uniform_realizes(self):
+    x = Tensor.uniform(16, 3, 3, 3).realize()
+    print(x.uop)
+    assert x.uop is not x.uop.base
+    assert x.uop.is_realized
+
+  def test_uniform_gradient(self):
+    x = Tensor.uniform(16, 3, 3, 3).realize()
+    y = x * 2
+    y.sum().gradient(x)[0].realize()
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_graph_rewrite.py

@@ -0,0 +1,537 @@
+import unittest, math
+from tinygrad import dtypes
+from tinygrad.helpers import all_same, Context
+from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, TrackedPatternMatcher, UPat
+from test.helpers import full_rewrite
+from hypothesis import given, strategies as strat
+
+# Helper function to apply the graph rewrite
+@Context(SPEC=0)
+def apply_rewrite(expr):
+  return full_rewrite(expr.sink()).src[0]
+
+@Context(SPEC=0)
+def apply_rewrite_values(expr):
+  srcs = full_rewrite(expr.sink()).src
+  if len(srcs) == 1:
+    if srcs[0].op is Ops.CONST: return (srcs[0].arg,)*srcs[0].dtype.count
+    if srcs[0].op is Ops.STACK: return tuple(s.arg for s in srcs[0].src)
+  return tuple(s.arg for s in srcs)
+
+def evaluate_uop(uop, variables):
+  if uop.op == Ops.CONST:
+    return uop.arg
+  elif uop.op == Ops.DEFINE_VAR:
+    var_name = uop.arg[0]
+    return variables[var_name]
+  elif uop.op in GroupOp.ALU:
+    src_values = [evaluate_uop(src, variables) for src in uop.src]
+    return exec_alu(uop.op, uop.dtype, src_values)
+  else:
+    raise NotImplementedError(f"Unsupported UOp {uop.op}")
+
+class TestArithmeticSimplifications(unittest.TestCase):
+  def test_full_graph_rewrite_division_by_zero(self):
+    optimized_div_uop = apply_rewrite(UOp.const(dtypes.float32, 10.0) / UOp.const(dtypes.float32, 0.0))
+    self.assertEqual(optimized_div_uop.op, Ops.CONST)
+    self.assertTrue(math.isinf(optimized_div_uop.arg) or math.isnan(optimized_div_uop.arg))
+
+  def test_full_graph_rewrite_redundant_operations(self):
+    optimized_uop = apply_rewrite((UOp.const(dtypes.float32, 10.0) + UOp.const(dtypes.float32, 0.0)) * UOp.const(dtypes.float32, 1.0))
+    self.assertEqual(optimized_uop.op, Ops.CONST)
+    self.assertEqual(optimized_uop.arg, 10.0)
+
+  def test_full_graph_rewrite_large_graph(self):
+    prev_uop = UOp.const(dtypes.int32, 0)
+    for i in range(1, 101):
+      prev_uop += UOp.const(dtypes.int32, i)
+    optimized_uop = apply_rewrite(prev_uop)
+    self.assertEqual(optimized_uop.op, Ops.CONST)
+    self.assertEqual(optimized_uop.arg, sum(range(1, 101)))
+
+  def test_full_graph_rewrite_division_by_one(self):
+    optimized_uop = apply_rewrite(UOp.const(dtypes.float32, 42.0) / UOp.const(dtypes.float32, 1.0))
+    self.assertEqual(optimized_uop.op, Ops.CONST)
+    self.assertEqual(optimized_uop.arg, 42.0)
+
+  def test_full_graph_rewrite_modulo_by_one(self):
+    optimized_uop = apply_rewrite(UOp.const(dtypes.int32, 42) % UOp.const(dtypes.int32, 1))
+    self.assertEqual(optimized_uop.op, Ops.CONST)
+    self.assertEqual(optimized_uop.arg, 0)
+
+
+class TestFoldingAndReduction(unittest.TestCase):
+  @unittest.skip("reduce is removed now")
+  def test_full_graph_rewrite_constant_reduction_folding(self):
+    const1 = UOp.const(dtypes.int32, 5)
+    const2 = UOp.const(dtypes.int32, 10)
+    const3 = UOp.const(dtypes.int32, 20)
+    optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD))
+    expected_sum = 5 + 10 + 20
+    self.assertEqual(optimized_sink.arg, expected_sum)
+
+  @unittest.skip("reduce is removed now")
+  def test_full_graph_rewrite_reduction_with_unused_range(self):
+    const1 = UOp.const(dtypes.int32, 15)
+    const2 = UOp.const(dtypes.int32, 25)
+    rng = UOp.range(10, idx=0)
+    optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng))
+    expected_sum = 10 * (15 + 25)
+    self.assertEqual(optimized_sink.arg, expected_sum)
+
+  @unittest.skip("currently failing")
+  def test_full_graph_rewrite_range_reduction(self):
+    simple_range = UOp.range(5, idx=0)
+    optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range))
+    expected_sum = sum(range(5))
+    self.assertEqual(optimized_sink.arg, expected_sum)
+
+  @unittest.skip("currently failing")
+  def test_full_graph_rewrite_simple_reduction_folding(self):
+    simple_range = UOp.range(4, idx=0)
+    add_uop = simple_range + UOp.const(dtypes.int32, 1)
+    optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range))
+    expected_sum = sum(i + 1 for i in range(4))
+    self.assertEqual(optimized_sink.arg, expected_sum)
+
+  @unittest.skip("currently failing")
+  def test_full_graph_rewrite_nested_loop_collapse(self):
+    outer_range = UOp.range(8, 0)
+    inner_range = UOp.range(4, 1)
+    expr = (outer_range * 10) + inner_range
+    optimized_reduce_uop = apply_rewrite(expr.reduce(Ops.ADD, outer_range, inner_range))
+    self.assertEqual(optimized_reduce_uop.op, Ops.CONST)
+    self.assertEqual(optimized_reduce_uop.arg, sum((i * 10) + j for i in range(8) for j in range(4)))
+
+
+class TestModuloAndDivisionFolding(unittest.TestCase):
+  def test_full_graph_rewrite_modulo_folding_with_define_var(self):
+    # index dtype because div-mod rules only work on index
+    x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint)
+    optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
+    self.assertEqual(optimized_mod_uop.op, Ops.CONST)
+    self.assertEqual(optimized_mod_uop.arg, 2)
+
+  def test_full_graph_rewrite_division_folding_with_define_var(self):
+    # index dtype because div-mod rules only work on index
+    n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint)
+    optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3)
+    self.assertEqual(optimized_div_uop.op, Ops.MUL)
+    self.assertEqual(optimized_div_uop.src[1].arg, 2)
+
+  def test_full_graph_rewrite_complex_mod_div_folding(self):
+    # index dtype because div-mod rules only work on index
+    k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint)
+    optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2)
+    self.assertEqual(optimized_div_uop.op, Ops.CONST)
+    self.assertEqual(optimized_div_uop.arg, 1)
+
+  def test_graph_rewrite_div_folding_bug(self):
+    lhs = UOp(Ops.ADD, dtypes.int.vec(4), src=(
+      UOp(Ops.STACK, dtypes.int.vec(4), arg=None, src=(UOp(Ops.SPECIAL, dtypes.int, arg='lidx0', src=(UOp.const(dtypes.int, 32),)),)*4),
+      UOp.const(dtypes.int.vec(4), (0, 256, 512, 768))))
+    rhs = UOp.const(dtypes.int.vec(4), 2)
+    unopt = lhs<rhs
+    opt = apply_rewrite(unopt)
+    print(unopt)
+    print(opt)
+    if opt.op is Ops.STACK: self.assertFalse(all_same(opt.src))
+
+  def test_full_graph_rewrite_modulo_large_divisor(self):
+    # index dtype because div-mod rules only work on index
+    x_var_uop = UOp.variable('x', 1, 5)
+    self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.weakint) % 10).render(simplify=False), x_var_uop.render(simplify=False))
+
+  def test_full_graph_rewrite_division_with_remainder(self):
+    x_var_uop = UOp.variable('x', 7, 9)
+    optimized_sink = apply_rewrite(x_var_uop // 2)
+    for x_value in range(7, 10):
+      self.assertEqual(x_value // 2, evaluate_uop(optimized_sink, {'x': x_value}))
+
+  def test_full_graph_rewrite_complex_mod_div_expression(self):
+    x_var_uop = UOp.variable('x', 1, 10)
+    optimized_sink = apply_rewrite(((x_var_uop * 5) % 3) // 2)
+    for x_value in range(1, 11):
+      original_result = ((x_value * 5) % 3) // 2
+      optimized_result = evaluate_uop(optimized_sink, {'x': x_value})
+      self.assertEqual(original_result, optimized_result)
+
+
+class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
+  def test_full_graph_rewrite_transcendental_edge_cases(self):
+    optimized_sink = full_rewrite(UOp.const(dtypes.float32, -1.0).log2().sink(UOp.const(dtypes.float32, 0.0).reciprocal()))
+    optimized_log2_neg, optimized_recip_zero = optimized_sink.src
+    self.assertTrue(math.isnan(optimized_log2_neg.arg), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.arg}")
+    self.assertTrue(math.isinf(optimized_recip_zero.arg) and optimized_recip_zero.arg > 0,
+                    f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.arg}")
+
+  @unittest.skip("broken")
+  def test_full_graph_rewrite_modulo_negative_dividend(self):
+    x_var_uop = UOp.variable('x', -5, -1)
+    optimized_sink = full_rewrite((x_var_uop % 3).sink())
+    for x_value in range(-5, 0):
+      self.assertEqual(x_value % 3, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
+
+  @unittest.skip("broken")
+  def test_full_graph_rewrite_division_negative_divisor(self):
+    x_var_uop = UOp.variable('x', 1, 5)
+    optimized_sink = full_rewrite((x_var_uop // -2).sink())
+    for x_value in range(1, 6):
+      self.assertEqual(x_value // -2, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
+
+class TestGEPAndVectorizeRewrite(unittest.TestCase):
+  def test_gep_single_element_extraction(self):
+    # GEP on a vector dtype to extract a single element
+    base_vector = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(apply_rewrite(base_vector.gep(2)).arg, 3.0)
+
+  def test_gep_tuple_extraction(self):
+    # GEP on a vector dtype to extract multiple elements as a vector
+    base_vector = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(list(apply_rewrite_values(base_vector.gep((2, 3)))), [3.0, 4.0])
+
+  def test_gep_on_const_stack(self):
+    # GEP on a const STACK to extract a single element
+    const_stack = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(apply_rewrite(const_stack.gep(2)).arg, 3.0)
+
+  def test_gep_tuple_on_const_stack(self):
+    # GEP on a const STACK using a tuple to extract multiple elements
+    const_stack = UOp.const(dtypes.float32.vec(4), (7.0, 8.0, 9.0, 10.0))
+    self.assertEqual(list(apply_rewrite_values(const_stack.gep((1, 3)))), [8.0, 10.0])
+
+  def test_gep_gep_simplification(self):
+    # Nested GEP simplification on a vector dtype
+    base_vector = UOp.const(dtypes.float32.vec(4), (10.0, 20.0, 30.0, 40.0))
+    gep_inner = base_vector.gep(1)  # Extract 2nd element (20.0)
+    self.assertEqual(apply_rewrite(gep_inner.gep(0)).arg, 20.0)
+
+  def test_vectorize_multiple_elements(self):
+    # Vectorizing multiple elements using GEP
+    base_vector = UOp.const(dtypes.float32.vec(4), (5.0, 10.0, 15.0, 20.0))
+    vectorized_uop = UOp(Ops.STACK, dtypes.float32.vec(4), src=(base_vector.gep(0), base_vector.gep(1), base_vector.gep(2), base_vector.gep(3)))
+    self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])
+
+
+import inspect
+from tinygrad.uop.ops import graph_rewrite, _substitute, track_rewrites
+from tinygrad.uop.symbolic import symbolic_simple
+
+class TestBottomUpRewrite(unittest.TestCase):
+  def test_const_folding(self):
+    a = UOp.const(dtypes.int, 5)
+    ret = (a*3) + (a*7)
+    gt = graph_rewrite(ret, symbolic_simple)
+    ret = graph_rewrite(ret, symbolic_simple, bottom_up=True)
+    self.assertIs(gt, ret)
+
+# normally .substitute would be fine, but it's not tracked
+@track_rewrites()
+def named_substitute(name:str, uop:UOp, rel:dict[UOp, UOp]): return graph_rewrite(uop, _substitute, rel, bottom_up=True)
+def substitute(uop:UOp, rel:dict[UOp, UOp]): return named_substitute(inspect.stack()[1].function, uop, rel)
+
+class TestSubstitute(unittest.TestCase):
+  # these work because the substituted things don't have parents
+  def test_simple(self):
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    ret = a + 4
+    ret = substitute(ret, {a:b})
+    self.assertIs(ret, b+4)
+
+  def test_double(self):
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    c = UOp.variable('c', 0, 10)
+    ret = (a + 4) + b
+    ret = substitute(ret, {a:c, b:c})
+    self.assertIs(ret, (c + 4) + c)
+
+  def test_diamond(self):
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    ret = (a + 4) + (a + 5)
+    ret = substitute(ret, {a:b})
+    self.assertIs(ret, (b + 4) + (b + 5))
+
+  # this works because there's nothing above the substituted node
+  def test_sin(self):
+    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
+    b = UOp.variable('b', 0, 10, dtype=dtypes.float)
+    ret = a.sin().sin()
+    ret = substitute(ret, {a.sin():b})
+    self.assertIs(ret, b.sin())
+
+  # broken due to infinite recursion
+  # NOTE: VIZ hangs and doesn't recover if you click this one
+  @unittest.skip("recursion error no longer raised")
+  def test_assert_inf_recurse(self):
+    a = UOp.variable('a', 0, 10)
+    n1 = a.sin()
+    ret = n1
+    with self.assertRaises(RecursionError):
+      ret = substitute(ret, {n1:n1.sqrt()})
+
+  def test_sin_to_sqrt(self):
+    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
+    n1 = a.sin()
+    ret = n1.sin()
+    ret = substitute(ret, {a.sin():a.sqrt()})
+    self.assertIs(ret, a.sqrt().sin())
+
+  def test_double_sin_to_sqrt(self):
+    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
+    n1 = a.sin()
+    ret = n1.sin()
+    # NOTE: this would work if it had gone in the opposite order
+    ret = substitute(ret, {a.sin():a.sqrt(), n1.sin():n1.sqrt()})
+    self.assertIs(ret, a.sqrt().sqrt())
+
+  def test_tagged_replace(self):
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    ret = (a+4).replace(tag=1)
+    ret = substitute(ret, {a:b})
+    # the srcs are rewritten but we keep tag
+    self.assertIs(ret, (b+4).replace(tag=1))
+
+matchers = strat.sampled_from([PatternMatcher, TrackedPatternMatcher])
+
+class TestRecurse(unittest.TestCase):
+  @given(matchers)
+  def test_no_inf_loop(self, PatternMatcher):
+    a = UOp.variable('a', 0, 10)
+    pm = PatternMatcher([(UPat(Ops.DEFINE_VAR, name="x"), lambda x: x)])
+    graph_rewrite(a, pm)
+
+  @given(matchers)
+  def test_no_inf_loop_bottom_up(self, PatternMatcher):
+    a = UOp.variable('a', 0, 10)
+    pm = PatternMatcher([(UPat(Ops.DEFINE_VAR, name="x"), lambda x: x)])
+    graph_rewrite(a, pm, bottom_up=True)
+
+  def test_inf_loop(self):
+    a = UOp.const(dtypes.int, 3)
+    pm = PatternMatcher([
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
+    ])
+    with self.assertRaises(RuntimeError):
+      graph_rewrite(a, pm)
+
+  def test_inf_loop_bottom_up(self):
+    a = UOp.const(dtypes.int, 3)
+    pm = PatternMatcher([
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
+    ])
+    with self.assertRaises(RuntimeError):
+      graph_rewrite(a, pm, bottom_up=True)
+
+def bidir_append(ctx, x, b): ctx.append((x.arg if x.op is Ops.CONST else "+", b))
+class TestBidirectional(unittest.TestCase):
+  def test_simple(self):
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
+    c = a + b
+    pm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, False)) ])
+    bpm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, True)) ])
+    ctx_list = []
+    graph_rewrite(c, pm, ctx=ctx_list, bpm=bpm)
+    self.assertListEqual(ctx_list, [('+', True), (1, True), (1, False), (2, True), (2, False), ('+', False)])
+
+class TestStopEarly(unittest.TestCase):
+  def test_stop_early(self):
+    a = UOp.const(dtypes.int, 3)
+    b = UOp.const(dtypes.int, 4)
+    c = a+b
+    cn = UOp.const(dtypes.int, 7)
+    d = UOp.const(dtypes.int, 2)
+    def visit_const(c:UOp):
+      print(f"visit {c.arg}")
+      assert c.arg not in (3,4)
+    pm_cvisit = PatternMatcher([(UPat(Ops.CONST, name="c"), visit_const),])
+    ret = (c+d).substitute({c:cn}, extra_pm=pm_cvisit)
+    assert ret == cn+d
+
+class TestWalkRewrite(unittest.TestCase):
+  """Tests for graph_rewrite with walk=True (MLIR Walk Pattern Rewrite Driver semantics).
+  walk=True gives a single-pass traversal that does NOT revisit or re-traverse into rewritten subtrees.
+  Supports both top-down (default) and bottom-up (bottom_up=True) modes."""
+
+  # *** top-down walk (default): process children first, then try pm on rebuilt node ***
+
+  def test_walk_topdown_simple_substitute(self):
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    ret = graph_rewrite(a + 4, _substitute, {a:b}, walk=True)
+    self.assertIs(ret, b+4)
+
+  def test_walk_topdown_does_not_traverse_into_replacement(self):
+    """Top-down walk: replacement subtrees are NOT re-entered."""
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    c = UOp.variable('c', 0, 10)
+    d = UOp.variable('d', 0, 10)
+    # a is replaced by b+c, but b inside the replacement is NOT further substituted to d
+    ret_walk = graph_rewrite(a + 4, _substitute, {a:b+c, b:d}, walk=True)
+    self.assertIs(ret_walk, (b+c)+4)
+    # contrast: greedy bottom_up WOULD replace b inside the replacement
+    ret_greedy = graph_rewrite(a + 4, _substitute, {a:b+c, b:d}, bottom_up=True)
+    self.assertIs(ret_greedy, (d+c)+4)
+
+  def test_walk_topdown_no_fixed_point(self):
+    """A bouncing pattern applies once and stops instead of looping."""
+    a = UOp.const(dtypes.int, 3)
+    pm = PatternMatcher([
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
+    ])
+    with self.assertRaises(RuntimeError):
+      graph_rewrite(a, pm, bottom_up=True)
+    ret = graph_rewrite(a, pm, walk=True)
+    self.assertIs(ret, UOp.const(dtypes.int, 4))
+
+  def test_walk_topdown_rewrites_children(self):
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    c = UOp.variable('c', 0, 10)
+    ret = graph_rewrite((a + 4) + (b + 5), _substitute, {a:c, b:c}, walk=True)
+    self.assertIs(ret, (c + 4) + (c + 5))
+
+  def test_walk_topdown_diamond(self):
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    ret = graph_rewrite((a + 4) + (a + 5), _substitute, {a:b}, walk=True)
+    self.assertIs(ret, (b + 4) + (b + 5))
+
+  def test_walk_topdown_children_rewritten_before_parent(self):
+    """Top-down walk processes children first: child substitution changes the rebuilt parent."""
+    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
+    n1 = a.sin()          # sin(a)
+    ret = n1.sin()         # sin(sin(a))
+    # sin(a)->sqrt(a) fires first (child), parent rebuilds to sin(sqrt(a)), which doesn't match sin(sin(a)) in dvars
+    ret_walk = graph_rewrite(ret, _substitute, {a.sin():a.sqrt(), n1.sin():n1.sqrt()}, walk=True)
+    self.assertIs(ret_walk, a.sqrt().sin())
+
+  def test_walk_topdown_self_referential_replacement(self):
+    """Replacement containing the replaced node works without infinite recursion."""
+    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
+    ret = graph_rewrite(a.sin() + 4, _substitute, {a.sin(): a.sin().sqrt()}, walk=True)
+    self.assertIs(ret, a.sin().sqrt() + 4)
+
+  def test_walk_topdown_visit_order(self):
+    """Top-down walk fires pm after children are processed (post-order)."""
+    visited = []
+    def track_visit(ctx, x):
+      ctx.append(x.arg if x.op is Ops.CONST else x.op)
+      return None
+    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
+    graph_rewrite(a + b, pm, ctx=visited, walk=True)
+    self.assertEqual(visited, [1, 2, Ops.ADD])
+
+  # *** bottom-up walk: try bpm on node first, skip children if it matches ***
+
+  def test_walk_bottomup_simple_substitute(self):
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    ret = graph_rewrite(a + 4, _substitute, {a:b}, bottom_up=True, walk=True)
+    self.assertIs(ret, b+4)
+
+  def test_walk_bottomup_does_not_traverse_into_replacement(self):
+    """Bottom-up walk: replacement subtrees are NOT entered."""
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    c = UOp.variable('c', 0, 10)
+    d = UOp.variable('d', 0, 10)
+    ret = graph_rewrite(a + 4, _substitute, {a:b+c, b:d}, bottom_up=True, walk=True)
+    self.assertIs(ret, (b+c)+4)
+
+  def test_walk_bottomup_parent_match_skips_children(self):
+    """Bottom-up walk matches parent first: if it matches, children are never visited."""
+    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
+    n1 = a.sin()
+    ret = n1.sin()         # sin(sin(a))
+    # sin(sin(a)) matches n1.sin()->n1.sqrt() immediately, children never visited, sin(a) inside replacement untouched
+    ret_walk = graph_rewrite(ret, _substitute, {a.sin():a.sqrt(), n1.sin():n1.sqrt()}, bottom_up=True, walk=True)
+    self.assertIs(ret_walk, a.sin().sqrt())
+
+  def test_walk_bottomup_no_fixed_point(self):
+    """Bottom-up walk also applies once per node, no fixed-point iteration."""
+    a = UOp.const(dtypes.int, 3)
+    pm = PatternMatcher([
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
+    ])
+    ret = graph_rewrite(a, pm, bottom_up=True, walk=True)
+    self.assertIs(ret, UOp.const(dtypes.int, 4))
+
+  def test_walk_bottomup_visit_order(self):
+    """Bottom-up walk fires bpm before descending (pre-order)."""
+    visited = []
+    def track_visit(ctx, x):
+      ctx.append(x.arg if x.op is Ops.CONST else x.op)
+      return None
+    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
+    graph_rewrite(a + b, pm, ctx=visited, bottom_up=True, walk=True)
+    # bpm fires on each node before children: +, 1, 2
+    self.assertEqual(visited, [Ops.ADD, 1, 2])
+
+  def test_walk_bottomup_unmatched_falls_through_to_children(self):
+    """Bottom-up walk: if bpm doesn't match a node, its children are still processed."""
+    a = UOp.variable('a', 0, 10)
+    b = UOp.variable('b', 0, 10)
+    c = UOp.variable('c', 0, 10)
+    # only a is in dvars, not a+4. bpm won't match a+4, so it descends and finds a.
+    ret = graph_rewrite((a + 4) + (b + 5), _substitute, {a:c, b:c}, bottom_up=True, walk=True)
+    self.assertIs(ret, (c + 4) + (c + 5))
+
+  # *** bidirectional walk: bpm fires before children, pm fires after rebuild ***
+
+  def test_walk_bidirectional_visit_order(self):
+    """Bidirectional walk: bpm fires pre-order, pm fires post-order."""
+    visited = []
+    def bpm_visit(ctx, x):
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm"))
+      return None
+    def pm_visit(ctx, x):
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
+      return None
+    bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_visit)])
+    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_visit)])
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
+    graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
+    # bpm fires pre-order, pm fires post-order
+    self.assertEqual(visited, [
+      (Ops.ADD, "bpm"), (1, "bpm"), (1, "pm"), (2, "bpm"), (2, "pm"), (Ops.ADD, "pm"),
+    ])
+
+  def test_walk_bidirectional_bpm_short_circuits(self):
+    """If bpm matches, children are skipped and pm never fires on that node."""
+    visited = []
+    def bpm_match(ctx, x):
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm"))
+      # rewrite const(1) -> const(10), short-circuiting its subtree
+      if x.op is Ops.CONST and x.arg == 1: return x.replace(arg=10)
+      return None
+    def pm_match(ctx, x):
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
+      return None
+    bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_match)])
+    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_match)])
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
+    ret = graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
+    # bpm matches const(1) and short-circuits it, so pm never fires on const(1)
+    self.assertNotIn((1, "pm"), visited)
+    # but pm still fires on const(2) and the rebuilt ADD
+    self.assertIn((2, "pm"), visited)
+    self.assertIs(ret, UOp.const(dtypes.int, 10) + b)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_hcq_iface.py

@@ -0,0 +1,105 @@
+import unittest, array, time
+from tinygrad.helpers import mv_address
+from tinygrad.runtime.support.hcq import MMIOInterface
+from tinygrad.runtime.support.usb import USBMMIOInterface
+from test.mockgpu.usb import MockUSB
+
+class TestHCQIface(unittest.TestCase):
+  def setUp(self):
+    self.size = 4 << 10
+    self.buffer = bytearray(self.size)
+    self.mv = memoryview(self.buffer).cast('I')
+    self.mmio = MMIOInterface(mv_address(self.mv), self.size, fmt='I')
+
+  def test_getitem_setitem(self):
+    self.mmio[1] = 0xdeadbeef
+    self.assertEqual(self.mmio[1], 0xdeadbeef)
+    values = array.array('I', [10, 20, 30, 40])
+    self.mmio[2:6] = values
+    read_slice = self.mmio[2:6]
+    # self.assertIsInstance(read_slice, array.array)
+    self.assertEqual(read_slice, values.tolist())
+    self.assertEqual(self.mv[2:6].tolist(), values.tolist())
+
+  def test_view(self):
+    full = self.mmio.view()
+    self.assertEqual(len(full), len(self.mmio))
+    self.mmio[0] = 0x12345678
+    self.assertEqual(full[0], 0x12345678)
+
+    # offset-only view
+    self.mmio[1] = 0xdeadbeef
+    off = self.mmio.view(offset=4)
+    self.assertEqual(off[0], 0xdeadbeef)
+
+    # offset + size view: write into sub-view and confirm underlying buffer
+    values = array.array('I', [11, 22, 33])
+    sub = self.mmio.view(offset=8, size=12)
+    sub[:] = values
+    self.assertEqual(sub[:], values.tolist())
+    self.assertEqual(self.mv[2:5].tolist(), values.tolist())
+
+  def test_speed(self):
+    start = time.perf_counter()
+    for i in range(10000):
+      self.mmio[3:100] = array.array('I', [i] * 97)
+      _ = self.mmio[3:100]
+    end = time.perf_counter()
+
+    mvstart = time.perf_counter()
+    for i in range(10000):
+      self.mv[3:100] = array.array('I', [i] * 97)
+      _ = self.mv[3:100].tolist()
+    mvend = time.perf_counter()
+    print(f"speed: hcq {end - start:.6f}s vs plain mv {mvend - mvstart:.6f}s")
+
+class TestUSBMMIOInterface(unittest.TestCase):
+  def setUp(self):
+    self.size = 256
+    self.buffer = bytearray(self.size)
+    self.usb = MockUSB(self.buffer)
+    self.mmio = USBMMIOInterface(self.usb, 0, self.size, fmt='B', pcimem=False)
+
+  def test_getitem_setitem_byte(self):
+    self.mmio[1] = 0xAB
+    self.assertEqual(self.mmio[1], 0xAB)
+    self.assertEqual(self.usb.mem[1], 0xAB)
+
+  def test_slice_getitem_setitem(self):
+    values = [1, 2, 3, 4]
+    self.mmio[10:14] = values
+    raw = self.mmio[10:14]
+    self.assertIsInstance(raw, bytes)
+    self.assertEqual(list(raw), values)
+    self.assertEqual(list(self.usb.mem[10:14]), values)
+
+  def test_view(self):
+    self.mmio[0] = 5
+    view = self.mmio.view(offset=1, size=3)
+    self.assertEqual(view[0], self.usb.mem[1])
+    view[:] = [7, 8, 9]
+    self.assertEqual(list(self.usb.mem[1:4]), [7, 8, 9])
+    full_view = self.mmio.view()
+    self.assertEqual(len(full_view), len(self.mmio))
+    self.mmio[2] = 0xFE
+    self.assertEqual(full_view[2], 0xFE)
+
+  def test_pcimem_byte(self):
+    usb2 = MockUSB(bytearray(self.size))
+    mmio_pci = USBMMIOInterface(usb2, 0, self.size, fmt='B', pcimem=True)
+    mmio_pci[3] = 0x11
+    self.assertEqual(mmio_pci[3], 0x11)
+    self.assertEqual(usb2.mem[3], 0x11)
+
+  def test_pcimem_slice(self):
+    usb3 = MockUSB(bytearray(self.size))
+    mmio_pci = USBMMIOInterface(usb3, 0, self.size, fmt='B', pcimem=True)
+    values = [2, 3, 4]
+    mmio_pci[4:7] = values
+    raw = mmio_pci[4:7]
+    self.assertIsInstance(raw, bytes)
+    self.assertEqual(list(raw), values)
+    self.assertEqual([mmio_pci[i] for i in range(4, 7)], values)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/null/test_helpers.py

@@ -0,0 +1,481 @@
+import ctypes, gzip, unittest, timeit, pickle
+from tinygrad import Variable
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction, count, all_same
+from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
+from tinygrad.helpers import ceildiv, ansistrip
+from tinygrad.tensor import Tensor, get_shape
+import numpy as np
+
+VARIABLE = ContextVar("VARIABLE", 0)
+
+class TestContextVars(unittest.TestCase):
+  # Ensuring that the test does not modify variables outside the tests.
+  ctx = Context()
+  def setUp(self): TestContextVars.ctx.__enter__()
+  def tearDown(self): TestContextVars.ctx.__exit__()
+
+  def test_initial_value_is_set(self):
+    _TMP = ContextVar("_TMP", 5)
+    self.assertEqual(_TMP.value, 5)
+
+  def test_cannot_recreate(self):
+    _TMP2 = ContextVar("_TMP2", 1)
+    with self.assertRaises(RuntimeError):
+      _TMP2 = ContextVar("_TMP2", 2)
+
+  def test_new_var_inside_context(self):
+    with Context(VARIABLE=1):
+      _TMP3 = ContextVar("_TMP3", 1)
+    with self.assertRaises(RuntimeError):
+      _TMP3 = ContextVar("_TMP3", 2)
+
+  def test_value_across_modules(self):
+    # Mocking module import by invoking the code but not in our globals().
+    exec('from tinygrad.helpers import ContextVar;C = ContextVar("C", 13)', {}) # pylint:disable=exec-used
+    # It should not matter that the first creation was in another module.
+    with self.assertRaises(RuntimeError):
+      _C = ContextVar("C", 0)
+
+  def test_assignment_across_modules(self):
+    B = ContextVar("B", 1)
+    # local assignment
+    B.value = 2
+    self.assertEqual(B.value, 2)
+    with self.assertRaises(RuntimeError):
+      # Assignment in another module.
+      exec('from tinygrad.helpers import ContextVar;B = ContextVar("B", 0);B.value = 3;', {}) # pylint:disable=exec-used
+
+  def test_context_assignment(self):
+    with Context(VARIABLE=1):
+      self.assertEqual(VARIABLE.value, 1)
+    self.assertEqual(VARIABLE.value, 0)
+
+  def test_unknown_param_to_context(self):
+    with self.assertRaises(KeyError):
+      with Context(SOMETHING_ELSE=1):
+        pass
+
+  def test_nested_context(self):
+    with Context(VARIABLE=1):
+      with Context(VARIABLE=2):
+        MORE = ContextVar("MORE", 2)
+        with Context(VARIABLE=3, MORE=3):
+          self.assertEqual(VARIABLE.value, 3)
+          self.assertEqual(MORE.value, 3)
+        self.assertEqual(VARIABLE.value, 2)
+        self.assertEqual(MORE.value, 2)
+      self.assertEqual(VARIABLE.value, 1)
+      self.assertEqual(MORE.value, 2)  # TODO: should this raise?
+    self.assertEqual(VARIABLE.value, 0)
+
+  def test_decorator(self):
+    @Context(VARIABLE=1, DEBUG=4)
+    def test():
+      self.assertEqual(VARIABLE.value, 1)
+
+    self.assertEqual(VARIABLE.value, 0)
+    test()
+    self.assertEqual(VARIABLE.value, 0)
+
+  def test_context_exit_reverts_updated_values(self):
+    D = ContextVar("D", 1)
+    D.value = 2
+    with Context(D=3):
+      ...
+    assert D.value == 2, f"Expected D to be 2, but was {D.value}. Indicates that Context.__exit__ did not restore to the correct value."
+
+class TestAllSame(unittest.TestCase):
+  def test_empty(self): self.assertTrue(all_same([]))
+  def test_single(self): self.assertTrue(all_same([1]))
+  def test_same(self): self.assertTrue(all_same([1, 1, 1]))
+  def test_different(self): self.assertFalse(all_same([1, 2, 1]))
+
+class TestMergeDicts(unittest.TestCase):
+  def test_merge_dicts(self):
+    a = {"a": 1, "b": 2}
+    b = {"a": 1, "c": 3}
+    c = {}
+    d = {"a": 2, "b": 2}
+    assert merge_dicts([a, b]) == {"a": 1, "b": 2, "c": 3}
+    assert merge_dicts([a, c]) == a
+    assert merge_dicts([a, b, c]) == {"a": 1, "b": 2, "c": 3}
+    with self.assertRaises(RuntimeError):
+      merge_dicts([a, d])
+
+class TestStripParens(unittest.TestCase):
+  def test_simple(self): self.assertEqual("1+2", strip_parens("(1+2)"))
+  def test_nested(self): self.assertEqual("1+(2+3)", strip_parens("(1+(2+3))"))
+  def test_casted_no_strip(self): self.assertEqual("(int)(1+2)", strip_parens("(int)(1+2)"))
+  def test_unmatched_parens(self): self.assertEqual("((c35+c39>>23&255)+-127).cast(dtypes.float)",
+    strip_parens("((c35+c39>>23&255)+-127).cast(dtypes.float)"))
+  def test_single_paren_left(self): self.assertEqual("(abc", strip_parens("(abc"))
+  def test_single_paren_right(self): self.assertEqual("abc)", strip_parens("abc)"))
+  def test_parens_at_different_depths(self): self.assertEqual("(a+(b))*(c)", strip_parens("(a+(b))*(c)"))
+
+class TestProd(unittest.TestCase):
+  def test_empty(self): self.assertEqual(1, prod(tuple()))
+  def test_ints(self): self.assertEqual(30, prod((2, 3, 5)))
+  def test_variable(self): self.assertEqual("(a*12)", prod((Variable("a", 1, 5), 3, 4)).render())
+  def test_variable_order(self): self.assertEqual("(a*12)", prod((3, 4, Variable("a", 1, 5))).render())
+
+class TestRoundUp(unittest.TestCase):
+  def test_round_up(self):
+    self.assertEqual(round_up(-3,4), 0)
+    self.assertEqual(round_up(-4,4), -4)
+    self.assertEqual(round_up(6,4), 8)
+    self.assertEqual(round_up(8,4), 8)
+    self.assertEqual(round_up(232, 24984), 24984)
+    self.assertEqual(round_up(24984, 232), 25056)
+
+class TestCeilDiv(unittest.TestCase):
+  def test_int(self):
+    self.assertEqual(ceildiv(10, 3), 4)
+    self.assertEqual(ceildiv(9, 3), 3)
+    self.assertEqual(ceildiv(0, 5), 0)
+    self.assertEqual(ceildiv(1, 5), 1)
+  def test_symbolic(self):
+    # tests that ceildiv with UOp uses (num + amt - 1) // amt formula for non-negative num
+    v = Variable('v', 0, 100)
+    result = ceildiv(v, 6)
+    self.assertEqual(result.render(), "((v+5)//6)")
+  def test_symbolic_negative_offset(self):
+    # tests ceildiv(v-5, 6) which is used in conv2d output shape
+    # old implementation incorrectly simplified -(x//-y) to ((v+1)//6-1) for v-5
+    # new implementation uses (v-5+5)//6 = v//6 which is correct
+    v = Variable('v', 11, 100)
+    result = ceildiv(v - 5, 6)
+    self.assertEqual(result.render(), "(v//6)")
+
+class TestCount(unittest.TestCase):
+  def test_count_basic(self):
+    c = count(3)
+    self.assertEqual(next(c), 3)
+    self.assertEqual(next(c), 4)
+
+  def test_count_step_pickle(self):
+    c = count(1, 2)
+    self.assertEqual(next(c), 1)
+    c2 = pickle.loads(pickle.dumps(c))
+    self.assertEqual(next(c2), 3)
+
+@unittest.skip("no fetch tests because they need internet")
+class TestFetch(unittest.TestCase):
+  def test_fetch_bad_http(self):
+    self.assertRaises(Exception, fetch, 'http://www.google.com/404', allow_caching=False)
+
+  def test_fetch_small(self):
+    assert (len(fetch('https://google.com', allow_caching=False).read_bytes())>0)
+
+  def test_fetch_img(self):
+    from PIL import Image
+    img = fetch("https://avatars.githubusercontent.com/u/132956020", allow_caching=False)
+    with Image.open(img) as pimg:
+      assert pimg.size == (77, 77), pimg.size
+
+  def test_fetch_subdir(self):
+    from PIL import Image
+    img = fetch("https://avatars.githubusercontent.com/u/132956020", allow_caching=False, subdir="images")
+    with Image.open(img) as pimg:
+      assert pimg.size == (77, 77), pimg.size
+    assert img.parent.name == "images"
+
+  def test_fetch_gunzip_valid(self):
+    # compare fetch(gunzip=True) to fetch(gunzip=False) plus decompressing afterwards
+    gzip_url: str = 'https://ftp.gnu.org/gnu/gzip/gzip-1.13.tar.gz'
+    fp_gz = fetch(gzip_url, gunzip=True)
+    fp_no_gz = fetch(gzip_url, gunzip=False)
+    with open(fp_gz, 'rb') as f: content_gz = f.read()
+    with open(fp_no_gz, 'rb') as f: content_no_gz = gzip.decompress(f.read())
+    assert fp_gz.stat().st_size > fp_no_gz.stat().st_size
+    assert isinstance(content_gz, bytes) and isinstance(content_no_gz, bytes)
+    assert len(content_gz) == len(content_no_gz)
+    assert content_gz == content_no_gz
+
+  def test_fetch_gunzip_invalid(self):
+    # given a non-gzipped file, fetch(gunzip=True) fails
+    no_gzip_url: str = 'https://ftp.gnu.org/gnu/gzip/gzip-1.13.zip'
+    with self.assertRaises(gzip.BadGzipFile):
+      fetch(no_gzip_url, gunzip=True)
+
+  def test_fetch_user_agent(self):
+    fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
+          allow_caching=False)
+
+  def test_fetch_half_and_full_file(self):
+    x = fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
+          headers={"Range": "bytes=0-10"}).read_bytes()
+    assert len(x) == 11, f"{len(x) != 11}"
+    x = fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
+          headers={"Range": "bytes=0-100"}).read_bytes()
+    assert len(x) == 101, f"{len(x) != 101}"
+
+  def test_fetch_sha(self):
+    self.assertRaises(Exception, fetch, "https://ftp.gnu.org/gnu/gzip/gzip-1.13.tar.gz", allow_caching=False, sha256="a")
+    fetch("https://ftp.gnu.org/gnu/gzip/gzip-1.13.tar.gz", allow_caching=False,
+          sha256="20fc818aeebae87cdbf209d35141ad9d3cf312b35a5e6be61bfcfbf9eddd212a")
+
+class TestFullyFlatten(unittest.TestCase):
+  def test_fully_flatten(self):
+    self.assertEqual(fully_flatten([[1, 3], [1, 2]]), [1, 3, 1, 2])
+    self.assertEqual(fully_flatten(((1, 3), (1, 2))), [1, 3, 1, 2])
+    self.assertEqual(fully_flatten([[[1], [3]], [[1], [2]]]), [1, 3, 1, 2])
+    self.assertEqual(fully_flatten([[[[1], 2], 3], 4]), [1, 2, 3, 4])
+    self.assertEqual(fully_flatten([[1, 2, [3, 4]], [5, 6], 7]), [1, 2, 3, 4, 5, 6, 7])
+    self.assertEqual(fully_flatten([[1, "ab"], [True, None], [3.14, [5, "b"]]]), [1, "ab", True, None, 3.14, 5, "b"])
+
+  def test_fully_flatten_numpy(self):
+    self.assertEqual(fully_flatten([np.array([])]), [])
+    self.assertEqual(fully_flatten([np.array(3)]), [3])
+    self.assertEqual(fully_flatten([np.array([3])]), [3])
+    self.assertEqual(fully_flatten([np.array([[3]])]), [3])
+    self.assertEqual(fully_flatten([np.array([1, 3]), np.array([1, 2])]), [1, 3, 1, 2])
+    self.assertEqual(fully_flatten((np.array([1, 3]), np.array([1, 2]))), [1, 3, 1, 2])
+    self.assertEqual(fully_flatten([np.array([[1], [3]]), np.array([[1], [2]])]), [1, 3, 1, 2])
+    self.assertEqual(fully_flatten([[1, "ab"], [True, None], np.array([[3.14], [6.28]])]), [1, "ab", True, None, 3.14, 6.28])
+
+class TestMemoryview(unittest.TestCase):
+  def test_from_mv_to_mv(self):
+    base = memoryview(bytearray(b"\x11\x22\x33"*40))
+    ct = from_mv(base)
+    mv = to_mv(ctypes.addressof(ct), len(base))
+    mv[0] = 2
+    assert base[0] == 2
+
+  @unittest.skip("allocates tons of memory")
+  def test_to_mv(self):
+    sizes = [
+      (16, "16 B"),
+      (64, "64 B"),
+      (256, "256 B"),
+      (1024, "1 KB"),
+      (4 * 1024, "4 KB"),
+      (16 * 1024, "16 KB"),
+      (64 * 1024, "64 KB"),
+      (256 * 1024, "256 KB"),
+      (1 * 1024 * 1024, "1 MB"),
+      (10 * 1024 * 1024, "10 MB"),
+      (200 * 1024 * 1024, "200 MB"),
+    ]
+
+    for sz, label in sizes:
+      buf = np.random.randint(0, 256, sz, dtype=np.uint8)
+      ptr = buf.ctypes.data
+
+      iters = 100_000
+      t_us = timeit.timeit(lambda: to_mv(ptr, sz), number=iters) * 1e6 / iters
+      print(f"Size {label:>9} | Time: {t_us:8.3f} µs")
+
+  def test_speed_from_mv_vs_mv_address(self):
+    x = memoryview(bytearray(1))
+
+    iters = 100000
+    fmv_us = timeit.timeit(lambda: from_mv(x), number=iters) * 1e6 / iters
+    mva_us = timeit.timeit(lambda: mv_address(x), number=iters) * 1e6 / iters
+    print(f"from_mv vs mv_address: {fmv_us:8.3f} µs vs {mva_us:8.3f} µs")
+
+class TestGetContraction(unittest.TestCase):
+  def test_contraction(self):
+    r = get_contraction((1,2,3,4), (2,3,4))
+    self.assertEqual(r, [[0, 1], [2], [3]])
+
+    r = get_contraction((2,1,3,4), (2,3,4))
+    self.assertEqual(r, [[0], [1, 2], [3]])
+
+    r = get_contraction((1,2,3,1,4), (1,2,3,4))
+    self.assertEqual(r, [[], [0, 1], [2], [3, 4]])
+
+    r = get_contraction((1,2,3,1,4,1,1), (2,3,4))
+    self.assertEqual(r, [[0, 1], [2], [3, 4, 5, 6]])
+
+    r = get_contraction((1,2,3,4), (1,2,3*4))
+    self.assertEqual(r, [[], [0, 1], [2, 3]])
+
+    r = get_contraction((1,2,3,4), (2,1,3,4))
+    self.assertEqual(r, [[0, 1], [], [2], [3]])
+
+    r = get_contraction((1,2,3,4), (1,1,2*3*4,1))
+    self.assertEqual(r, [[], [], [0,1,2,3], []])
+
+    r = get_contraction((2,1,3,4), (1,2,3,4))
+    self.assertEqual(r, [[], [0], [1, 2], [3]])
+
+    r = get_contraction((1,2,3,4), (2*3*4,1,1,1))
+    self.assertEqual(r, [[0, 1, 2, 3], [], [], []])
+
+    r = get_contraction((4,4,4,4), (16,1,16))
+    self.assertEqual(r, [[0, 1], [], [2, 3]])
+
+    r = get_contraction((1,2,3,4,1,1,1), (2,3,4))
+    self.assertEqual(r, [[0, 1], [2], [3, 4, 5, 6]])
+
+    r = get_contraction((1,2,3,4), (1,2,3,4,1))
+    self.assertEqual(r, [[], [0, 1], [2], [3], []])
+
+    r = get_contraction((14,1,384,14,1,1,1,1), (1,14,384,14))
+    self.assertEqual(r, [[], [0], [1,2], [3,4,5,6,7]])
+
+    r = get_contraction((14,1,384,1,14,1,1,1,1), (1,14,384,14))
+    self.assertEqual(r, [[], [0], [1,2], [3,4,5,6,7,8]])
+
+    r = get_contraction((512, 512), (1, 1, 512, 1, 1, 1, 1, 512))
+    self.assertEqual(r, [[], [], [0], [], [], [], [], [1]])
+
+    r = get_contraction((1,2,3,4), (1,2,6,2))
+    self.assertEqual(r, None)
+
+  def test_contraction_ones(self):
+    r = get_contraction((1,), (1,1,1))
+    self.assertEqual(r, [[], [], [0]])
+
+    r = get_contraction((1,1), (1,1,1))
+    self.assertEqual(r, [[], [], [0, 1]])
+
+    r = get_contraction((1,1,1,1), (1,))
+    self.assertEqual(r, [[0,1,2,3]])
+
+    r = get_contraction((1,1,1,1), (1,1))
+    self.assertEqual(r, [[], [0,1,2,3]])
+
+    r = get_contraction((1,1,1,1), (1,1,1))
+    self.assertEqual(r, [[], [], [0,1,2,3]])
+
+    r = get_contraction((1,1,1,1), (1,1,1,1))
+    self.assertEqual(r, [[], [], [], [0,1,2,3]])
+
+class TestGetShape(unittest.TestCase):
+  def test_get_shape(self):
+    assert get_shape(2) == ()
+    assert get_shape([]) == (0,)
+    assert get_shape([[]]) == (1, 0)
+    assert get_shape([[1, 2]]) == (1, 2)
+    assert get_shape([[1, 2], (3, 4)]) == (2, 2)
+
+  def test_inhomogeneous_shape(self):
+    with self.assertRaises(ValueError): get_shape([[], [1]])
+    with self.assertRaises(ValueError): get_shape([[1, [2]], [1]])
+
+class TestPolyN(unittest.TestCase):
+  def test_float(self):
+    np.testing.assert_allclose(polyN(1.0, [1.0, -2.0, 1.0]), 0.0)
+    np.testing.assert_allclose(polyN(2.0, [1.0, -2.0, 1.0]), 1.0)
+    np.testing.assert_allclose(polyN(3.0, [1.0, -2.0, 1.0]), 4.0)
+    np.testing.assert_allclose(polyN(4.0, [1.0, -2.0, 1.0]), 9.0)
+
+  def test_uop(self):
+    from tinygrad.dtype import dtypes
+    from tinygrad.uop.ops import UOp
+    from test.helpers import eval_uop
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 1.0), [1.0, -2.0, 1.0])), 0.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 2.0), [1.0, -2.0, 1.0])), 1.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 3.0), [1.0, -2.0, 1.0])), 4.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 4.0), [1.0, -2.0, 1.0])), 9.0)
+
+class TestTimeToStr(unittest.TestCase):
+  def test_seconds(self):           self.assertEqual("   10.01s ", time_to_str(10.01))
+  def test_boundary_sec_ms(self):   self.assertEqual("10000.00ms", time_to_str(10))
+  def test_milliseconds(self):      self.assertEqual("  500.00ms", time_to_str(0.5))
+  def test_boundary_ms_us(self):    self.assertEqual("10000.00us", time_to_str(0.01))
+  def test_microseconds(self):      self.assertEqual("  100.00us", time_to_str(0.0001))
+  def test_zero(self):              self.assertEqual("    0.00us", time_to_str(0))
+  def test_width_formatting(self):  self.assertEqual(" 10.01s ", time_to_str(10.01, w=6))
+
+class TestCStyleDivMod(unittest.TestCase):
+  def test_div_pos(self):
+    self.assertEqual(cdiv(-9, 5), -1)
+    self.assertEqual(cdiv(-4, 5), 0)
+    self.assertEqual(cdiv(0, 5), 0)
+    self.assertEqual(cdiv(4, 5), 0)
+    self.assertEqual(cdiv(9, 5), 1)
+  def test_div_neg(self):
+    self.assertEqual(cdiv(-9, -5), 1)
+    self.assertEqual(cdiv(-4, -5), 0)
+    self.assertEqual(cdiv(0, -5), 0)
+    self.assertEqual(cdiv(4, -5), 0)
+    self.assertEqual(cdiv(9, -5), -1)
+  def test_mod_pos(self):
+    self.assertEqual(cmod(-9, 5), -4)
+    self.assertEqual(cmod(-4, 5), -4)
+    self.assertEqual(cmod(0, 5), 0)
+    self.assertEqual(cmod(4, 5), 4)
+    self.assertEqual(cmod(9, 5), 4)
+  def test_mod_neg(self):
+    self.assertEqual(cmod(-9, -5), -4)
+    self.assertEqual(cmod(-4, -5), -4)
+    self.assertEqual(cmod(0, -5), 0)
+    self.assertEqual(cmod(4, -5), 4)
+    self.assertEqual(cmod(9, -5), 4)
+
+class TestGetBits(unittest.TestCase):
+  def test_low_bits(self):
+    self.assertEqual(getbits(0b11010110, 0, 3), 0b0110)
+
+  def test_high_bits(self):
+    self.assertEqual(getbits(0b11010110, 4, 7), 0b1101)
+
+  def test_middle_bits(self):
+    self.assertEqual(getbits(0b11010110, 3, 5), 0b010)
+
+  def test_full_range(self):
+    self.assertEqual(getbits(0b11010110, 0, 7), 0b11010110)
+
+  def test_single_bit(self):
+    self.assertEqual(getbits(0b100000000, 8, 8), 1)
+
+class TestArgFix(unittest.TestCase):
+  def test_none(self):
+    self.assertEqual(argfix(None), (None, ))
+    self.assertEqual(argfix(None, None), (None, None))
+  def test_positional_arguments(self):
+    self.assertEqual(argfix(1, 2, 3), (1, 2, 3))
+  def test_tuple(self):
+    self.assertEqual(argfix((1., 2., 3.)), (1., 2., 3.))
+  def test_list(self):
+    self.assertEqual(argfix([True, False]), (True, False))
+
+class TestWordWrap(unittest.TestCase):
+  def test_wrap_simple(self):
+    wrap = 10
+    st = "x"*wrap*2
+    st2 = word_wrap(st, wrap)
+    self.assertEqual(len(st2.splitlines()), 2)
+
+  def test_wrap_colored(self):
+    wrap = 10
+    st = colored("x"*wrap*2, "red")
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(len(st2.splitlines()), 2)
+
+  def test_wrap_colored_at_boundary(self):
+    wrap = 10
+    st = "x"*(wrap-2) + colored("yyy", "red")
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(ansistrip(st2), "x"*(wrap-2)+"yy\ny")
+    self.assertNotIn("\x1b[\n", st2)
+
+  def test_wrap_explicit_newline(self):
+    wrap = 10
+    st = "\n".join(["x"*wrap, "x"*wrap, "x"*wrap])
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(len(st2.splitlines()), len(st.splitlines()))
+
+    st = "\n".join(["x"*(wrap+1), "x"*wrap, "x"*wrap])
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(len(st2.splitlines()), len(st.splitlines())+1)
+
+    st = "\n".join(["x"*(wrap+1), "x"*(wrap+1), "x"*(wrap+1)])
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(len(st2.splitlines()), len(st.splitlines())+3)
+
+class TestIsNumpyNdarray(unittest.TestCase):
+  def test_ndarray(self):
+    self.assertTrue(is_numpy_ndarray(np.array([1, 2, 3])))
+  def test_ndarray_tolist(self):
+    self.assertFalse(is_numpy_ndarray(np.array([1, 2, 3]).tolist()))
+  def test_list(self):
+    self.assertFalse(is_numpy_ndarray([1, 2, 3]))
+  def test_tensor(self):
+    self.assertFalse(is_numpy_ndarray(Tensor([1, 2, 3])))
+    self.assertFalse(is_numpy_ndarray(Tensor(np.array([1, 2, 3]))))
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_indexing.py

@@ -0,0 +1,100 @@
+# test cases are modified from pytorch test_indexing.py
+
+import unittest
+
+from tinygrad import Tensor
+
+class TestIndexing(unittest.TestCase):
+  def test_single_int(self):
+    v = Tensor.randn(5, 7, 3)
+    self.assertEqual(v[4].shape, (7, 3))
+
+  def test_multiple_int(self):
+    v = Tensor.randn(5, 7, 3)
+    self.assertEqual(v[4].shape, (7, 3))
+    self.assertEqual(v[4, :, 1].shape, (7,))
+
+  def test_none(self):
+    v = Tensor.randn(5, 7, 3)
+    self.assertEqual(v[None].shape, (1, 5, 7, 3))
+    self.assertEqual(v[:, None].shape, (5, 1, 7, 3))
+    self.assertEqual(v[:, None, None].shape, (5, 1, 1, 7, 3))
+    self.assertEqual(v[..., None].shape, (5, 7, 3, 1))
+
+  def test_int_indices(self):
+    v = Tensor.randn(5, 7, 3)
+    self.assertEqual(v[[0, 4, 2]].shape, (3, 7, 3))
+    self.assertEqual(v[:, [0, 4, 2]].shape, (5, 3, 3))
+    self.assertEqual(v[:, [[0, 1], [4, 3]]].shape, (5, 2, 2, 3))
+
+  def test_index_src_datatype(self):
+    src = Tensor.ones(3, 2, 4)
+    # test index
+    res = src[[0, 2, 1], :, :]
+    self.assertEqual(res.shape, src.shape)
+
+  def test_empty_slice(self):
+    x = Tensor.randn(2, 3, 4, 5)
+    y = x[:, :, :, 1]
+    z = y[:, 1:1, :]
+    self.assertEqual((2, 0, 4), z.shape)
+
+  def test_invalid_index(self):
+    x = Tensor.arange(0, 16).reshape(4, 4)
+    self.assertRaises(TypeError, lambda: x["0":"1"])
+
+  def test_out_of_bound_index(self):
+    x = Tensor.arange(0, 100).reshape(2, 5, 10)
+    self.assertRaises(IndexError, lambda: x[0, 5])
+    self.assertRaises(IndexError, lambda: x[4, 5])
+    self.assertRaises(IndexError, lambda: x[0, 1, 15])
+    self.assertRaises(IndexError, lambda: x[:, :, 12])
+
+class TestNumpy(unittest.TestCase):
+  def test_index_no_floats(self):
+    a = Tensor([[[5.]]])
+
+    self.assertRaises(IndexError, lambda: a[0.0])
+    self.assertRaises(IndexError, lambda: a[0, 0.0])
+    self.assertRaises(IndexError, lambda: a[0.0, 0])
+    self.assertRaises(IndexError, lambda: a[0.0, :])
+    self.assertRaises(IndexError, lambda: a[:, 0.0])
+    self.assertRaises(IndexError, lambda: a[:, 0.0, :])
+    self.assertRaises(IndexError, lambda: a[0.0, :, :])
+    self.assertRaises(IndexError, lambda: a[0, 0, 0.0])
+    self.assertRaises(IndexError, lambda: a[0.0, 0, 0])
+    self.assertRaises(IndexError, lambda: a[0, 0.0, 0])
+    self.assertRaises(IndexError, lambda: a[-1.4])
+    self.assertRaises(IndexError, lambda: a[0, -1.4])
+    self.assertRaises(IndexError, lambda: a[-1.4, 0])
+    self.assertRaises(IndexError, lambda: a[-1.4, :])
+    self.assertRaises(IndexError, lambda: a[:, -1.4])
+    self.assertRaises(IndexError, lambda: a[:, -1.4, :])
+    self.assertRaises(IndexError, lambda: a[-1.4, :, :])
+    self.assertRaises(IndexError, lambda: a[0, 0, -1.4])
+    self.assertRaises(IndexError, lambda: a[-1.4, 0, 0])
+    self.assertRaises(IndexError, lambda: a[0, -1.4, 0])
+    # these two trigger slice internal type verification first
+    self.assertRaises(TypeError, lambda: a[0.0:, 0.0])
+    self.assertRaises(TypeError, lambda: a[0.0:, 0.0,:])
+
+  def test_none_index(self):
+    # `None` index adds newaxis
+    a = Tensor([1, 2, 3])
+    self.assertEqual(a[None].ndim, a.ndim+1)
+
+  def test_everything_returns_views(self):
+    # Before `...` would return a itself.
+    a = Tensor([5])
+
+    self.assertIs(a, a[()])
+    self.assertIs(a, a[...])
+    self.assertIs(a, a[:])
+
+  def test_broaderrors_indexing(self):
+    a = Tensor.zeros(5, 5)
+    self.assertRaises(IndexError, a.__getitem__, ([0, 1], [0, 1, 2]))
+    self.assertRaises(IndexError, a.contiguous().__setitem__, ([0, 1], [0, 1, 2]), 0)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_linearizer_failures.py

@@ -0,0 +1,25 @@
+# ruff: noqa: E501
+import unittest
+from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
+from tinygrad.dtype import dtypes
+from tinygrad.device import Device
+from tinygrad.codegen import to_program
+
+class TestLinearizerFailures(unittest.TestCase):
+  def test_fail_1(self):
+    c0 = UOp.param(0, dtypes.float.ptr(64))
+    c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.LOOP)
+    c3 = ((c1*UOp.const(dtypes.weakint, 32))+c2)
+    c4 = UOp.param(1, dtypes.float.ptr(163840))
+    c5 = UOp.range(UOp.const(dtypes.weakint, 2560), 0, AxisType.REDUCE)
+    c6 = c4.index(((((((c5//UOp.const(dtypes.weakint, 8))%UOp.const(dtypes.weakint, 8))*UOp.const(dtypes.weakint, 8))+(c5%UOp.const(dtypes.weakint, 8)))+(((c2*UOp.const(dtypes.weakint, 40))+(c5//UOp.const(dtypes.weakint, 64)))*UOp.const(dtypes.weakint, 64)))+(c1*UOp.const(dtypes.weakint, 81920))))
+    c7 = UOp.param(2, dtypes.float.ptr(64))
+    c8 = c7.index(c3)
+    c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
+    c10 = c0.index(c3).store(c9).end(c1, c2)
+    ast = c10.sink(arg=KernelInfo())
+    to_program(ast, renderer=Device[Device.DEFAULT].renderer)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_linearizer_rewrite.py

@@ -0,0 +1,45 @@
+import unittest
+from tinygrad import Tensor, Context, Device
+from tinygrad.codegen import to_program
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.uop.ops import KernelInfo
+
+class TestLinearizerRewrite(unittest.TestCase):
+  def test_reduction(self):
+    t = Tensor.ones((64,64), device="NULL").contiguous().realize()
+    out = (t*2).sum(axis=1)
+    with Context(SPLIT_REDUCEOP=0):
+      si = out.schedule_linear().src[-1]
+      opts_to_apply = []
+      opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
+      opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4))
+      ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
+      prg = to_program(ast, Device["CPU"].renderer)
+      print(prg.src[3].arg)
+
+  def test_arange(self):
+    out = Tensor.arange(32, device="NULL")
+    with Context(SPLIT_REDUCEOP=0):
+      si = out.schedule_linear().src[-1]
+      opts_to_apply = []
+      opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
+      ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
+      prg = to_program(ast, Device["CPU"].renderer)
+      print(prg.src[3].arg)
+
+  def test_kernel_info(self):
+    out = Tensor.arange(4, device="NULL")
+    si = out.schedule_linear().src[-1]
+
+    ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=()))
+    prg = to_program(ast, Device["CPU"].renderer)
+    assert prg.src[0].arg.applied_opts == (), f"expected no opts, got {prg}"
+
+    #prg = to_program(ast.replace(arg=KernelInfo()), Device["CPU"].renderer)
+    #assert prg.src[0].arg.applied_opts != (), f"expected opts to apply, got {prg.src[0].arg.applied_opts}"
+
+    prg = to_program(ast.replace(arg=KernelInfo(name="custom")), Device["CPU"].renderer)
+    self.assertEqual(prg.arg.name, "custom")
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_llm_server.py

@@ -0,0 +1,206 @@
+import unittest, threading, time
+from unittest.mock import Mock
+
+class TestLLMServer(unittest.TestCase):
+  """Integration tests using the real OpenAI client."""
+
+  @classmethod
+  def setUpClass(cls):
+    cls.mock_tok = Mock()
+    cls.mock_tok.role = Mock(return_value=[100, 101])
+    cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
+    cls.mock_tok.decode = Mock(return_value="Hello")
+    cls.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: "Hello" if tid is not None else "")
+    cls.mock_tok.end_turn = Mock(return_value=[998])
+    cls.mock_tok.prefix = Mock(return_value=[1])
+    cls.mock_tok.preset = "llama3"
+    cls.mock_tok.bos_id = 1
+    cls.mock_tok.eos_id = 999
+    cls.mock_tok.eot_id = None
+    cls.mock_tok.is_end = Mock(side_effect=lambda tid: tid in (999,))
+
+    cls.mock_model = Mock()
+    cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
+    cls.mock_model.get_start_pos = Mock(return_value=0)
+
+    from tinygrad.llm.cli import LLMServer
+
+    cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "test-model", cls.mock_tok)
+    cls.port = cls.server.server_address[1]
+    cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
+    cls.server_thread.start()
+    time.sleep(0.1)
+
+    from openai import OpenAI
+    cls.client = OpenAI(base_url=f"http://127.0.0.1:{cls.port}/v1", api_key="test")
+
+  @classmethod
+  def tearDownClass(cls):
+    cls.server.shutdown()
+    cls.server.server_close()
+
+  def test_chat_completion_stream(self):
+    stream = self.client.chat.completions.create(
+      model="test",
+      messages=[{"role": "user", "content": "Hello"}],
+      stream=True
+    )
+
+    chunks = list(stream)
+    self.assertGreater(len(chunks), 0)
+    self.assertEqual(chunks[0].choices[0].delta.role, "assistant")
+    self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
+
+  def test_openai_response_structure(self):
+    stream = self.client.chat.completions.create(
+      model="test-model",
+      messages=[{"role": "user", "content": "Test"}],
+      stream=True
+    )
+
+    for chunk in stream:
+      self.assertTrue(chunk.id.startswith("chatcmpl-"))
+      self.assertEqual(chunk.object, "chat.completion.chunk")
+      self.assertIsNotNone(chunk.choices)
+      self.assertIsNotNone(chunk.created)
+      self.assertIsInstance(chunk.created, int)
+      self.assertEqual(chunk.model, "test-model")
+
+  def test_stream_with_usage(self):
+    stream = self.client.chat.completions.create(
+      model="test",
+      messages=[{"role": "user", "content": "Hello"}],
+      stream=True,
+      stream_options={"include_usage": True}
+    )
+
+    chunks = list(stream)
+    last_chunk = chunks[-1]
+
+    self.assertIsNotNone(last_chunk.usage)
+    self.assertIsNotNone(last_chunk.usage.prompt_tokens)
+    self.assertIsNotNone(last_chunk.usage.completion_tokens)
+    self.assertIsNotNone(last_chunk.usage.total_tokens)
+
+  def test_multi_turn_conversation(self):
+    stream = self.client.chat.completions.create(
+      model="test",
+      messages=[
+        {"role": "system", "content": "You are helpful."},
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Hi!"},
+        {"role": "user", "content": "How are you?"}
+      ],
+      stream=True
+    )
+
+    chunks = list(stream)
+    self.assertGreater(len(chunks), 0)
+    self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
+
+  def test_content_is_streamed(self):
+    stream = self.client.chat.completions.create(
+      model="test",
+      messages=[{"role": "user", "content": "Hello"}],
+      stream=True
+    )
+
+    contents = []
+    for chunk in stream:
+      if chunk.choices and chunk.choices[0].delta.content:
+        contents.append(chunk.choices[0].delta.content)
+
+    self.assertGreater(len(contents), 0)
+
+  def test_non_streaming(self):
+    resp = self.client.chat.completions.create(
+      model="test-model",
+      messages=[{"role": "user", "content": "Hello"}],
+      stream=False
+    )
+
+    self.assertTrue(resp.id.startswith("chatcmpl-"))
+    self.assertEqual(resp.object, "chat.completion")
+    self.assertEqual(resp.model, "test-model")
+    self.assertIsNotNone(resp.created)
+    self.assertEqual(len(resp.choices), 1)
+    self.assertEqual(resp.choices[0].message.role, "assistant")
+    self.assertIsNotNone(resp.choices[0].message.content)
+    self.assertEqual(resp.choices[0].finish_reason, "stop")
+    self.assertIsNotNone(resp.usage)
+    self.assertIsNotNone(resp.usage.prompt_tokens)
+    self.assertIsNotNone(resp.usage.completion_tokens)
+
+  def test_max_tokens_streaming(self):
+    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 302, 303, 999]))
+    stream = self.client.chat.completions.create(
+      model="test", messages=[{"role": "user", "content": "Hello"}], stream=True, max_tokens=2
+    )
+    chunks = list(stream)
+    content_chunks = [c for c in chunks if c.choices and c.choices[0].delta.content]
+    self.assertEqual(len(content_chunks), 2)
+    self.assertEqual(chunks[-1].choices[0].finish_reason, "length")
+
+  def test_max_tokens_non_streaming(self):
+    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 302, 303, 999]))
+    resp = self.client.chat.completions.create(
+      model="test", messages=[{"role": "user", "content": "Hello"}], stream=False, max_tokens=2
+    )
+    self.assertEqual(resp.choices[0].finish_reason, "length")
+    self.assertEqual(resp.usage.completion_tokens, 2)
+
+  def test_assistant_prefill(self):
+    """Last assistant message should be treated as prefill (not a completed turn)."""
+    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 999]))
+    captured_ids = []
+    orig_generate = self.mock_model.generate.side_effect
+    def capture_generate(ids, **kwargs):
+      captured_ids.extend(ids)
+      return orig_generate(ids, **kwargs)
+    self.mock_model.generate = Mock(side_effect=capture_generate)
+
+    resp = self.client.chat.completions.create(
+      model="test", messages=[
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Sure"}
+      ], stream=False
+    )
+    # prefill tokens should be in ids: role("assistant") + encode("Sure") but NO end_turn after it
+    # and NO extra role("assistant") appended
+    role_tokens = self.mock_tok.role.call_args_list
+    # last role() call should be for "assistant" (the prefill message), not an extra one
+    self.assertEqual(role_tokens[-1], unittest.mock.call("assistant"))
+    # end_turn should be called once less than role() — the prefill assistant msg doesn't get end_turn
+    self.assertEqual(self.mock_tok.end_turn.call_count, self.mock_tok.role.call_count - 1)
+    self.assertIsNotNone(resp.choices[0].message.content)
+
+  def test_assistant_prefill_not_last(self):
+    """Assistant message that's NOT last should be a normal completed turn."""
+    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 999]))
+    self.mock_tok.role.reset_mock()
+    self.mock_tok.end_turn.reset_mock()
+    self.client.chat.completions.create(
+      model="test", messages=[
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Sure"},
+        {"role": "user", "content": "Continue"}
+      ], stream=False
+    )
+    # all messages get end_turn, plus an extra role("assistant") at the end
+    # roles: user, assistant, user, assistant(generation prompt) = 4 role calls
+    # end_turns: user, assistant, user = 3 end_turn calls (one per message)
+    self.assertEqual(self.mock_tok.end_turn.call_count, 3)
+    self.assertEqual(self.mock_tok.role.call_count, 4)
+
+  def test_models_endpoint(self):
+    import requests as req
+    resp = req.get(f"http://127.0.0.1:{self.port}/v1/models")
+    self.assertEqual(resp.status_code, 200)
+    data = resp.json()
+    self.assertEqual(data["object"], "list")
+    self.assertEqual(len(data["data"]), 1)
+    self.assertEqual(data["data"][0]["id"], "test-model")
+    self.assertEqual(data["data"][0]["object"], "model")
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_llm_tokenizer.py

@@ -0,0 +1,72 @@
+import unittest, base64, functools, sys
+from tinygrad.llm.cli import SimpleTokenizer
+from tinygrad.helpers import fetch
+
+@unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
+class TestLLMTokenizer(unittest.TestCase):
+  @functools.cached_property
+  def llama_tok(self):
+    # from https://github.com/tinygrad/tinygrad/blob/e0106b6b257ebc003eb3694144e3e198f7d8cc37/examples/llama3.py#L14
+    model_file = fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model")
+    with open(model_file, "rt") as fd:
+      str_vocab = [line.split(maxsplit=1) for line in fd.read().splitlines() if line]
+
+      # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
+      bs = [*range(33, 127), *range(161, 173), *range(174, 256)]  # bytes that map to themselves
+      _byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
+      _byte_encoder = {v:k for k,v in _byte_decoder.items()}
+      normal_tokens = {''.join([_byte_encoder[x] for x in base64.b64decode(stok)]): int(srank) for stok, srank in str_vocab}
+
+    special_tokens = [
+      "<|begin_of_text|>",
+      "<|end_of_text|>",
+      "<|reserved_special_token_0|>",
+      "<|reserved_special_token_1|>",
+      "<|reserved_special_token_2|>",
+      "<|reserved_special_token_3|>",
+      "<|start_header_id|>",
+      "<|end_header_id|>",
+      "<|reserved_special_token_4|>",
+      "<|eot_id|>",
+    ] + [ f"<|reserved_special_token_{i}|>" for i in range(5, 256 - 5) ]
+    return SimpleTokenizer(normal_tokens, {token: len(normal_tokens) + i for i, token in enumerate(special_tokens)})
+
+  def _test_coding(self, tok: SimpleTokenizer, text: str, expected_tokens: list[int]):
+    self.assertEqual(tok.encode(text), expected_tokens)
+    self.assertEqual(tok.decode(expected_tokens), text)
+
+  # NOTE: the correct tokenization for this can only be found by looking up the text chunk in the vocab, not by applying merges
+  def test_llama_early_tokenize(self): self._test_coding(self.llama_tok, " например", [ 111797 ])
+
+  def test_llama_basic(self): self._test_coding(self.llama_tok, "hello world", [ 15339, 1917 ])
+  def test_llama_control_char(self): self._test_coding(self.llama_tok, " \x850", [ 220, 116360, 15 ])
+  def test_llama_bytes(self): self._test_coding(self.llama_tok, " \xec\x8b\xa4\xed", [ 1717, 105, 116174, 82638, 2483 ])
+  def test_llama_special1(self): self._test_coding(self.llama_tok, "hello <|end_of_text|>", [ 15339, 220, 128001 ])
+  def test_llama_special2(self): self._test_coding(self.llama_tok, "<|start_header_id|>user<|end_header_id|>\n\n", [ 128006, 882, 128007, 271 ])
+  def test_llama_repeat(self): self._test_coding(self.llama_tok, "00000000000000000", [ 931, 931, 931, 931, 931, 410 ])
+  def test_llama_pat(self): self._test_coding(self.llama_tok, "today\n  \n", [ 31213, 14211 ])
+
+  def test_tekken_from_gguf_kv(self):
+    kv = {
+      "tokenizer.ggml.tokens": ["<unk>", "<s>", "</s>", "[INST]", "[/INST]", "hello"],
+      "tokenizer.ggml.token_type": [3, 3, 3, 3, 3, 1],
+      "tokenizer.ggml.pre": "tekken",
+      "tokenizer.ggml.eos_token_id": 2,
+    }
+    tok = SimpleTokenizer.from_gguf_kv(kv)
+    self.assertEqual(tok.role("user"), [3])
+    self.assertEqual(tok.encode("hello"), [5])
+    self.assertEqual(tok.end_turn(), [4])
+    self.assertEqual(tok.role("assistant"), [])
+
+  def test_stream_decoder(self):
+    """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""
+    bs = [*range(33, 127), *range(161, 173), *range(174, 256)]
+    be = {b: chr(b) for b in bs} | {b: chr(256+i) for i,b in enumerate(b for b in range(256) if b not in bs)}
+    token_bytes = {25677: b'\x20\xf0\x9f\x98', 138: b'\x8a'}  # ' ' + 3/4 emoji | 1/4 emoji (qwen3.5)
+    tok = SimpleTokenizer({"".join(be[b] for b in v): k for k, v in token_bytes.items()}, {})
+    dec = tok.stream_decoder()
+    self.assertEqual(dec(25677) + dec(138) + dec(), " 😊")
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_memory_planner.py

@@ -0,0 +1,247 @@
+import unittest
+from tinygrad import dtypes
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.schedule.memory import memory_plan_rewrite
+
+global_map = {}
+held_bufs: set[UOp] = set()
+def b(i, base=None, offset=0, pin=False, size=16):
+  global global_map
+  if i in global_map: return global_map[i]
+  if base is not None:
+    global_map[i] = global_map[base]
+    return global_map[i]
+  global_map[i] = UOp.new_buffer("NULL", size, dtypes.int8)
+  if pin: held_bufs.add(global_map[i])
+  return global_map[i]
+
+def _make_linear(buffer_lists, copies=None):
+  copy_pairs = {frozenset((id(dst), id(src))) for dst, src in copies} if copies else set()
+  calls = []
+  for bufs in buffer_lists:
+    is_copy = len(bufs) == 2 and frozenset((id(bufs[0]), id(bufs[1]))) in copy_pairs
+    calls.append(UOp(Ops.CALL, dtypes.void, (UOp(Ops.COPY if is_copy else Ops.SINK), *bufs)))
+  return UOp(Ops.LINEAR, src=tuple(calls))
+
+def _get_arena(buf, linear, result):
+  for orig_si, new_si in zip(linear.src, result.src):
+    for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
+      if orig is buf and new.op is Ops.SLICE: return new.src[0]
+  return None
+
+def check_assign(buffer_lists, copies=None):
+  linear = _make_linear(buffer_lists, copies)
+  result = memory_plan_rewrite(linear, held_bufs)
+
+  # build mapping: original buf -> (arena, offset_bytes, nbytes) from the result
+  replace_map: dict[int, tuple[UOp, int, int]] = {}
+  for orig_si, new_si in zip(linear.src, result.src):
+    for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
+      if new.op is Ops.SLICE and id(orig) not in replace_map:
+        replace_map[id(orig)] = (new.src[0], new.src[1].arg * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)
+
+  # verify pinned buffers are not planned
+  for buf in held_bufs:
+    assert id(buf) not in replace_map, "pinned buffer was planned"
+
+  # compute lifetimes
+  first_appearance, last_appearance = {}, {}
+  for i, bufs in enumerate(buffer_lists):
+    for buf in bufs:
+      if buf in held_bufs: continue
+      if id(buf) not in first_appearance: first_appearance[id(buf)] = i
+      last_appearance[id(buf)] = i
+
+  # verify non-overlapping: no two live buffers share the same arena region
+  taken_parts: set[tuple[int, int, int, int]] = set()  # (id(arena), offset, nbytes, id(buf))
+  for i, bufs in enumerate(buffer_lists):
+    for buf in bufs:
+      if buf in held_bufs or id(buf) not in replace_map: continue
+      arena, off, nb = replace_map[id(buf)]
+      for part in taken_parts:
+        assert id(buf) == part[3] or part[0] != id(arena) or part[1] + part[2] <= off or part[1] >= off + nb, \
+          f"overlap at step {i}: [{off}, {off+nb}) conflicts with [{part[1]}, {part[1]+part[2]})"
+      if first_appearance.get(id(buf)) == i: taken_parts.add((id(arena), off, nb, id(buf)))
+      if last_appearance.get(id(buf)) == i: taken_parts.discard((id(arena), off, nb, id(buf)))
+
+class TestMemoryPlanner(unittest.TestCase):
+  def setUp(self):
+    global global_map
+    held_bufs.clear()
+    global_map = {}
+
+  def test_simple_buffer(self):
+    bs = [
+      [b(0), b(1), b(2)],
+      [b(1), b(2), b(3)],
+      [b(4), b(3)],
+      [b(5), b(2)],
+    ]
+    check_assign(bs)
+
+  def test_simple_pinned(self):
+    bs = [
+      [b(0, pin=True), b(1), b(2, pin=True)],
+      [b(1), b(2), b(3)],
+      [b(4), b(3)],
+      [b(5), b(2)],
+    ]
+    check_assign(bs)
+
+  def test_all_pinned(self):
+    bs = [
+      [b(0, pin=True), b(1, pin=True)],
+      [b(1), b(2, pin=True)],
+      [b(4, pin=True), b(3, pin=True)],
+    ]
+    check_assign(bs)
+
+  def test_simple_buffer_offset(self):
+    bs = [
+      [b(0, pin=True), b(1, base=0, offset=1, size=8), b(2)],
+      [b(1), b(2), b(3, base=0, offset=1, size=8)],
+      [b(4), b(3)],
+    ]
+    check_assign(bs)
+
+  def test_buffer_offset(self):
+    bs = [
+      [b(0, pin=True), b(1, base=0, offset=1, size=8), b(2)],
+      [b(1), b(2), b(3, base=0, offset=1, size=8)],
+      [b(4), b(3)],
+      [b(5, base=2, offset=2, size=8), b(3)],
+      [b(6), b(5), b(0)],
+      [b(7), b(8, pin=True)],
+      [b(8), b(9, base=2, offset=2, size=8)],
+      [b(9), b(3), b(5)],
+    ]
+    check_assign(bs)
+
+  def test_buffer_offset2(self):
+    bs = [
+      [b(0, pin=True), b(1), b(2)],
+      [b(1), b(2), b(3)],
+      [b(4), b(3)],
+      [b(5), b(3)],
+      [b(6), b(5), b(0)],
+      [b(7), b(8, pin=True)],
+      [b(8), b(9)],
+      [b(9), b(3), b(5)],
+      [b(11), b(0)],
+      [b(11), b(10), b(5)],
+      [b(12), b(11), b(0)],
+      [b(6), b(12), b(7)],
+      [b(13), b(6), b(11)],
+    ]
+    check_assign(bs)
+
+  def test_all_offsets_of_one(self):
+    bs = [
+      [b(0, pin=True), b(1)],
+      [b(3, base=1, offset=0, size=8), b(2, base=0, offset=0, size=8)],
+      [b(5, base=1, offset=8, size=8), b(4, base=0, offset=8, size=8)],
+      [b(7, base=1, offset=4, size=8), b(6, base=0, offset=4, size=8)],
+
+      [b(4), b(5), b(2)],
+      [b(3), b(7)],
+      [b(10), b(6), b(7)],
+      [b(11), b(3), b(2)],
+      [b(12), b(5), b(4), b(3), b(2)],
+      [b(13), b(6), b(12), b(7)],
+    ]
+    check_assign(bs)
+
+  def test_very_small_buffers(self):
+    bs = [
+      [b(0, pin=True), b(1, size=32)],
+      [b(3, size=4), b(4, size=6)],
+    ]
+    check_assign(bs)
+
+  def test_very_big_buffers(self):
+    bs = [
+      [b(0, pin=True), b(1, size=34359738368000)],
+      [b(3, size=1 << 128), b(4, size=1 << 64)],
+    ]
+    check_assign(bs)
+
+  def test_copy_bufs_separate_from_compute(self):
+    bs = [
+      [b(0), b(1)],
+      [b(1), b(2)],
+      [b(3), b(2)],
+    ]
+    linear = _make_linear(bs, copies=[(b(1), b(0))])
+    result = memory_plan_rewrite(linear)
+    r1_arena, r2_arena = _get_arena(b(1), linear, result), _get_arena(b(2), linear, result)
+    assert r1_arena is not None and r2_arena is not None
+    assert r1_arena is not r2_arena
+
+  def test_copy_bufs_reuse_among_copies(self):
+    bs = [
+      [b(0), b(1)],
+      [b(2), b(1)],
+      [b(3), b(2)],
+    ]
+    linear = _make_linear(bs, copies=[(b(1), b(0)), (b(2), b(1))])
+    result = memory_plan_rewrite(linear)
+    r1_arena, r2_arena = _get_arena(b(1), linear, result), _get_arena(b(2), linear, result)
+    assert r1_arena is not None and r2_arena is not None
+    assert r1_arena is r2_arena
+
+  def test_compute_bufs_reuse_among_compute(self):
+    bs = [
+      [b(0), b(1)],
+      [b(2), b(1)],
+      [b(3), b(2)],
+      [b(4), b(3)],
+    ]
+    linear = _make_linear(bs, copies=[(b(1), b(0))])
+    result = memory_plan_rewrite(linear)
+    r2_arena, r3_arena = _get_arena(b(2), linear, result), _get_arena(b(3), linear, result)
+    assert r2_arena is not None and r3_arena is not None
+    assert r2_arena is r3_arena
+
+  def test_copy_and_compute_no_cross_reuse(self):
+    bs = [
+      [b(0), b(1)],
+      [b(2), b(1)],
+      [b(3), b(2)],
+    ]
+    linear = _make_linear(bs, copies=[(b(2), b(1))])
+    result = memory_plan_rewrite(linear)
+    r0_arena, r2_arena = _get_arena(b(0), linear, result), _get_arena(b(2), linear, result)
+    assert r0_arena is not None and r2_arena is not None
+    assert r0_arena is not r2_arena
+
+  def test_multiple_copy_bufs_with_offsets(self):
+    bs = [
+      [b(0, pin=True), b(1), b(2)],
+      [b(3, base=0, offset=1, size=8), b(1), b(2)],
+      [b(4), b(3)],
+      [b(5), b(4)],
+    ]
+    check_assign(bs, copies=[(b(1), b(0)), (b(2), b(0))])
+
+  def test_copy_bufs_pinned_mixed(self):
+    bs = [
+      [b(0, pin=True), b(1), b(2)],
+      [b(1), b(3), b(2)],
+      [b(4), b(3)],
+      [b(5), b(4), b(0)],
+    ]
+    check_assign(bs, copies=[(b(1), b(0)), (b(3), b(1))])
+
+  def test_deferred_copy_frees_chain(self):
+    bs = []
+    copies = []
+    for i in range(6):
+      copy_buf, compute_buf = b(i * 2 + 1), b(i * 2 + 2)
+      bs.append([copy_buf, b(0, pin=True)])
+      bs.append([compute_buf, copy_buf])
+      copies.append((copy_buf, b(0, pin=True)))
+    bs.append([b(100, pin=True)])
+    check_assign(bs, copies=copies)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/null/test_method_cache.py

@@ -0,0 +1,48 @@
+import unittest
+from tinygrad import Tensor, Device, Variable
+from examples.gpt2 import Transformer
+from tinygrad.nn.state import get_state_dict
+
+class TestMethodCache(unittest.TestCase):
+  def setUp(self):
+    self.backup_compiler = Device[Device.DEFAULT].compiler.compile_cached
+  def tearDown(self):
+    Device[Device.DEFAULT].compiler.compile_cached = self.backup_compiler
+
+  def test_simple_methodcache(self):
+    a = Tensor([1])
+    b = Tensor([2])
+    c = Tensor([3])
+    d = Tensor([4])
+    (a+b).realize()
+    Device[Device.DEFAULT].compiler.compile_cached = None
+    (c+d).realize()
+
+  def test_nested_methodcache(self):
+    a,b,c,d = Tensor([1]), Tensor([2]), Tensor([3]), Tensor([4])
+    ((a+b)+(a+b)).realize()
+    Device[Device.DEFAULT].compiler.compile_cached = None
+    ((c+d)+(c+d)).realize()
+
+  def test_nested_methodcache_swap(self):
+    a,b,c,d = Tensor([1]), Tensor([2]), Tensor([3]), Tensor([4])
+    ((a+b)+(c+d)).realize()
+    Device[Device.DEFAULT].compiler.compile_cached = None
+    ((c+d)+(a+b)).realize()
+
+  @unittest.skip("incorrect use of transformer")
+  def test_small_transformer(self):
+    args_tiny = {"dim": 16, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 10}
+    model = Transformer(**args_tiny)
+    for v in get_state_dict(model).values(): v.assign(Tensor.empty(*v.shape, dtype=v.dtype).realize())
+    # NOTE: you have to do this twice due to the k-v cache
+    for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize()
+    for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize()
+    Device[Device.DEFAULT].compiler.compile_cached = None
+    for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize()
+
+if __name__ == '__main__':
+  unittest.main()
+
+
+

tinygrad_repo/test/null/test_microbenchmarks.py

@@ -0,0 +1,100 @@
+import unittest, time
+from tinygrad import dtypes, Tensor, UOp, getenv
+from tinygrad.helpers import Profiling
+
+PYPROFILE = getenv("PYPROFILE")
+class TestBench(unittest.TestCase):
+  @staticmethod
+  def setUpClass():
+    # no fixed cost
+    Tensor.empty(10,10)
+    Tensor.randn(10,10)
+
+  def start_time(self): self.st = time.perf_counter()
+  def setUp(self):
+    # it's about 1 ms per 1k UOps on M3
+    if PYPROFILE:
+      self.prof = Profiling()
+      self.prof.__enter__()
+    else:
+      self.prof = None
+    self.N = 10000
+    self.start_time()
+
+  def tearDown(self):
+    et = (time.perf_counter() - self.st)
+    if self.prof is not None: self.prof.__exit__()
+    print(f"{self._testMethodName:30s} {et*1e6/self.N:.2f} us")
+
+  def test_uop_instant_creation(self):
+    for i in range(self.N): UOp.const(dtypes.int, 100+i)
+
+  def test_uop_list_creation(self):
+    [UOp.const(dtypes.int, 100+i) for i in range(self.N)]
+
+  def test_uop_add_2n(self):
+    a = UOp.const(dtypes.int, 2)
+    for _ in range(self.N): a = a + a
+
+  def test_uop_toposort(self):
+    a = UOp.const(dtypes.int, 0)
+    for i in range(self.N): a = a + UOp.const(dtypes.int, 100+i)
+    self.start_time()
+    self.assertEqual(len(a.toposort()), 2*self.N+1)
+
+  def test_uop_toposort_2n(self):
+    a = UOp.const(dtypes.int, 0)
+    for _ in range(self.N): a = a + a
+    self.start_time()
+    self.assertEqual(len(a.toposort()), self.N+1)
+
+  def test_uop_simplify(self):
+    a = UOp.const(dtypes.int, 2)
+    for _ in range(self.N): (a+a).simplify()
+
+  def test_uop_simplify_complex(self):
+    self.N //= 10 # this test is slow
+    x = UOp.variable("x", 0, 10)
+    y = UOp.variable("y", 0, 10)
+    expr = (x*2)+5+(x*4)+(y*2)+y
+    for _ in range(self.N): expr.simplify()
+
+  def test_uop_simplify_div(self):
+    self.N //= 10 # this test is slow
+    x = UOp.variable("x", 0, 10)
+    y = UOp.variable("y", 0, 10)
+    z = UOp.variable("z", 0, 10)
+    expr = (x*4+y*8)//(z*2)
+    for _ in range(self.N): expr.simplify()
+
+  def test_uop_chain_free(self):
+    a = UOp.const(dtypes.int, 2)
+    for _ in range(self.N): a = a + a
+    self.start_time()
+    del a
+
+  def test_tensor_zeros(self):
+    self.N //= 10 # this test is slow
+    for _ in range(self.N): Tensor.zeros(10, 10)
+
+  def test_tensor_add(self):
+    self.N //= 10 # this test is slow
+    a = Tensor.zeros(10, 10)
+    b = Tensor.zeros(10, 10)
+    for _ in range(self.N): a+b
+
+  def test_tensor_empty(self):
+    self.N //= 10 # this test is slow
+    for _ in range(self.N): Tensor.empty(10, 10)
+
+  def test_tensor_rand(self):
+    self.N //= 100 # this test is very slow
+    for _ in range(self.N): Tensor.rand(10, 10)
+
+  def test_tensor_randn(self):
+    self.N //= 100 # this test is very slow
+    for _ in range(self.N): Tensor.randn(10, 10)
+
+if __name__ == '__main__':
+  unittest.main()
+

tinygrad_repo/test/null/test_mnist_dataset.py

@@ -0,0 +1,14 @@
+import unittest
+from tinygrad.helpers import GlobalCounters
+from tinygrad.nn.datasets import mnist
+
+class TestDataset(unittest.TestCase):
+  def test_dataset_is_realized(self):
+    X_train, _, _, _ = mnist()
+    X_train[0].contiguous().realize()
+    GlobalCounters.reset()
+    X_train[0].contiguous().realize()
+    self.assertLessEqual(GlobalCounters.kernel_count, 1)  # 0 if SLICE (zero-copy), 1 otherwise
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_multitensor.py

@@ -0,0 +1,216 @@
+import gc, unittest
+from tinygrad import Tensor, GlobalCounters, dtypes
+from tinygrad.engine.jit import TinyJit
+from tinygrad.helpers import Context
+
+class TestMultiRamUsage(unittest.TestCase):
+  def setUp(self):
+    gc.collect()
+    self.baseline = GlobalCounters.mem_used
+    self.baseline_per_device = dict(GlobalCounters.mem_used_per_device)
+    self.N = 100
+  def assertUsed(self, amt, strict=True):
+    gc.collect()
+    used = GlobalCounters.mem_used - self.baseline
+    print(f"used {used} bytes")
+    if strict: self.assertEqual(used, amt)
+    else: self.assertLessEqual(used, amt)
+  def assertDeviceUsed(self, expected:dict[str, int]):
+    gc.collect()
+    for dev, amt in expected.items():
+      used = GlobalCounters.mem_used_per_device[dev] - self.baseline_per_device.get(dev, 0)
+      self.assertEqual(used, amt, f"device {dev}: expected {amt} bytes used, got {used}")
+
+  def test_zeros(self):
+    _ = Tensor.zeros(self.N, self.N).contiguous().realize()
+    self.assertUsed(self.N*self.N*4)
+
+  def test_zeros_del(self):
+    _ = Tensor.zeros(self.N, self.N).contiguous().realize()
+    del _
+    self.assertUsed(0)
+
+  def test_zeros_copy(self):
+    devices_2 = ("NULL:1", "NULL:2")
+    _ = Tensor.zeros(self.N, self.N).contiguous().to(devices_2).realize()
+    # NOTE: the first one on the DEFAULT device should be freed
+    self.assertUsed(self.N*self.N*4*2)
+
+  def test_zeros_shard(self, devices=("NULL:1", "NULL:2")):
+    _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices, axis=0).realize()
+    self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
+  def test_zeros_shard_self(self): self.test_zeros_shard(("NULL:0", "NULL:1"))
+
+  def test_zeros_contiguous_shard(self):
+    devices_2 = ("NULL:1", "NULL:2")
+    _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).contiguous().realize()
+    self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
+
+  def test_sharded_memory_replicated(self):
+    devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
+    X = Tensor.ones(256).contiguous().realize()
+    self.assertUsed(256 * 4)
+    X.shard_(devices_4).realize()
+    self.assertUsed(256 * 4 * 4)
+
+  def test_sharded_memory_replicated_const(self):
+    devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
+    X = Tensor.ones(256, buffer=False).realize()
+    self.assertUsed(0)
+    X.shard_(devices_4).realize()
+    self.assertUsed(256 * 4 * 4)  # TODO: can be zero
+
+  def test_sharded_memory_axis_const(self):
+    devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
+    X = Tensor.ones(256, buffer=False).realize()
+    self.assertUsed(0)
+    X.shard_(devices_4, axis=0).realize()
+    self.assertUsed(256 * 4)  # TODO: can be zero
+
+  def test_zeros_per_device(self):
+    _ = Tensor.zeros(self.N, self.N, device="NULL").contiguous().realize()
+    self.assertDeviceUsed({"NULL": self.N*self.N*4})
+
+  def test_zeros_del_per_device(self):
+    _ = Tensor.zeros(self.N, self.N, device="NULL").contiguous().realize()
+    del _
+    self.assertDeviceUsed({"NULL": 0})
+
+  def test_zeros_copy_per_device(self):
+    devices_2 = ("NULL:1", "NULL:2")
+    _ = Tensor.zeros(self.N, self.N).contiguous().to(devices_2).realize()
+    self.assertDeviceUsed({"NULL:1": self.N*self.N*4, "NULL:2": self.N*self.N*4})
+
+  def test_zeros_shard_per_device(self):
+    devices_2 = ("NULL:1", "NULL:2")
+    _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).realize()
+    self.assertDeviceUsed({"NULL:1": self.N*(self.N//2)*4, "NULL:2": self.N*(self.N//2)*4})
+
+  def test_sharded_memory_replicated_per_device(self):
+    devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
+    X = Tensor.ones(256, device="NULL").contiguous().realize()
+    self.assertDeviceUsed({"NULL": 256*4})
+    X.shard_(devices_4).realize()
+    for d in devices_4:
+      self.assertDeviceUsed({d: 256*4})
+
+  def _test_matmul_half(self, dev_count:int):
+    N = 32
+    total_mem = {}
+    devs = tuple(f"NULL:{i}" for i in range(dev_count))
+    for dtype in {dtypes.float, dtypes.half}:
+      GlobalCounters.reset()
+      a = Tensor.empty((N, N), dtype=dtype, device=devs[0]).shard(devs, axis=0)
+      b = Tensor.empty((N, N), dtype=dtype, device=devs[0]).shard(devs, axis=None)
+      (a @ b).realize()
+      total_mem[dtype] = GlobalCounters.global_mem
+    self.assertEqual(total_mem[dtypes.half], total_mem[dtypes.float] // 2)
+
+  def test_matmul_half(self): self._test_matmul_half(dev_count=2)
+  def test_matmul_half_alt(self): self._test_matmul_half(dev_count=4)
+
+  def test_multi_layer_allreduce(self):
+    N = 32
+    devices_2 = ("NULL:1", "NULL:2")
+
+    def make_inp():
+      x = Tensor.zeros(N, N).contiguous().shard(devices_2, axis=None).realize()
+      w1 = Tensor.zeros(N, N).contiguous().shard(devices_2, axis=1).realize()
+      w2 = Tensor.zeros(N, N).contiguous().shard(devices_2, axis=0).realize()
+      return x, w1, w2
+
+    def run_layers(n_layers):
+      GlobalCounters.reset()
+
+      @TinyJit
+      def f(x, w1, w2):
+        for _ in range(n_layers):
+          x = (x @ w1 @ w2)
+        return x.contiguous()
+
+      for _ in range(3):
+        a = make_inp()
+        r = f(*a)
+        del a, r
+
+      gc.collect()
+      return GlobalCounters.mem_used
+
+    mem_2 = run_layers(2)
+    mem_4 = run_layers(4)
+    self.assertEqual(mem_2, mem_4, f"graph memory should not grow with layers: 2 layers={mem_2}, 4 layers={mem_4}")
+
+  def test_allreduce_cast_dtype_memory(self):
+    N = 32
+    devices_2 = ("NULL:1", "NULL:2")
+    mem = {}
+    for allreduce_cast in (0, 1):
+      GlobalCounters.reset()
+      with Context(ALLREDUCE_CAST=allreduce_cast, SCACHE=0):
+        x = Tensor.empty((N, N), dtype=dtypes.bfloat16, device="NULL:1").shard(devices_2, axis=0)
+        x.sum(0).realize()
+      mem[allreduce_cast] = GlobalCounters.global_mem
+    # with ALLREDUCE_CAST, allreduce copies happen in bf16 (2 bytes) instead of fp32 (4 bytes)
+    self.assertLess(mem[1], mem[0])
+
+class TestMultiScalarALU(unittest.TestCase):
+  """Test that tuple-device scalars work correctly in ALU with MULTI tensors (_shard scalar fix)."""
+  def test_multi_times_replicated_scalar(self):
+    devices = ("NULL:0", "NULL:1")
+    x = Tensor.ones(4).contiguous().shard(devices, axis=0)
+    s = Tensor(2.0).to(devices)
+    result = x * s
+    self.assertEqual(result.shape, (4,))
+    self.assertEqual(result.uop.axis, 0)
+
+  def test_multi_add_replicated_scalar(self):
+    devices = ("NULL:0", "NULL:1")
+    x = Tensor.ones(4).contiguous().shard(devices, axis=0)
+    s = Tensor(1.0).to(devices)
+    result = x + s
+    self.assertEqual(result.shape, (4,))
+    self.assertEqual(result.uop.axis, 0)
+
+  def test_multi_times_call_scalar(self):
+    """Per-device scalar from a CALL (like FP8 local amax) used in ALU with MULTI."""
+    import functools
+    from tinygrad.uop.ops import Ops
+    devices = ("NULL:0", "NULL:1")
+    x = Tensor.ones(4, 4).contiguous().shard(devices, axis=0)
+    # simulate per-device scalar via CALL (strips MULTI from param body → no allreduce)
+    @functools.cache
+    def _fxn(x_p, device):
+      t = Tensor(x_p, device=device)
+      inner = Tensor(t.uop.src[0]) if t.uop.op is Ops.MULTI else t
+      return (inner.sum(),)
+    param = x.as_param(0)
+    fxn = _fxn(param.uop, x.device)
+    per_dev_scalar = Tensor(fxn[0].uop.call(x.uop).gettuple(0))
+    result = x * per_dev_scalar
+    self.assertEqual(result.shape, (4, 4))
+    self.assertEqual(result.uop.axis, 0)
+    result.realize()
+
+class TestMultiAxis(unittest.TestCase):
+  def test_reshape_shard_invalid(self):
+    devices = ("NULL:0", "NULL:1")
+    t = Tensor.ones(4, 3).shard(devices, axis=0)
+    with self.assertRaises(RuntimeError, msg="reshape cannot move items between shards"):
+      t.reshape(3, 4).uop.axis
+
+  def test_reshape_shard_valid(self):
+    devices = ("NULL:0", "NULL:1")
+    t = Tensor.ones(4, 8).shard(devices, axis=0)
+    self.assertEqual(t.reshape(2, 16).uop.axis, 0)
+    self.assertEqual(t.reshape(2, 2, 8).uop.axis, 0)
+
+  def test_empty_like_sharded(self):
+    t = Tensor.ones(4, 8).shard(("NULL:0", "NULL:1"), axis=0)
+    e = t.empty_like()
+    self.assertEqual(e.shape, t.shape)
+    self.assertEqual(e.device, t.device)
+    self.assertEqual(e.uop.axis, 0)
+    self.assertTrue(e.uop.has_buffer_identity())
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_pattern_matcher.py

@@ -0,0 +1,215 @@
+import unittest, itertools
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import Ops, UOp, GroupOp # noqa: F401
+from tinygrad.uop.ops import PatternMatcher, UPat
+
+class TestPatternMatcher(unittest.TestCase):
+  def test_simple_match(self):
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.int, 1)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c2), None)
+
+  def test_upat_any(self):
+    def test(a, x=None, y=None, z=None):
+      #print(x,y,z)
+      if y is not None: return (a+y).rtag()
+    matcher = PatternMatcher([
+      (UPat.var("a")+UPat.any(UPat.var("x"), UPat.var("y"), UPat.var("z")), test),
+    ])
+    v1 = UOp.variable("a", 0, 10)
+    v2 = UOp.variable("b", 0, 10)
+    c1 = v1+v2
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+
+  def test_minimum_len(self):
+    matcher = PatternMatcher([
+      (UPat(Ops.NOOP, src=(UPat(Ops.NOOP), UPat(Ops.NOOP)), allow_any_len=True), lambda: True),
+    ])
+    self.assertTrue(matcher.rewrite(UOp(Ops.NOOP, src=(UOp(Ops.NOOP), UOp(Ops.NOOP), UOp(Ops.NOOP),))))
+    self.assertTrue(matcher.rewrite(UOp(Ops.NOOP, src=(UOp(Ops.NOOP), UOp(Ops.NOOP)))))
+    self.assertIsNone(matcher.rewrite(UOp(Ops.NOOP, src=(UOp(Ops.NOOP),))))
+
+  @unittest.skip("closures aren't supported on pattern matchers")
+  def test_match_sz_0(self):
+    match_cnt = 0
+    def fxn(x):
+      nonlocal match_cnt
+      match_cnt += 1
+      assert len(x.src) == 0
+      return UOp(Ops.CONST, src=(UOp(Ops.CONST),))
+    matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
+    c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
+    # second rewrite shouldn't match anything
+    c1 = matcher.rewrite(c1)
+    c1 = matcher.rewrite(c1)
+    self.assertEqual(match_cnt, 1)
+
+  def test_match_sz_0_ctx(self):
+    def fxn(ctx, x):
+      ctx.append(True)
+      assert len(x.src) == 0
+      return x.replace(src=(UOp(Ops.DEVICE, arg="blah"),))
+    matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
+    c1 = UOp.const(dtypes.float, 1.0)
+    # second rewrite shouldn't match anything
+    ctx = []
+    c1 = matcher.rewrite(c1, ctx)
+    c1 = matcher.rewrite(c1, ctx)
+    self.assertEqual(len(ctx), 1)
+
+  def test_uop(self):
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp(Ops.ADD, dtypes.float, (c1, c1))
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c2), None)
+
+  def test_uop_set(self):
+    matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.bool, False)
+    c2 = UOp(Ops.CAST, dtypes.int, (c1,))
+    c3 = UOp.const(dtypes.float, 1.0)
+    c4 = UOp(Ops.ADD, dtypes.float, (c3, c3))
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c2), c2.rtag())
+    self.assertEqual(matcher.rewrite(c4), None)
+
+  def test_arg(self):
+    matcher = PatternMatcher([
+      (UPat(Ops.CONST, arg=0, name="x"), lambda x: x.rtag()),
+      (UPat(Ops.CONST, arg=False, name="x"), lambda x: x.rtag()),
+      (UPat(Ops.MAX, name="x"), lambda x: x.rtag()),
+    ])
+    c1 = UOp.const(dtypes.float, 0.0)
+    c2 = UOp.const(dtypes.bool, False)
+    c3 = UOp(Ops.MAX, dtypes.float, (c1, c1))
+    c4 = UOp(Ops.MUL, dtypes.float, (c1, c1))
+    c5 = UOp.const(dtypes.int, -1)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c2), c2.rtag())
+    self.assertEqual(matcher.rewrite(c3), c3.rtag())
+    self.assertEqual(matcher.rewrite(c4), None)
+    self.assertEqual(matcher.rewrite(c5), None)
+
+  def test_filter_arg(self):
+    matcher = PatternMatcher([
+      (UPat(Ops.MUL, src=[UPat(Ops.CONST, name="c"), UPat(Ops.CONST, arg=2)], name="x"),
+       lambda x,c: x.rtag() if c.arg in {1, -1} else None)
+    ])
+    y1 = UOp.const(dtypes.int, 1)
+    y2 = UOp.const(dtypes.int, 2)
+    y3 = UOp.const(dtypes.int, -1)
+    c1 = UOp(Ops.MUL, dtypes.int, (y1, y2))
+    c2 = UOp(Ops.MUL, dtypes.int, (y2, y2))
+    c3 = UOp(Ops.MUL, dtypes.int, (y3, y2))
+    c4 = UOp(Ops.MUL, dtypes.int, (y2, y1))
+    c5 = UOp(Ops.MUL, dtypes.int, (y2, y3))
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c2), None)
+    self.assertEqual(matcher.rewrite(c3), c3.rtag())
+    self.assertEqual(matcher.rewrite(c4), c4.rtag())
+    self.assertEqual(matcher.rewrite(c5), c5.rtag())
+
+  def test_dup_name(self):
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x.rtag())])
+    y1 = UOp.const(dtypes.float, 1.0)
+    y2 = UOp.const(dtypes.float, 1.0)
+    c1 = UOp(Ops.ADD, dtypes.float, (y1, y1))
+    c2 = UOp(Ops.ADD, dtypes.float, (y1, y2))
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c2), c1.rtag())
+
+  def test_dtype(self):
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float32), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float64, 1.0)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c2), None)
+
+  def test_dtype_set(self):
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float64, 1.0)
+    c3 = UOp.const(dtypes.float16, 1.0)
+    c4 = UOp.const(dtypes.int, 1)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c2), c2.rtag())
+    self.assertEqual(matcher.rewrite(c3), None)
+    self.assertEqual(matcher.rewrite(c4), None)
+
+  def test_src_one(self):
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp(Ops.ADD, dtypes.float, (c1,c2))
+    self.assertEqual(matcher.rewrite(c3), c3.rtag())
+    self.assertEqual(matcher.rewrite(c2), None)
+    # that CONST/ALU -> ALU/CONST rewrite is now instant
+    """
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(GroupOp.ALU))), lambda x: x)])
+    c4 = UOp(Ops.ADD, dtypes.float, (c1,c3))
+    c5 = UOp(Ops.ADD, dtypes.float, (c3,c1))
+    self.assertEqual(matcher.rewrite(c3), None)
+    self.assertEqual(matcher.rewrite(c4), c4)
+    self.assertEqual(matcher.rewrite(c5), None)
+    """
+
+  def test_src_permutations(self):
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp(Ops.ADD, dtypes.float, (c1,c2))
+    c4 = UOp(Ops.ADD, dtypes.float, (c3,c2))
+    c5 = UOp(Ops.ADD, dtypes.float, (c2,c3))
+    c6 = UOp(Ops.ADD, dtypes.float, (c3,c4))
+    self.assertEqual(matcher.rewrite(c3), None)
+    self.assertEqual(matcher.rewrite(c4), c4.rtag())
+    self.assertEqual(matcher.rewrite(c5), c5.rtag())
+    self.assertEqual(matcher.rewrite(c6), None)
+
+  def test_src_repeat(self):
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp(Ops.ADD, dtypes.float, (c1,c2))
+    c4 = UOp(Ops.ADD, dtypes.float, (c2,c3))
+    self.assertEqual(matcher.rewrite(c3), c3.rtag())
+    self.assertEqual(matcher.rewrite(c4), None)
+
+  def test_allow_len(self):
+    matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
+    c4 = UOp(Ops.EXP2, dtypes.float, (c1,))
+    c5 = UOp(Ops.ADD, dtypes.float, (c1,c2))
+    c6 = UOp(Ops.MULACC, dtypes.float, (c1,c2,c3))
+    self.assertEqual(matcher.rewrite(c4), None)
+    self.assertEqual(matcher.rewrite(c5), None)
+    self.assertEqual(matcher.rewrite(c6), c6.rtag())
+
+  def test_deep_src_permutations(self):
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    u1 = (c1 + c2) + c1
+    u2 = (c2 + c1) + c1
+    matcher = PatternMatcher([
+      (UPat(GroupOp.ALU, src=[UPat(GroupOp.ALU, src=[UPat(name='a'), UPat(name='b')]), UPat(name='b')]), lambda a,b: b)
+    ])
+    self.assertIsNotNone(matcher.rewrite(u1))
+    self.assertIsNotNone(matcher.rewrite(u2))
+
+  def _assert_eq_upat(self, a:UPat, b:UPat):
+    assert (sorted(map(str,a.op)) if a.op else [] == (sorted(map(str,b.op)) if b.op else []))
+    assert (sorted(a.match_dtype) if a.match_dtype else [] == (sorted(b.match_dtype) if b.match_dtype else []))
+    assert (a.name, type(a.src)) == (b.name, type(b.src))
+    def simple_src(u:UPat):
+      if u.src is None: return []
+      if isinstance(u.src, itertools.repeat): return next(u.src[0])
+      return u.src[0]
+    for a,b in zip(simple_src(a), simple_src(b)): self._assert_eq_upat(a, b)
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)

tinygrad_repo/test/null/test_process_replay.py

@@ -0,0 +1,44 @@
+import unittest
+from tinygrad import Tensor, Device, Context
+from tinygrad.codegen import do_to_program
+from tinygrad.codegen.opt import Opt, OptOps
+from test.external.process_replay.process_replay import replay_to_program
+from test.helpers import replace_opts
+
+N = 16
+class TestProcessReplay(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
+    cls.renderer = Device[Device.DEFAULT].renderer
+
+  def test_replay_no_opts(self):
+    # opts=None means use default heuristic path
+    p = do_to_program(self.ast, self.renderer)
+    good, compare, _ = replay_to_program(p, self.ast, self.renderer)
+    self.assertEqual(good, compare)
+
+  def test_replay_empty_opts(self):
+    # opts=[] means explicitly apply zero opts (unoptimized)
+    ast = replace_opts(self.ast, [])
+    p = do_to_program(ast, self.renderer)
+    good, compare, _ = replay_to_program(p, ast, self.renderer)
+    self.assertEqual(good, compare)
+
+  def test_replay_with_opt(self):
+    # opts=[Opt(...)] means apply a specific opt
+    opts = [Opt(OptOps.UPCAST, 0, 4)]
+    ast = replace_opts(self.ast, opts)
+    p = do_to_program(ast, self.renderer)
+    good, compare, _ = replay_to_program(p, ast, self.renderer)
+    self.assertEqual(good, compare)
+
+  def test_beam(self):
+    with Context(BEAM=1):
+      ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
+    p = do_to_program(ast, self.renderer)
+    good, compare, _ = replay_to_program(p, ast, self.renderer)
+    self.assertEqual(good, compare)
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)

tinygrad_repo/test/null/test_real_world.py

@@ -0,0 +1,189 @@
+import unittest, time, gc
+import numpy as np
+from tinygrad.nn import optim
+from tinygrad.nn.state import get_parameters
+from tinygrad.engine.jit import TinyJit
+from tinygrad import Tensor, Device, GlobalCounters, dtypes, Variable
+from tinygrad.helpers import Context
+from test.helpers import slow, jit_cache_count
+from extra.lr_scheduler import OneCycleLR
+from test.helpers import derandomize_model
+
+from examples.gpt2 import Transformer as GPT2Transformer
+from examples.hlb_cifar10 import SpeedyResNet, hyp
+from extra.models.llama import Transformer as LLaMaTransformer
+from examples.stable_diffusion import UNetModel, unet_params
+from extra.models.unet import ResBlock
+from extra.models.bert import BertForPretraining
+
+global_mem_used = 0
+def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jitted=False):
+  with Context(JIT=2):
+    tms = []
+    for _ in range(2):
+      early_gen = [x.realize() if isinstance(x, Tensor) else x for x in gen()]
+      GlobalCounters.reset()
+      Device[Device.DEFAULT].synchronize()
+      st = time.perf_counter_ns()
+      model(*early_gen)
+      Device[Device.DEFAULT].synchronize()
+      tms.append(time.perf_counter_ns() - st)
+    mem_used = (GlobalCounters.mem_used - global_mem_used) / 1e9
+
+    kernels_used = jit_cache_count(model.captured.linear) if getattr(model, "captured", None) is not None else None
+    print(f"{nm}: used {mem_used/1e9:.2f} GB and {kernels_used} kernels in {min(tms)/1e6:.2f} ms")
+    assert mem_used < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.3f} GB - {mem_used:.3} GB used"
+    assert (max_memory_allowed - mem_used) / max_memory_allowed < 0.2, f"{max_memory_allowed:.3f} GB is too far from {mem_used:.3} GB used"
+    if kernels_used:
+      assert kernels_used <= max_kernels_allowed, f"{nm} used more than {max_kernels_allowed} kernels, it used {kernels_used}"
+      assert (max_kernels_allowed - kernels_used) / max_kernels_allowed < 0.2, f"{max_kernels_allowed=} is too far from {kernels_used=} used"
+    if all_jitted:
+      assert kernels_used > 0 and kernels_used == GlobalCounters.kernel_count or (kernels_used <= GlobalCounters.kernel_count and getattr(Device[Device.DEFAULT], "graph", None)), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted"  # noqa: E501
+
+supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
+
+class TestRealWorld(unittest.TestCase):
+  def setUp(self):
+    gc.collect()
+    global global_mem_used
+    global_mem_used = GlobalCounters.mem_used
+    self.old_float = dtypes.default_float
+    np.random.seed(2002)
+
+  def tearDown(self):
+    dtypes.default_float = self.old_float
+
+  @slow
+  @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
+  def test_stable_diffusion(self):
+    params = unet_params
+    params["model_ch"] = 8
+    params["ctx_dim"] = 8
+    params["num_res_blocks"] = 1
+    params["n_heads"] = 2
+    model = UNetModel(**params)
+    derandomize_model(model)
+    @TinyJit
+    def test(t, t2): return model(t, Tensor([801]), t2).realize()
+    helper_test("test_sd", lambda: (Tensor.randn(1, 4, 32, 32), Tensor.randn(1, 77, params["ctx_dim"])), test, 0.011, 515)
+
+  def test_unet_resblock(self):
+    model = [ResBlock(16, 24, 16) for _ in range(4)]
+    derandomize_model(model)
+    @TinyJit
+    def test(t, t2):
+      for l in model: t = l(t, t2)
+      return t.realize()
+
+    # TODO: support _offset on CL to get mem down to 0.0002
+    exp_mem = 0.00037 if Device.DEFAULT == "CL" else 0.0002
+    helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, exp_mem, 37)
+
+  @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
+  def test_llama(self):
+    dtypes.default_float = dtypes.float16
+
+    args_tiny = {"dim": 1024, "hidden_dim": 2048, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
+    model = LLaMaTransformer(**args_tiny)
+    derandomize_model(model)
+    @TinyJit
+    def test(t): return model(t, 0).realize()
+    # TODO: test first token vs rest properly
+    helper_test("test_llama", lambda: (Tensor([[1,2,3,4]]),), test, 0.23, 118, all_jitted=True)
+
+  @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
+  def test_gpt2(self):
+    dtypes.default_float = dtypes.float16
+
+    args_tiny = {"dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-5, "vocab_size": 1000}
+    model = GPT2Transformer(**args_tiny)
+    derandomize_model(model)
+    @TinyJit
+    def test(t, v):
+      with Context(JIT=0): return model(t, v).realize()
+    helper_test("test_gpt2", lambda: (Tensor([[1,]]),Variable("pos", 1, 100).bind(1)), test, 0.23, 168, all_jitted=True)
+
+  @slow
+  def test_train_mnist(self):
+    from examples.beautiful_mnist import Model
+    with Tensor.train():
+      model = Model()
+      optimizer = optim.Adam(get_parameters(model))
+      BS = 32
+
+      @TinyJit
+      def train(X):
+        out = model(X)
+        loss = out.mean()
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.017, 103)
+
+  @slow
+  def test_forward_cifar(self):
+    BS = 32
+    # with training batchnorm still though
+    with Tensor.train():
+      model = SpeedyResNet(Tensor.ones((12,3,2,2)))
+      @TinyJit
+      def run(X): return model(X)
+      helper_test("forward_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), run, 0.033, 27)
+
+  @slow
+  def test_train_cifar(self):
+    with Tensor.train():
+      model = SpeedyResNet(Tensor.ones((12,3,2,2)))
+      optimizer = optim.SGD(get_parameters(model), lr=0.01, momentum=0.8, nesterov=True, weight_decay=0.15)
+      BS = 32
+
+      @TinyJit
+      def train(X):
+        out = model(X)
+        loss = out.mean()
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+      helper_test("train_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), train, 0.12, 126)
+
+  @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
+  def test_train_cifar_hyp(self):
+    dtypes.default_float = dtypes.float16
+    with Tensor.train():
+      model = SpeedyResNet(Tensor.ones((12,3,2,2)))
+      optimizer = optim.SGD(get_parameters(model), lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['bias_decay'])
+      initial_div_factor = hyp['opt']['initial_div_factor']
+      final_lr_ratio = hyp['opt']['final_lr_ratio']
+      pct_start = hyp['opt']['percent_start']
+      lr_scheduler = OneCycleLR(optimizer, max_lr=hyp['opt']['bias_lr'], pct_start=pct_start, div_factor=initial_div_factor,
+                                final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=4)
+      assert not np.isnan(lr_scheduler.min_lr), "lr too small or initial_div_facotr too big for half"
+
+  @slow
+  def test_bert(self):
+    with Tensor.train():
+      args_tiny = {"attention_probs_dropout_prob": 0.0, "hidden_dropout_prob": 0.0, "vocab_size": 30522, "type_vocab_size": 2,
+                  "max_position_embeddings": 512, "hidden_size": 128, "intermediate_size": 512, "num_attention_heads": 2, "num_hidden_layers": 2}
+      model = BertForPretraining(**args_tiny)
+      optimizer = optim.LAMB(get_parameters(model))
+
+      @TinyJit
+      def train(input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
+                masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
+        lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
+        loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+      from examples.mlperf.helpers import get_fake_data_bert
+      data = get_fake_data_bert(BS=4)
+      for v in data.values(): v.to_(Device.DEFAULT)
+
+      helper_test("train_bert", lambda: (data["input_ids"], data["segment_ids"], data["input_mask"], data["masked_lm_positions"], \
+          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.31, 400)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_rearrange_einops.py

@@ -0,0 +1,171 @@
+# modified from
+# https://github.com/arogozhnikov/einops/blob/master/tests/test_examples.py
+# https://github.com/arogozhnikov/einops/blob/master/tests/test_ops.py
+# https://github.com/arogozhnikov/einops/blob/master/tests/test_parsing.py
+
+import numpy as np
+import unittest
+from tinygrad import Tensor
+
+
+class test_rearrange_examples(unittest.TestCase):
+  def test1(self):
+    # transpose
+    x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
+    y = x.rearrange("b c h w -> b h w c")
+    assert tuple(y.shape) == (10, 30, 40, 20)
+
+  def test2(self):
+    # view / reshape
+    x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
+    y = x.rearrange("b c h w -> b (c h w)")
+    assert tuple(y.shape) == (10, 20 * 30 * 40)
+
+  def test3(self):
+    # depth-to-space
+    x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
+    y = x.rearrange("b (c h1 w1) h w -> b c (h h1) (w w1)", h1=2, w1=2)
+    assert tuple(y.shape) == (10, 5, 30 * 2, 40 * 2)
+
+  def test4(self):
+    # space-to-depth
+    x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
+    y = x.rearrange("b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=2, w1=2)
+    assert tuple(y.shape) == (10, 20 * 4, 30 // 2, 40 // 2)
+
+  def test5(self):
+    # simple transposition
+    x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
+    y = x.rearrange("b1 sound b2 letter -> b1 b2 sound letter")
+    assert tuple(y.shape) == (10, 30, 20, 40)
+
+  def test6(self):
+    # parsing parameters
+    x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
+    t = x.rearrange("b c h w -> (b h w) c")
+    t = t[:, ::2]  # replacement for dot-product, just changes size of second axis
+    assert tuple(t.shape) == (10 * 30 * 40, 10)
+
+  def test7(self):
+    x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
+    # split of embedding into groups
+    y1, y2 = x.rearrange("b (c g) h w -> g b c h w", g=2)
+    assert tuple(y1.shape) == (10, 10, 30, 40)
+    assert tuple(y2.shape) == (10, 10, 30, 40)
+
+  def test8(self):
+    x = Tensor(np.arange(10 * 20 * 1 * 1, dtype=np.int32).reshape([10, 20, 1, 1]))
+    # squeeze - unsqueeze
+    y = x.rearrange("b c () () -> b c")
+    assert tuple(y.shape) == (10, 20)
+    y = y.rearrange("b c -> c b () ()")
+    assert tuple(y.shape) == (20, 10, 1, 1)
+
+  def test9(self):
+    x = Tensor(np.arange(10 * 20 * 1 * 1, dtype=np.int32).reshape([10, 20, 1, 1]))
+    # squeeze - unsqueeze
+    y = x.rearrange("b c 1 1 -> b c")
+    assert tuple(y.shape) == (10, 20)
+    y = y.rearrange("b1 c -> c b1 1 1")
+    assert tuple(y.shape) == (20, 10, 1, 1)
+
+
+class test_rearrange_ops(unittest.TestCase):
+  def test_rearrange_errors(self):
+    x = Tensor.zeros([1, 1, 1, 1, 1])
+    x.rearrange("a b c d ... ->  a b c ... d")
+    bad_patterns = [
+      "a b c d (...) ->  a b c ... d",  # collapsed ellipsis on input
+      "a b (c d ... ->  a b c ... d",   # unbalanced brackets
+      "a b* c d ... ->  a b c ... d",   # not alphanumeric
+      "a b c d ->  a b c d -> a b c d", # two "->"
+      "a ... c ... ->  ... a ... c",    # two "..."
+      "a b c d e -> f b c d e",         # name mismatch
+    ]
+    for pattern in bad_patterns:
+      with self.assertRaises(AssertionError):
+        x.rearrange(pattern)
+
+    x.rearrange("... ->  (...)")
+    with self.assertRaises(AssertionError):
+      x.rearrange("(...) -> (...)")
+
+    y = Tensor.zeros([8, 1])
+    y.rearrange("(a1 a2 a3) b -> b a3 a2 a1", a1=2, a2=2)
+    with self.assertRaises(RuntimeError):
+      ## should fail as not enough dimensions specified
+      y.rearrange("(a1 a2 a3) b -> b a3 a2 a1", a1=2)
+    with self.assertRaises(ValueError):
+      ## should fail as 6 does not divide 8
+      y.rearrange("(a1 a2 a3) b -> b a3 a2 a1", a1=3, a2=2)
+    with self.assertRaises(AssertionError):
+      ## incorrect dimension provided for an axis that is only permuted
+      y.rearrange("(a1 a2 a3) b -> b a3 a2 a1", a1=2, a2=2, b=2)
+    with self.assertRaises(AssertionError):
+      ## unused axis provided
+      y.rearrange("(a b c) d -> a b c d", b=2, c=2, e=2)
+
+
+class test_rearrange_parsing(unittest.TestCase):
+  def test_elementary_axis_name(self):
+    for name in [
+      "a",
+      "b",
+      "h",
+      "dx",
+      "h1",
+      "zz",
+      "i9123",
+      "somelongname",
+      "Alex",
+      "camelCase",
+      "u_n_d_e_r_score",
+      "unreasonablyLongAxisName",
+    ]:
+      Tensor.ones((1,)).rearrange(f"{name} -> {name}")
+
+    for name in ["2b", "12", "_startWithUnderscore", "endWithUnderscore_", "_"]:
+      with self.assertRaises(AssertionError):
+        Tensor.ones((1,)).rearrange(f"{name} -> {name}")
+
+    with self.assertRaises(RuntimeError):
+      Tensor.ones((1,)).rearrange(" -> ")
+
+  def test_invalid_expressions(self):
+    # double ellipsis should raise an error
+    def _test_expression(expression: str):
+      Tensor.ones((2, 3, 4, 5, 6)).rearrange(f"{expression} -> {expression}")
+
+    _test_expression("... a b c d")
+    with self.assertRaises(AssertionError):
+      _test_expression("... a b c d ...")
+    with self.assertRaises(AssertionError):
+      _test_expression("... a b c (d ...)")
+    with self.assertRaises(AssertionError):
+      _test_expression("(... a) b c (d ...)")
+
+    # double/missing/enclosed parenthesis
+    Tensor.ones((2, 3, 4, 5, 6)).rearrange("a b c d ... -> (a) b c (d ...)")
+    with self.assertRaises(AssertionError):
+      _test_expression("(a)) b c (d ...)")
+    with self.assertRaises(AssertionError):
+      _test_expression("(a b c (d ...)")
+    with self.assertRaises(AssertionError):
+      _test_expression("(a) (()) b c (d ...)")
+    with self.assertRaises(AssertionError):
+      _test_expression("(a) ((b c) (d ...))")
+
+    # invalid identifiers
+    _test_expression("camelCase under_scored cApiTaLs ß ...")
+    with self.assertRaises(AssertionError):
+      _test_expression("1a")
+    with self.assertRaises(AssertionError):
+      _test_expression("_pre")
+    with self.assertRaises(AssertionError):
+      _test_expression("...pre")
+    with self.assertRaises(AssertionError):
+      _test_expression("pre...")
+
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/null/test_resnet.py

@@ -0,0 +1,21 @@
+import unittest
+from extra.models import resnet
+
+class TestResnet(unittest.TestCase):
+  def test_model_load(self):
+    model = resnet.ResNet18()
+    model.load_from_pretrained()
+
+    model = resnet.ResNeXt50_32X4D()
+    model.load_from_pretrained()
+
+  def test_model_load_no_fc_layer(self):
+    model = resnet.ResNet18(num_classes=None)
+    model.load_from_pretrained()
+
+    model = resnet.ResNeXt50_32X4D(num_classes=None)
+    model.load_from_pretrained()
+
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_rewrite_bottom_up_gate.py

@@ -0,0 +1,28 @@
+import unittest
+from tinygrad.uop.ops import PatternMatcher, UOp, graph_rewrite, Ops, UPat, BottomUpGate
+
+def assert_not_reached(): assert False, "This function should not be reached"
+def gate(): raise BottomUpGate
+
+class TestBottomUpGate(unittest.TestCase):
+  def test_basic_bottom_up_gate(self):
+    """Test that BottomUpGate stops bottom-up"""
+    pm = PatternMatcher([
+      (UPat(Ops.ADD), gate),
+      (UPat(Ops.MUL), assert_not_reached)
+    ])
+
+    a,b,c = UOp.variable("a",0,10), UOp.variable("b",0,10), UOp.variable("c",0,10)
+    graph_rewrite((a*a)+(b*c), pm, bottom_up=True)
+
+  def test_bottom_up_gate_with_rewriting(self):
+    pm = PatternMatcher([
+      (UPat.var("a")+UPat.var("a"), lambda a: 2*a),
+      (UPat(Ops.MUL), gate),
+      (UPat(Ops.CONST), assert_not_reached)
+    ])
+    a = UOp.variable("a",0,10)
+    graph_rewrite(a+a, pm, bottom_up=True)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/null/test_schedule_cache.py

@@ -0,0 +1,41 @@
+import unittest
+from tinygrad import Tensor, Variable, Context
+from tinygrad.helpers import cpu_events
+from tinygrad.schedule import schedule_cache
+
+def schedule_one():
+  Tensor([1]).schedule_linear()
+
+class TestScheduleCache(unittest.TestCase):
+  def test_bound_variable_var_vals(self):
+    v = Variable('pos', 1, 100)
+    x = Tensor.ones(10).contiguous().realize()
+
+    t = x + Tensor(v.bind(42))
+    _, var_vals = t.linear_with_vars()
+    self.assertEqual(var_vals, {'pos': 42})
+
+  def test_disable_schedule_cache(self):
+    schedule_cache.clear()
+
+    # test write
+    with Context(SCACHE=0): schedule_one()
+    self.assertEqual(len(schedule_cache), 0)
+    with Context(SCACHE=1):
+      schedule_one()
+      schedule_one()
+    self.assertEqual(len(schedule_cache), 1)
+
+    # test read
+    with Context(PROFILE=1):
+      cpu_events.clear()
+      with Context(SCACHE=0): schedule_one()
+      num_events_no_cache = len(cpu_events)
+
+      cpu_events.clear()
+      with Context(SCACHE=1): schedule_one()
+      num_events_cache = len(cpu_events)
+    self.assertLess(num_events_cache, num_events_no_cache)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/null/test_simplify_valid_idx.py

@@ -0,0 +1,600 @@
+import unittest, itertools
+
+from tinygrad.codegen.late.devectorizer import load_store_indexing
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite
+from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
+from tinygrad.helpers import Context
+from test.helpers import full_rewrite
+from test.null.test_uop_symbolic import check_uop_against_string
+
+# symbolic-only idx + valid simplification (no late lowering of FLOORDIV/FLOORMOD)
+def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load, name="simplify_valid_idx")
+# image-aware idx + valid simplification: adds the codegen-layer matcher that drops provably in-bounds gates
+def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+load_store_indexing, name="simplify_image_idx")
+
+def get_gated_load_uop(valid:UOp, idx:UOp):
+  return UOp(Ops.LOAD, dtypes.float, (
+    UOp.param(0, dtypes.float.ptr()).index(idx.valid(valid), ptr=True),
+    UOp.const(dtypes.float, 0.0)
+  ))
+
+def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
+  return UOp(Ops.LOAD, dtypes.float.vec(4), (
+    UOp.param(0, dtypes.imagef(image_shape)).index(idx[1].valid(valid), idx[0].valid(valid), ptr=True),
+    UOp(Ops.STACK, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4)
+  ))
+
+def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.weakint, (UOp.const(dtypes.weakint, nmax),), expr)
+def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
+def Range(n, nmax): return UOp.range(nmax, n)
+
+class TestHelpers(unittest.TestCase):
+  def test_is_increasing(self):
+    idx1 = Special("idx1", 32)
+    idx2 = Special("idx2", 64)
+    ridx0 = Variable("ridx0", 0, 5)
+    ridx1 = Variable("ridx1", 0, 2)
+    ridx2 = Variable("ridx2", 0, 2)
+    # (ridx0+(idx1*48)+(ridx2*6)+(-6)),((idx2*2)+ridx1+(-1)))
+    f0 = ((idx1*24)+(ridx2*3)+ridx0+765)%768
+    f1 = ridx0+(idx1*48)+(ridx2*6)+(-6)
+    f2 = (idx2*2)+ridx1+((idx1+((ridx2+7)//8)+31)//32)+(-2)
+    f3 = (idx2*2)+ridx1+(-1)
+
+    self.assertFalse(f0.is_increasing())
+    self.assertTrue(f1.is_increasing())
+    self.assertTrue(f2.is_increasing())
+    self.assertTrue(f3.is_increasing())
+
+    rng = UOp.range(5, 2)
+    self.assertTrue(rng.is_increasing())
+    self.assertTrue((rng+2).is_increasing())
+
+class TestValidIdxSimplification(unittest.TestCase):
+  def check(self, load, sidx, svalid, extra=()):
+    load = simplify_valid_idx(UOp.sink(load, *extra)).src[0]
+    off = load.src[0].src[1]
+    check_uop_against_string(self, off.get_idx(), sidx)
+    check_uop_against_string(self, off.get_valid(), svalid)
+
+  def test_cumsum(self):
+    gidx0 = Special("gidx0", 5)
+    lidx0 = Special("lidx0", 4)
+    gate = (gidx0*4+lidx0<19).ne(True)
+    idx = gidx0*4+lidx0-19
+    load = get_gated_load_uop(gate, idx)
+    self.check(load,
+      "0",
+      "(((lidx0+(gidx0*4))<19)!=True)")
+
+  def test_simplify_within_valid1(self):
+    ridx0 = Range(0, 4)
+    ridx1 = Range(1, 4)
+    ridx2 = Range(2, 4)
+    ridx3 = Range(3, 4)
+    valid = ((ridx0*3+ridx1)<8) & ((((ridx0*3+ridx1)//8+ridx2*3+ridx3)%4)<2)
+    idx = ridx0+ridx1+ridx2+ridx3
+    load = get_gated_load_uop(valid, idx)
+    self.check(load,
+      "(((r0+r1)+r2)+r3)",
+      "((((r0*3)+r1)<8)&((((r2*3)+r3)%4)<2))")
+
+  def test_simplify_within_valid2(self):
+    gidx0 = Special("gidx0", 56)
+    ridx0 = Range(0, 3)
+    alu0 = gidx0+ridx0
+    valid = (alu0 < 57) & (alu0 >= 1)
+    self.assertIsNone(simplify_valid(valid))
+
+  def test_valid_order_matters1(self):
+    ridx0 = Range(0, 2)
+    v0 = ridx0<1
+    v1 = ((ridx0*5+1)%6)<5
+    self.assertEqual(simplify_valid(v0&v1).render(), "(r0<1)")
+    self.assertEqual(simplify_valid(v1&v0).render(), "(r0<1)")
+
+  def test_valid_order_matters2(self):
+    gidx0 = Special("gidx0", 13)
+    gidx1 = Special("gidx1", 13)
+    ridx0 = Range(0, 4)
+    alu0 = (gidx1+(ridx0*13))
+    v0 = (gidx0+11)%14<11
+    v1 = (alu0+((gidx0+39)//42))%14<11
+    v2 = gidx0<3
+    v3 = alu0<42
+
+    for v in itertools.permutations([v0,v1,v2,v3]):
+      self.assertEqual(simplify_valid(v[0]&v[1]&v[2]&v[3]).render(), "False")
+
+  def test_simplify_valid_from_div(self):
+    x = Variable("x", -100, 100)
+    valid = ((x<0)&((100%x).cast(dtypes.bool)))
+    # NOTE: this simplifies the (100%x) part somehow, still has two clauses
+    self.assertIsNotNone(simplify_valid(valid))
+    self.assertEqual(len(list(valid.split_uop(Ops.AND))), 2)
+
+  @unittest.expectedFailure  # TODO: fix
+  def test_from_merge_views(self):
+    # taken from test_merges_from_fuzzer1
+    # generated by
+    # v0 = View(shape=(2, 4), strides=(2, 1), offset=-2, mask=((0, 2), (2, 4)), contiguous=False)
+    # v1 = View(shape=(2, 4, 2, 2), strides=(4, 0, -2, -1), offset=3, mask=None, contiguous=False)
+    # s = ShapeTracker((v0, v1))
+    # idx, valid = s.to_indexed_uops()
+    # print(f"{idx.render()=}")
+    # print(f"{valid.render()=}")
+
+    # s = ShapeTracker((View(shape=(2, 4, 2, 2), strides=(2, 0, 0, -1), offset=1, mask=((0, 2), (0, 4), (0, 1), (0, 2)), contiguous=False),))
+    # idx, valid = s.to_indexed_uops()
+    # print(f"{idx.render()=}")
+    # print(f"{valid.render()=}")
+    ridx0 = Range(0, 2)
+    ridx2 = Range(2, 2)
+    ridx3 = Range(3, 2)
+    idx = (((ridx0*2)+((((ridx2*2)+(ridx3*3))+3)%4))+-2)
+    valid = ((((((ridx2*2)+(ridx3*3))+3)%4)<2)!=True) # noqa: E712
+    load = get_gated_load_uop(valid, idx)
+    self.check(load,
+      "(((r0*2)+(r3*-1))+1)",
+      "(r2<1)")
+
+  def test_load_in_valid(self):
+    # from FUSE_ARANGE=1 python test/test_ops.py TestOps.test_scatter_add
+    # can lead to OOB
+    ridx2 = Range(2, 4)
+    lidx0 = Special("lidx0", 3)
+    gidx0 = Special("gidx0", 2)
+    idx=(((lidx0+(gidx0*3))+(ridx2*5))+40)
+    valid = (lidx0+(gidx0*3)) < 5
+    val7 = get_gated_load_uop(valid, idx)
+    valid2 = valid & val7.cast(dtypes.bool).logical_not()
+    self.assertIsNone(simplify_valid(valid2))
+
+  def test_valid_becomes_const1(self):
+    # from DSP mobilenetv2
+    ridx0 = Range(0, 30)
+    ridx1 = Range(1, 7)
+    ridx2 = Range(2, 2)
+    alu11 = (ridx1+ridx2)
+    alu15 = ((alu11+1)//7)
+    idx = (alu15*-31)+(((((alu11+218)//224)+ridx0)%30)*1568)
+    valid = (ridx2<1)&(ridx1<6)
+    load = get_gated_load_uop(valid, idx)
+    # prevent ridx1 and ridx2 from being shrunk
+    red = load.reduce(ridx1, ridx2, arg=Ops.ADD)
+    self.check(load,
+      "(r0*1568)",
+      "((r2<1)&(r1<6))",
+      extra=(red,))
+
+  def test_valid_becomes_const1_z3(self):
+    from z3 import Ints, Solver, And, If, Not, unsat
+    ridx0, ridx1, ridx2, alu11, alu15 = Ints('ridx0 ridx1 ridx2 alu11 alu15')
+    alu11 = (ridx1+ridx2)
+    alu15 = ((alu11+1)/7)
+    idx = (alu15*-31)+(((((alu11+218)/224)+ridx0)%30)*1568)
+    valid = (ridx2<1)&(ridx1<6)
+    load = If(valid, idx, 0)
+
+    # correct simplification
+    s = Solver()
+    s.add(And(0<=ridx0, ridx0<30, 0<=ridx1, ridx1<7, 0<=ridx2, ridx2<2))
+    simplifed_idx = (ridx0*1568)
+    simplifed_load = If(valid, simplifed_idx, 0)
+    s.add(Not(load == simplifed_load))  # Check if they are NOT equivalent
+    assert s.check() == unsat, f"The expressions are not equivalent. {s.model()=}"
+
+    # new solver for a wrong simplified expression
+    s = Solver()
+    s.add(And(0<=ridx0, ridx0<30, 0<=ridx1, ridx1<7, 0<=ridx2, ridx2<2))
+    wrong_simplifed_idx = (ridx0*1567)+ridx1
+    wrong_simplifed_load = If(valid, wrong_simplifed_idx, 0)
+    s.add(Not(load == wrong_simplifed_load))  # Check if they are NOT equivalent
+    assert s.check() != unsat, "The expressions are equivalent??"
+    print("The expressions are not equivalent.")
+    print(s.model())
+
+  def test_valid_becomes_const2(self):
+    ridx0 = Range(0, 4)
+    ridx1 = Range(1, 4)
+    ridx2 = Range(2, 4)
+    ridx3 = Range(3, 4)
+    # TODO: this should also work without the extra nesting
+    idx = (((ridx0+ridx1)+(ridx2+ridx3)+28)//30)
+    valid = ((ridx0+ridx1)<1).ne(True) & ((ridx2+ridx3)<1).ne(True)
+    load = get_gated_load_uop(valid, idx)
+    self.check(load,
+      "1",
+      "((((r0+r1)<1)!=True)&(((r2+r3)<1)!=True))")
+
+  def test_valid_with_non_const_rhs(self):
+    ridx0 = Range(0, 1024)
+    ridx1 = Range(1, 4)
+    ridx2 = Range(2, 4)
+    valid = (ridx0<(ridx1*4 + ridx2))&(ridx0<-1).ne(True)
+    idx = ridx0
+    load = get_gated_load_uop(valid, idx)
+    self.check(load,
+      "r0",
+      "(r0<((r1*4)+r2))")
+
+class TestImageSimplification(unittest.TestCase):
+  def check(self, load, svalid, sidx0, sidx1):
+    load = simplify_image_idx(load.sink()).src[0]
+    off = load.src[0]
+    self.assertEqual(len(off.src), 3)
+    idx0, idx1 = off.src[2].get_idx(), off.src[1].get_idx()
+    check_uop_against_string(self, idx0, sidx0)
+    check_uop_against_string(self, idx1, sidx1)
+    self.assertEqual(off.src[1].get_valid(), off.src[2].get_valid())
+    if svalid is not None:
+      check_uop_against_string(self, off.src[1].get_valid(), svalid)
+    else:
+      self.assertEqual(off.src[1].get_valid(), UOp.const(dtypes.bool, True), "svalid is None but valid is not True")
+
+  def test_idx_gt_c(self):
+    # (idx1 < c+1).ne(True) ? (..., idx1-1+c) : 0 can drop the valid
+    # (idx1 < c+1).ne(True) -> idx > c
+    gidx0 = Special("gidx0", 32)
+    gidx1 = Special("gidx1", 32)
+    shape = (10, 10, 4)
+    load = get_load_image_uop(shape, (gidx1<1).ne(True), (gidx0, gidx1-1))
+    self.check(load, None, "gidx0", "(gidx1+-1)")
+    load = get_load_image_uop(shape, (gidx1<1).ne(True), (gidx0, gidx1-2))
+    self.check(load, None, "gidx0", "(gidx1+-2)")
+
+    # should match any one of the AND clause and drop the matched statement from valid
+    valid = (gidx0<1).ne(True) & (gidx1<1).ne(True)
+    load = get_load_image_uop(shape, valid, (gidx0+1, gidx1-1))
+    self.check(load, "((gidx0<1)!=True)", "(gidx0+1)", "(gidx1+-1)")
+
+    valid = (gidx0<1).ne(True) & (gidx1<1).ne(True)
+    load = get_load_image_uop(shape, valid, (gidx0, gidx1-1))
+    self.check(load, "((gidx0<1)!=True)", "gidx0", "(gidx1+-1)")
+
+  def test_idx_lt_bound(self):
+    # (idx1 < image_bound) ? (..., idx1) : 0 can drop the valid
+    gidx0 = Special("gidx0", 32)
+    gidx1 = Special("gidx1", 32)
+    load = get_load_image_uop((10, 10, 4), gidx1<10, (gidx0, gidx1))
+    self.check(load, None, "gidx0", "gidx1")
+
+    # same thing, valid has a div
+    load = get_load_image_uop((10, 10, 4), gidx1//2<5, (gidx0, gidx1))
+    self.check(load, None, "gidx0", "gidx1")
+
+    # 10x20 image, not out of bound
+    load = get_load_image_uop((20, 10, 4), gidx1<10, (gidx0, gidx1))
+    self.check(load, "(gidx1<10)", "gidx0", "gidx1")
+
+  def test_generic_idx_lt_bound(self):
+    # (idx1 < image_bound - c) ? (..., idx1 + c) : 0 can drop the valid
+    gidx0 = Special("gidx0", 32)
+    gidx1 = Special("gidx1", 32)
+    shape = (10, 10, 4)
+    load = get_load_image_uop(shape, (gidx1<8), (gidx0, gidx1+2))
+    self.check(load, None, "gidx0", "(gidx1+2)")
+
+    load = get_load_image_uop(shape, (gidx1<5), (gidx0, gidx1+5))
+    self.check(load, None, "gidx0", "(gidx1+5)")
+
+  def test_valid_empty_set(self):
+    gidx0 = Special("gidx0", 32)
+    gidx1 = Special("gidx1", 32)
+    shape = (32, 32, 4)
+    idx = (gidx0%2, gidx1+2)
+    # not empty
+    load = get_load_image_uop(shape, gidx0<8, idx)
+    self.check(load, "(gidx0<8)", "(gidx0%2)", "(gidx1+2)")
+
+    # empty -> invalid
+    load = get_load_image_uop(shape, (gidx0<8) & (gidx0<8).ne(True), idx)
+    with Context(NOOPT=1, SPEC=0):
+      load = full_rewrite(load.sink()).src[0]
+    self.assertEqual(load.op, Ops.STACK)
+    self.assertEqual(load.dtype.count, 4)
+
+  def test_openpilot_conv1(self):
+    # first conv in openpilot
+    # kernel in tinygrad ae5d1407ee844a97a52ad3756835d38e7e2b9e1b https://gist.github.com/chenyuxyz/39c2d4e9a076b46731c67d345ff066b6
+    idx1 = Special("idx1", 32)
+    idx2 = Special("idx2", 64)
+    # ridx0 = Variable("ridx0", 0, 5)
+    # ridx1 = Variable("ridx1", 0, 2)
+    # ridx2 = Variable("ridx2", 0, 2)
+    ridx0 = Range(0, 6)
+    ridx1 = Range(1, 3)
+    ridx2 = Range(2, 3)
+
+    alu1 = ((idx2*2)+ridx1)
+    alu4 = ((idx1*48)+(ridx2*6)+ridx0)
+
+    valid = ((((idx2*2)+(ridx1))<1).ne(True))&((((idx1*8)+(ridx2))<1).ne(True))
+    shape = (128, 1536, 4)
+    idx = ((alu4+1530)%1536, alu1+((idx1+((ridx2+7)//8)+31)//32)+(-2))
+
+    load = get_load_image_uop(shape, valid, idx)
+    self.check(load, None, "((((idx1*48)+(r2*6))+r0)+-6)", "(((idx2*2)+r1)+-1)")
+
+  def test_openpilot_conv2(self):
+    # conv in test/external/external_test_valid_remove.py
+    idx1 = Special("idx1", 32)
+    idx2 = Special("idx2", 64)
+    # ridx0 = Variable("ridx0", 0, 2)
+    # ridx1 = Variable("ridx1", 0, 2)
+    # ridx2 = Variable("ridx2", 0, 2)
+    ridx0 = Range(0, 3)
+    ridx1 = Range(1, 3)
+    ridx2 = Range(2, 3)
+
+    alu1 = ((idx2*2)+ridx1)
+    alu3 = ((idx1*24)+(ridx2*3)+ridx0)
+
+    valid = ((((idx2*2)+ridx1)<1).ne(True))&((((idx1*8)+ridx2)<1).ne(True))
+    shape = (128, 768, 4)
+    idx = ((alu3+765)%768, alu1+((idx1+((ridx2+7)//8)+31)//32)+(-2))
+    load = get_load_image_uop(shape, valid, idx)
+
+    self.check(load, None, "((((idx1*24)+(r2*3))+r0)+-3)", "(((idx2*2)+r1)+-1)")
+
+  def test_openpilot_conv3(self):
+    # in openpilot 0.9.7
+    idx0 = Special("idx0", 64)
+    idx1 = Special("idx1", 2)
+    idx2 = Special("idx2", 4)
+    ridx0 = Range(0, 7)
+    ridx1 = Range(1, 7)
+
+    alu2 = ((idx2*2)+ridx0)
+    alu4 = ((idx1*8)+ridx1)
+    alu6 = ((idx1*512)+(ridx1*64)+idx0)
+
+    valid = (alu2<11)&(alu4<3).ne(True)
+    shape = (8, 1024, 4)
+    idx = (((alu6+832)%1024),(alu2+((idx1+((ridx1+5)//8)+1)//2)+(-4)))
+
+    load = get_load_image_uop(shape, valid, idx)
+
+    self.check(load,
+               "((((idx2*2)+r0)<11)&((((idx1*8)+r1)<3)!=True))",
+               "(idx0+(idx1*512+r1*64)+-192)",
+               "((((idx2*2)+r0)+(((idx1+((r1+5)//8))+1)//2))+-4)")
+
+  def test_simplify1(self):
+    # idx has the form (A % m, A // m + k) and valid has (c0 < A) and (A < c1)
+    gidx = Special("gidx", 512)
+    valid = (gidx<488) & (gidx<480).ne(True)
+    idx = ((gidx*3+18)%26, (gidx*3+18)//26-56)
+    load = get_load_image_uop((1, 26, 4), valid, idx)
+    self.check(load, None, "((gidx*3)+-1438)", "0")
+
+  def test_simplify2(self):
+    # from DEV=CL DEBUG=4 FORWARD_ONLY=1 IMAGE=2 python3 test/test_ops.py TestOps.test_simple_padding_conv2d
+    lidx = Special("lidx", 4)
+    valid = (lidx<3) & (lidx<1).ne(True)
+    idx = ((lidx+1)%2, (lidx+1)//2-1)
+    load = get_load_image_uop((1, 2, 4), valid, idx)
+    self.check(load, None, "(lidx+-1)", "0")
+
+  def test_simplify3(self):
+    # from openpilot
+    idx0 = Special("idx0", 265)
+    valid = (idx0<201).ne(True)
+    idx = ((idx0+55)%64, (idx0+55)//64-4)
+    load = get_load_image_uop((1, 64, 4), valid, idx)
+    self.check(load, None, "(idx0+-201)", "0")
+
+  def test_simplify4(self):
+    idx0 = Special("idx0", 512)
+    shape = (4, 64, 4)
+    alu2 = ((idx0*4+1)%32)
+    alu3 = ((idx0*4+2)%32)
+    alu4 = ((idx0*4+3)%32)
+    alu5 = (idx0*4%32)
+    alu8 = (idx0//8%32//4)
+    alu9 = idx0<256
+
+    load = get_load_image_uop(shape, alu9, (((alu8+(alu2*8))%64),(alu2//8)))
+    self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+8)", "(idx0//2%4)")
+
+    load = get_load_image_uop(shape, alu9, (((alu8+(alu3*8))%64),(alu3//8)))
+    self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+16)", "(idx0//2%4)")
+
+    load = get_load_image_uop(shape, alu9, (((alu8+(alu4*8))%64),(alu4//8)))
+    self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+24)", "(idx0//2%4)")
+
+    load = get_load_image_uop(shape, alu9, (((alu8+(alu5*8))%64),(alu5//8)))
+    self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32)", "(idx0//2%4)")
+
+  def test_simplify5(self):
+    # openpilot 0.9.7, chunk replacement to simplify
+    shape = (10, 384, 4)
+    idx0 = Special("idx0", 16)
+    idx1 = Special("idx1", 24)
+    alu0 = idx0*4
+    alu1 = (idx1*256)+alu0
+    alu2 = idx1//3
+    alu3 = ((alu1+1)%768)
+    idx = ((idx0+((((alu3//640)+alu2)%8)*16)+128),((alu3//64)%10))
+    valid = alu3<640
+
+    load = get_load_image_uop(shape, valid, idx)
+    self.check(load, None, "((idx0+((idx1//3)*16))+128)", "((idx1%3)*4)")
+
+  def test_simplify6(self):
+    # from openpilot
+    # the valid implies the numerator of the div/mod is positive and can be simplified with floordiv rules
+    idx1 = Special("idx1", 16)
+    idx2 = Special("idx2", 64)
+    ridx3 = Range(3, 3)
+    ridx4 = Range(4, 3)
+    ridx5 = Range(5, 3)
+    alu0 = ((idx2*1536)+(ridx4*768)+ridx3+(idx1*24)+(ridx5*3)+-771)%768
+    alu1 = ((idx2*1536)+(ridx4*768)+ridx3+(idx1*24)+(ridx5*3)+-771)//768
+    valid = (((idx2+ridx4)<1)!=1)&(((idx1+ridx5)<1)!=1)
+    load = get_load_image_uop((128, 768, 4), valid, (alu0, alu1))
+    self.check(load, None, "((((idx1*24)+r3)+(r5*3))+-3)", "(((idx2*2)+r4)+-1)")
+
+  def test_simplify7(self):
+    # DEBUG=2 ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916 # noqa: E501
+    # kernel 143
+    gidx0 = Special("gidx0", 32)
+    lidx0 = Special("lidx0", 16)
+    lidx1 = Special("lidx1", 8)
+    r0 = Range(0, 7)
+
+    # buf.render()='UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), arg=1, src=())'
+    alu0 = ((gidx0*2+(lidx0*128+r0*64+lidx1*8+-183)%64*64+(lidx0*128+r0*64+lidx1*8+-183)//64%32*4096+1)//4%1024)
+    alu1 = ((gidx0*2+(lidx0*128+r0*64+lidx1*8+-183)%64*64+(lidx0*128+r0*64+lidx1*8+-183)//64%32*4096+1)//4096)
+    valid = ((lidx1<7)&((((lidx0*2+r0)<3)!=1)&((lidx0*2+r0)<35)))
+    load = get_load_image_uop((32, 1024, 4), valid, (alu0, alu1))
+    self.check(load, None, "(lidx1*128+gidx0//2+144)", "(lidx0*2+r0+-3)")
+
+    # same idx, written without the inline simplification of the inner div/mod
+    alu0 = ((gidx0*2+lidx1*512+(lidx0*8192+r0*4096)+-11711)//4%1024)
+    alu1 = (lidx0*2+r0+-3)
+    valid = ((lidx1<7)&((((lidx0*2+r0)<3)!=1)&((lidx0*2+r0)<35)))
+    load = get_load_image_uop((32, 1024, 4), valid, (alu0, alu1))
+    self.check(load, None, "(lidx1*128+gidx0//2+144)", "(lidx0*2+r0+-3)")
+
+  def test_simplify8(self):
+    # from openpilot compile3, kernel r_4_16_8_16_4_4_3_3n1
+    # valid guarantees A >= 0, so divmod simplifies and gate is removed
+    gidx0 = Special("gidx0", 16)
+    gidx1 = Special("gidx1", 4)
+    lidx0 = Special("lidx0", 8)
+    lidx1 = Special("lidx1", 16)
+    A = gidx0 + gidx1*8192 + lidx0*1024 + lidx1*64 - 1040
+    valid = ((lidx1 < 1).ne(True)) & (((gidx1 + lidx0) < 1).ne(True))
+    load = get_load_image_uop((32, 1024, 4), valid, (A % 1024, A // 1024))
+    self.check(load, None, "(gidx0+lidx1*64+-16)", "(lidx0+gidx1*8+-1)")
+
+  def test_simplify9(self):
+    # from openpilot compile3, kernel r_32_16_8_4_4_7_7 (image 1x16384)
+    # valid guarantees A1 >= 0 and A1 < 512, gate should be removable
+    gidx0 = Special("gidx0", 32)
+    lidx0 = Special("lidx0", 16)
+    lidx1 = Special("lidx1", 8)
+    r0 = Range(0, 7)
+    A1 = lidx0*32 + r0*32 + lidx1*4 - 99
+    valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 19)
+    alu0 = gidx0 + (A1 % 32)*32 + (A1 // 32 % 16)*1024
+    load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(dtypes.weakint, 0)))
+    try:
+      self.check(load, None, "(gidx0+lidx0*1024+r0*1024+lidx1*128+-3168)", "0")
+    except AssertionError:
+      # TODO: fold valid
+      self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<19))",
+                       "(gidx0+lidx1*128+(lidx0*1024+r0*1024)+-3168)", "0")
+
+  def test_simplify10(self):
+    # from openpilot compile3, kernel r_16_8_4_4_4_4_7_7 (image 1x8192)
+    # valid guarantees A1 >= 0 and A1 < 128, gate should be removable
+    gidx0 = Special("gidx0", 16)
+    lidx0 = Special("lidx0", 8)
+    lidx1 = Special("lidx1", 4)
+    lidx2 = Special("lidx2", 4)
+    r0 = Range(0, 7)
+    A1 = lidx0*16 + r0*16 + lidx1*4 - 51
+    valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 11)
+    alu0 = lidx2 + gidx0*4 + (A1 % 16)*64 + (A1 // 16 % 8)*1024
+    load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(dtypes.weakint, 0)))
+    try:
+      self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0")
+    except AssertionError:
+      # TODO: fold valid
+      self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<11))",
+                       "(lidx2+gidx0*4+lidx1*256+(lidx0*1024+r0*1024)+-3264)", "0")
+
+class TestDropTrueGate(unittest.TestCase):
+  def test_drop_true_gate_on_index(self):
+    # test that INDEX with a constant True valid gets simplified to drop the valid
+    from tinygrad.codegen.late.devectorizer import load_store_indexing
+    from tinygrad.uop.ops import graph_rewrite
+    from tinygrad.uop.symbolic import sym
+    buf = UOp.param(0, dtypes.int.ptr())
+    idx = UOp.const(dtypes.weakint, 0)
+    true_gate = UOp.const(dtypes.bool, True)
+    index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx.valid(true_gate)))
+    # apply the optimization
+    result = graph_rewrite(index_with_gate, sym+load_store_indexing)
+    # the True valid should be dropped (INDEX should only have 2 sources)
+    self.assertEqual(len(result.src), 2, "True valid should be dropped from INDEX")
+
+class TestRangeShrink(unittest.TestCase):
+  def get_ranges(self, sink):
+    with Context(NOOPT=1, SPEC=0):
+      result = full_rewrite(sink)
+    return [u for u in result.toposort() if u.op is Ops.RANGE]
+
+  def test_range_shrink_single_guard(self):
+    # range 0..203 guarded by r < 4 everywhere -> shrink to 0..3
+    r = Range(0, 204)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
+    ranges = self.get_ranges(load.sink())
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 4)
+
+  def test_range_shrink_picks_max_guard(self):
+    # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8
+    r = Range(0, 204)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
+    load2 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 8), r)
+    ranges = self.get_ranges(UOp.sink(load1, load2))
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 8)
+
+  def test_range_no_shrink_guard_ge_max(self):
+    # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain)
+    r = Range(0, 204)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 300), r)
+    ranges = self.get_ranges(load.sink())
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 204)
+
+  def test_range_no_shrink_when_unguarded_elsewhere(self):
+    # one load guards r < 4, but another load uses r without a gate -> no shrink
+    r = Range(0, 204)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
+    load2 = UOp(Ops.LOAD, dtypes.float, (UOp.param(1, dtypes.float.ptr()).index(r, ptr=True),))
+    ranges = self.get_ranges(UOp.sink(load1, load2))
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 204)
+
+  def test_range_no_shrink_when_used_in_reduce(self):
+    # range used in both a gated load AND directly in the reduce expression -> no shrink
+    r = Range(0, 204)
+    gated_load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
+    red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
+    ranges = self.get_ranges(red.sink())
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 204)
+
+  def test_range_shrink_to_single_iteration(self):
+    # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely
+    r = Range(0, 204)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 1), r)
+    ranges = self.get_ranges(load.sink())
+    self.assertEqual(len(ranges), 0)
+
+  def test_range_shrink_store_where_invalid(self):
+    # emulates mask.where(x.pad_to(mask.shape), Invalid): range should shrink accordingly
+    from tinygrad.dtype import Invalid
+    r = Range(0, 204)
+    x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
+    ranges = self.get_ranges(UOp.param(0, dtypes.float.ptr()).index(r).store((r < 4).where(x, 0)).sink())
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 4)
+
+  def test_range_shrink_store_where_invalid_flipped(self):
+    # above, but flipped
+    from tinygrad.dtype import Invalid
+    r = Range(0, 204)
+    x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
+    ranges = self.get_ranges(UOp.param(0, dtypes.float.ptr()).index(r).store((r < 4).where(0, x)).sink())
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 4)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_symbolic_failures.py

@@ -0,0 +1,120 @@
+import unittest
+from tinygrad import Variable
+
+
+class TestFuzzFailure(unittest.TestCase):
+  def test_fuzz_failure1(self):
+    v1=Variable('v1', 0, 8)
+    v2=Variable('v2', 0, 2)
+    v3=Variable('v3', 0, 1)
+    expr = (((((((((((((((((((((((0//4)%2)//8)+-2)+-4)+-3)+v1)+-4)+v2)+-2)+v3)+v2)//3)%7)*1)//2)+v2)*-1)+2)+1)+0)+-3)+v3)
+    v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(0), v3.const_like(0)
+    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    self.assertEqual(num, rn)
+
+  def test_fuzz_failure2(self):
+    v1=Variable('v1', 0, 16)
+    v2=Variable('v2', 0, 5)
+    v3=Variable('v3', 0, 3)
+    expr = (((((((((((((((((((((((((0*4)//5)*2)*-1)*-2)+-4)*4)*2)*3)*4)+-4)*4)+v2)+v2)+v3)//3)+v2)+v1)//9)+3)+1)//1)+-4)//4)*2)
+    expr = (((((v1+(v2+(((v3+(v2*2))+1)//3)))+4)//9)+-57)//(9*4))
+    v1_val, v2_val, v3_val = v1.const_like(6), v2.const_like(0), v3.const_like(0)
+    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    self.assertEqual(num, rn)
+
+  def test_fuzz_failure3(self):
+    v1=Variable('v1', 0, 2)
+    v2=Variable('v2', 0, 1)
+    v3=Variable('v3', 0, 2)
+    expr = (((((((((((((((((((0//2)//3)+v3)+0)+-4)*-2)*-2)+-1)+2)+3)+v3)+0)//8)*-3)+0)*-2)*-4)*-2)//5)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0)
+    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    self.assertEqual(num, rn)
+
+  def test_fuzz_failure4(self):
+    v1=Variable('v1', 0, 2)
+    v2=Variable('v2', 0, 3)
+    v3=Variable('v3', 0, 4)
+    expr = (((((((((((((((((((((((((((((0*-2)+0)*-1)//9)//6)//8)+v1)*-4)+v2)//4)//8)+4)*3)+v1)+v3)//8)//7)+4)+v3)*-4)+1)+v1)*3)+4)*2)//5)//2)//3)*-4)
+    v1_val, v2_val, v3_val = v1.const_like(2), v2.const_like(0), v3.const_like(2)
+    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    self.assertEqual(num, rn)
+
+  def test_fuzz_failure5(self):
+    v1=Variable('v1', 0, 1)
+    v2=Variable('v2', 0, 1)
+    v3=Variable('v3', 0, 3)
+    expr = ((((((((((((((0+v2)+v1)*0)+v2)//1)//7)+-2)+v2)+v1)*4)+-3)//5)+v2)+1)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0)
+    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    self.assertEqual(num, rn)
+
+  def test_fuzz_failure6(self):
+    v1=Variable('v1', 0, 8)
+    v2=Variable('v2', 0, 64)
+    v3=Variable('v3', 0, 128)
+    expr = (((((((((((((((((((((((((((((0//3)+4)+v1)//2)+-1)//1)*1)*-1)*4)//5)+v1)//6)+v1)*-1)+-4)+v2)+-2)*-3)+v3)+-4)+-2)*-1)//8)//4)*-4)+3)+v3)*
+            -2)+v2)
+    v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(3), v3.const_like(2)
+    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    self.assertEqual(num, rn)
+
+  def test_fuzz_failure7(self):
+    v1=Variable('v1', 0, 64)
+    v2=Variable('v2', 0, 5)
+    v3=Variable('v3', 0, 128)
+    expr = (((((((((((((((((((((((((((((0+v2)*-4)+0)//9)+-4)*-2)*3)*4)//9)+v3)+v1)//4)+v1)+v3)+-1)*4)//4)+v2)//7)//3)+v1)+v2)+v3)+1)*2)//4)*3)+-1)*1)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(2), v3.const_like(65)
+    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    self.assertEqual(num, rn)
+
+  def test_fuzz_failure8(self):
+    v1=Variable('v1', 0, 2)
+    v2=Variable('v2', 0, 8)
+    v3=Variable('v3', 0, 9)
+    expr = (((((((0+-1)+2)+v1)*-2)//3)+v1)*-4)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0)
+    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    self.assertEqual(num, rn)
+
+  def test_fuzz_failure9(self):
+    v1=Variable('v1', 0, 256)
+    v2=Variable('v2', 0, 1)
+    v3=Variable('v3', 0, 8)
+    expr = (((((((((((((((((((((((((((((0*-2)//1)+3)*-2)+-3)*-4)*1)+v1)+0)%2)%8)%9)+v2)%9)+-4)//4)+-1)*-2)+0)+v1)+v1)+3)+v1)+4)+-4)+0)*2)+-3)%6)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(1), v3.const_like(0)
+    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    self.assertEqual(num, rn)
+
+  def test_fuzz_failure10(self):
+    v1=Variable("v1", 0, 256)
+    v2=Variable("v2", 0, 32)
+    v3=Variable("v3", 0, 32)
+    x5 = (v1 <= 9).where(v1 * -4 - 4, v1 // 9) // 9
+    expr = ((x5 >= -4).where(x5, (v2 % 3 + v2) // 5) * -1).maximum(((v1 * -2) % 6 + v3 % 1) * -1) * -1
+    v1_val, v2_val, v3_val = v1.const_like(9), v2.const_like(0), v3.const_like(0)
+    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    self.assertEqual(num, rn)
+
+  def test_fuzz_failure11(self):
+    v1=Variable("v1", 0, 16)
+    v2=Variable("v2", 0, 128)
+    v3=Variable("v3", 0, 5)
+    expr = (((v2 * 0).maximum(8) - v2 * 2) % 5 + v1 // 6 + v1 + 5) % 5
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(7), v3.const_like(0)
+    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    self.assertEqual(num, rn)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_symbolic_tensor.py

@@ -0,0 +1,112 @@
+import unittest
+from tinygrad import Variable
+from tinygrad.tensor import Tensor
+
+class TestSymbolic(unittest.TestCase):
+  def assert_tuple_equal(self, x, y):
+    for a,b in zip(x,y): self.assertFalse(a != b)
+
+  def test_cat_dim0_is_expanded(self):
+    i = Variable("i", 1, 5).bind(3)
+    j = Variable("j", 1, 5).bind(3)
+    k = Variable("k", 1, 5).bind(3)
+    t = Tensor.rand(5, 4)[:i].cat(Tensor.rand(5, 4)[:j], dim=0).cat(Tensor.rand(5, 4)[:k], dim=0)
+    self.assert_tuple_equal(t.shape, (i+j+k, 4))
+    t = Tensor.rand(5, 3)[:i].cat(Tensor.rand(5, 3)[:i], dim=0).cat(Tensor.rand(3, 3), dim=0)
+    self.assert_tuple_equal(t.shape, (2*i+3, 3))
+
+  def test_cat_dim1_strides(self):
+    i = Variable("i", 1, 5).bind(4)
+    j = Variable("j", 1, 5).bind(4)
+    k = Variable("k", 1, 5).bind(4)
+    t = Tensor.rand(3, 5)[:, :i].cat(Tensor.rand(3, 5)[:, :j], dim=1).cat(Tensor.rand(3, 5)[:, :k], dim=1)
+    self.assert_tuple_equal(t.shape, (3, i+j+k))
+
+class TestSymbolicVarVals(unittest.TestCase):
+  def assert_equal(self, x, y): self.assertFalse(x != y)
+
+  def test_shrink_unbind(self):
+    v = Variable("v", 1, 100)
+    bv = Variable("v", 1, 100).bind(2)
+    t = Tensor.rand(3, 4).shrink(((0,bv),(0,4)))
+    unbound_st, var_val = t.uop.unbind_all()
+    assert var_val == {v: 2}
+    t = Tensor.rand(3, 4).shrink(((bv, bv+1), (0, 4)))
+    unbound_st, var_val = t.uop.unbind_all()
+    assert var_val == {v: 2}
+
+class TestSymbolicReshape(unittest.TestCase):
+  def test_reshape(self):
+    a = Tensor.rand(5, 4)
+    b = Tensor.rand(5, 6)
+    for i in range(1, 6):
+      vi = Variable("i", 1, 5).bind(i)
+      ret = a[:vi]
+      ret = ret.reshape((vi, 4))
+      assert ret.shape == (vi, 4)
+      ret = b[:vi]
+      ret = ret.reshape((vi, 2, 3))
+      assert ret.shape == (vi, 2, 3)
+
+  def test_two_symbol_reshape(self):
+    t = Tensor.rand(5, 5)
+    for i in range(1, 6):
+      for j in range(1, 6):
+        vi = Variable("i", 1, 5).bind(i)
+        vj = Variable("j", 1, 5).bind(j)
+        ret = t[:vi, :vj]
+        ret = ret.reshape(vj, vi)
+        assert ret.shape == (vj, vi)
+        ret = ret.reshape(vi, vj)
+        assert ret.shape == (vi, vj)
+        ret = ret.reshape(1, vi*vj)
+        assert ret.shape == (1, vi*vj)
+
+class TestSymbolicExpand(unittest.TestCase):
+  def test_expand_into_symbols(self):
+    vi = Variable("i", 1, 5).bind(3)
+    vj = Variable("j", 1, 5).bind(3)
+    a = Tensor([[1], [2], [3]]).expand((3, vi))
+    assert a.shape == (3, vi)
+    a = a.reshape(3, vi, 1).expand((3, vi, vj))
+    assert a.shape == (3, vi, vj)
+
+  def test_plus_expands_constant(self):
+    a = Tensor.rand(3, 5)
+    for i in range(1, 6):
+      vi = Variable("i", 1, 5).bind(i)
+      ret = a[:, :vi]
+      ret = ret + 1
+      self.assertTupleEqual(ret.shape, (3, vi))
+
+  def test_pad_then_expand_into_symbols(self):
+    vi = Variable("i", 1, 10).bind(3)
+    a = Tensor(1).unsqueeze(0).pad((0, 24)).unsqueeze(0).expand((vi, 25))
+    self.assertEqual(a.shape, (vi, 25))
+    self.assertEqual(a.reshape(25*vi).shape, (vi*25,))
+    self.assertEqual(a.reshape(vi*25).shape, (vi*25,))
+
+class TestSymbolicShrink(unittest.TestCase):
+  def test_shrink_symbols_simple(self):
+    vi = Variable("i", 1, 5)
+    t = Tensor.rand(5, 5).shrink(((0, 5),(0,vi)))
+    assert t.shape == (5, vi)
+
+  def test_shrink_symbols(self):
+    vi = Variable("i", 1, 5)
+    t = Tensor.rand(3, 5).shrink(((0, 2), (vi, vi+1)))
+    assert t.shape == (2, 1)
+
+class TestSymbolicContiguousViewOffset(unittest.TestCase):
+  def test_shrink_from_start(self):
+    v = Variable("v", 1, 10).bind(5)
+    t = Tensor.rand(10).realize().shrink(((0, v),))
+    self.assertEqual(t.uop.contiguous_view_offset(), 0)
+
+  def test_shrink_with_offset(self):
+    v = Variable("v", 1, 7).bind(4)
+    t = Tensor.rand(10).realize().shrink(((3, 3+v),))
+    self.assertEqual(t.uop.contiguous_view_offset(), 3)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_tensor.py

@@ -0,0 +1,219 @@
+# tensor tests that pass on NULL backend (no copyout needed)
+import numpy as np
+import unittest
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.uop.ops import Ops, UOp
+from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
+from tinygrad.codegen import to_program
+from tinygrad.dtype import DType
+
+x_init = np.random.randn(1,3).astype(np.float32)
+W_init = np.random.randn(3,3).astype(np.float32)
+m_init = np.random.randn(1,3).astype(np.float32)
+
+class TestTrainMode(unittest.TestCase):
+  def test_train_mode(self):
+    assert not Tensor.training
+    @Tensor.train()
+    def f():
+      assert Tensor.training
+    f()
+    assert not Tensor.training
+
+class TestInferenceMode(unittest.TestCase):
+  def test_inference(self):
+    x = Tensor(x_init)
+    m = Tensor(m_init)
+    W = Tensor(W_init)
+    tmp = x.mul(m)
+    mm = tmp.matmul(W)
+    out = mm.relu()
+    out = out.sum()
+    #out.backward()
+    assert x.grad is None
+    assert m.grad is None
+    assert tmp.grad is None
+    assert mm.grad is None
+    assert W.grad is None
+
+  def test_no_grad_mode_context_manager(self):
+    x = Tensor(x_init)
+    m = Tensor(m_init)
+    W = Tensor(W_init)
+    def f(x, m, W):
+      tmp = x.mul(m)
+      mm = tmp.matmul(W)
+      out = mm.relu()
+      out = out.sum()
+      #out.backward()
+      assert x.grad is None
+      assert m.grad is None
+      assert tmp.grad is None
+      assert mm.grad is None
+      assert W.grad is None
+    f(x, m, W)
+
+class TestIdxUpcast(unittest.TestCase):
+  def _find_op(self, ast: UOp, op: Ops):
+    if ast.op is op: return ast
+    for src in ast.src:
+      if (ret:=self._find_op(src, op)) is not None: return ret
+  def _schedule_render(self, a: Tensor):
+    linear, _ = a.linear_with_vars()
+    for si in linear.src:
+      ast = si.src[0]
+      if ast.op is Ops.SINK:
+        renderer = Device[si.src[1].buffer.device].renderer
+        prg = to_program(ast, renderer)
+        return tuple(prg.src[2].src)
+
+  def _assert(self, dtype: DType, a: Tensor):
+    uops = self._schedule_render(a)
+    # Assert the dtype of the INDEX value, This will need be updated if UOp spec changes
+    store = next(uop for uop in uops if uop.op is Ops.STORE)
+    assert store.op is Ops.STORE
+    idx = self._find_op(store, Ops.INDEX)
+    # PTX and NIR turn Ops.INDEX into pointer arithmetic earlier than cstyle, plus it's already cast to int64
+    if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)):
+      assert idx.op is Ops.INDEX
+      idx_val = idx.src[1]
+      self.assertIs(idx_val.dtype, dtype)
+
+  # use expand to generate kernel that uses large idx
+  def do_op_then_assert(self, dtype: DType, dim1, dim2, dim3):
+    self._assert(dtype, Tensor.empty(dim1, dim2, 1).expand(-1, -1, dim3).contiguous())
+
+  @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "int64 is supported")
+  def test_overflow(self):
+    # 2**11, 2**11, 2**11 -> 2**33 will overflow when indexed
+    self.do_op_then_assert(dtypes.long, 2048, 2048, 2048)
+
+  @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "int64 is supported")
+  def test_overflow_sym(self):
+    self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 1, 2048).bind(32))
+
+  def test_regular(self):
+    self.do_op_then_assert(dtypes.int, 64, 64, 64)
+
+  def test_regular_sym(self):
+    self.do_op_then_assert(dtypes.int, 256, 256, UOp.variable("dim3", 1, 64).bind(32))
+
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "PTX and NIR always converts Ops.INDEX to int64")
+  def test_symfold(self):
+    # This would cause an overflow, but after sym fold it's within int32
+    a = Tensor.arange(65535)
+    uops = self._schedule_render(a)
+    assert all(uop.dtype is not dtypes.long for uop in uops)
+
+  @unittest.skipIf(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "int64 is supported")
+  def test_int64_unsupported_overflow_sym(self):
+    with self.assertRaises((KeyError, RuntimeError)):
+      self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 1, 2048).bind(32))
+
+  @unittest.skipIf(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "int64 is supported")
+  @unittest.expectedFailure  # bug in gpu dims limiting
+  def test_int64_unsupported_overflow(self):
+    with self.assertRaises((KeyError, RuntimeError)):
+      self.do_op_then_assert(dtypes.long, 2048, 2048, 2048)
+
+  @unittest.skip("This is kept for reference, it requires large memory to run")
+  def test_overflow_kernel_run(self):
+    # This creates a total of 2**31+10 elements, requiring at least 2147 MB memory to run
+    # Modified example from issue 3271
+    a = Tensor.empty(2**11, 2**11, 1, dtype=dtypes.int8).permute((2, 0, 1)).expand((2**9+10, -1, -1)).contiguous()
+    a.realize()
+
+class TestTensorUnique(unittest.TestCase):
+  def test_empty_bufs_unique(self):
+    a = Tensor.empty(10, 10).contiguous()
+    b = Tensor.empty(10, 10).contiguous()
+    Tensor.realize(a,b)
+    self.assertIsNot(a.uop.buffer, b.uop.buffer)
+
+  def test_zeros_bufs_unique_sep(self):
+    a = Tensor.zeros(10, 10).contiguous()
+    Tensor.realize(a)
+    b = Tensor.zeros(10, 10).contiguous()
+    Tensor.realize(b)
+    self.assertIsNot(a.uop.buffer, b.uop.buffer)
+
+  def test_zeros_bufs_unique(self):
+    a = Tensor.zeros(10, 10).contiguous()
+    b = Tensor.zeros(10, 10).contiguous()
+    Tensor.realize(a,b)
+    self.assertIsNot(a.uop.buffer, b.uop.buffer)
+
+  def test_times_2_not_unique(self):
+    a = Tensor.zeros(10, 10).contiguous()
+    b = a * 2
+    c = a * 2
+    Tensor.realize(b,c)
+    self.assertIs(b.uop.buffer, c.uop.buffer)
+
+class TestRand(unittest.TestCase):
+  def test_rand_large_tensor(self):
+    # large tensor rand (num > uint32.max) should not crash in frontend
+    Tensor.manual_seed(0)
+    Tensor.rand(2**17, 2**17).schedule_linear()
+    Tensor.rand(2**17, 2**17).schedule_linear()
+    Tensor.rand(2**17, 2**17).schedule_linear()
+
+class TestTensorConstLike(unittest.TestCase):
+  def test_const_like_shape(self):
+    t = Tensor.ones(3, 4)
+    c = t.const_like(0)
+    self.assertEqual(c.shape, (3, 4))
+    self.assertEqual(c.dtype, t.dtype)
+
+  def test_const_like_multi_device(self):
+    devs = ("NULL:0", "NULL:1")
+    t = Tensor.ones(8, 4).shard(devs, axis=0)
+    c = t.const_like(5)
+    self.assertEqual(c.shape, (8, 4))
+    self.assertEqual(c.device, t.device)
+    self.assertEqual(c.uop.axis, 0)
+
+  def test_full_like_device_on_multi_raises(self):
+    t = Tensor.ones(8, 4).shard(("NULL:0", "NULL:1"), axis=0)
+    with self.assertRaises(RuntimeError): t.full_like(5, device="NULL")
+
+class TestTensorDevice(unittest.TestCase):
+  def test_create_from_single_device_tuple(self):
+    (Tensor([1.0], device=(Device.DEFAULT,)) + Tensor([2.0])).realize()
+
+class TestTensorPad(unittest.TestCase):
+  # padding int tensor with float-only value (like -inf) must promote dtype to fit value
+  def test_pad_int_with_neg_inf(self):
+    t = Tensor.arange(9).reshape(1, 1, 3, 3)
+    self.assertEqual(t.dtype, dtypes.int)
+    r = t.pad((1, 2, 0, -1), value=-float('inf'))
+    self.assertEqual(r.dtype, dtypes.float)
+    self.assertEqual(r.shape, (1, 1, 2, 6))
+
+class TestTensorDeviceMismatch(unittest.TestCase):
+  def test_gather(self):
+    x = Tensor.empty(3, 4, device="NULL")
+    idx = Tensor.zeros(3, 4, dtype=dtypes.int32, device="NULL:1")
+    with self.assertRaises(RuntimeError): x.gather(0, idx)
+  def test_scatter_index(self):
+    x = Tensor.zeros(3, 4, device="NULL")
+    idx = Tensor.zeros(3, 4, dtype=dtypes.int32, device="NULL:1")
+    src = Tensor.ones(3, 4, device="NULL")
+    with self.assertRaises(RuntimeError): x.scatter(0, idx, src)
+  def test_scatter_src(self):
+    x = Tensor.zeros(3, 4, device="NULL")
+    idx = Tensor.zeros(3, 4, dtype=dtypes.int32, device="NULL")
+    src = Tensor.ones(3, 4, device="NULL:1")
+    with self.assertRaises(RuntimeError): x.scatter(0, idx, src)
+  def test_getitem_tensor_index(self):
+    x = Tensor.empty(4, 5, device="NULL")
+    idx = Tensor([0, 1], dtype=dtypes.int32, device="NULL:1")
+    with self.assertRaises(RuntimeError): x[idx]
+  def test_sparse_categorical_crossentropy(self):
+    x = Tensor.zeros(2, 3, device="NULL")
+    Y = Tensor([0, 1], dtype=dtypes.int32, device="NULL:1")
+    with self.assertRaises(RuntimeError): x.sparse_categorical_crossentropy(Y)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_tensor_io.py

@@ -0,0 +1,24 @@
+import unittest
+from tinygrad import Tensor, dtypes
+from tinygrad.nn.state import TensorIO
+
+class TestTensorIO(unittest.TestCase):
+  def test_create(self):
+    with self.assertRaises(ValueError):
+      TensorIO(Tensor(b"Hello World").reshape(1, -1))
+    with self.assertRaises(ValueError):
+      TensorIO(Tensor([], dtype=dtypes.int64).reshape(1, -1))
+
+  def test_seek(self):
+    t = Tensor(b"Hello World!")
+    fobj = TensorIO(t)
+    self.assertEqual(fobj.tell(), 0)
+    self.assertEqual(fobj.seek(1), 1)
+    self.assertEqual(fobj.seek(-2, 2), len(t) - 2)
+    self.assertEqual(fobj.seek(1, 1), len(t) - 1)
+    self.assertEqual(fobj.seek(10, 1), len(t))
+    self.assertEqual(fobj.seek(10, 2), len(t))
+    self.assertEqual(fobj.seek(-10, 0), 0)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_tensor_metadata.py

@@ -0,0 +1,127 @@
+import unittest
+from tinygrad import Tensor, dtypes
+from tinygrad.tensor import _METADATA
+from tinygrad.engine.realize import capturing
+from tinygrad.helpers import Context
+
+@unittest.skip("tensor metadata is no longer supported")
+class TestTensorMetadata(unittest.TestCase):
+  def setUp(self) -> None:
+    _METADATA.set(None)
+    self._ctx = Context(SCACHE=0)
+    self._ctx.__enter__()
+  def tearDown(self) -> None:
+    self._ctx.__exit__(None, None, None)
+
+  @unittest.skip("why would this be true?")
+  def test_exclude_noop_metadata(self):
+    a = Tensor.rand(4, 4)*1
+    self.assertEqual(a.uop.metadata[0].name, "__mul__")
+    k = a.schedule_linear().src[-1]
+    self.assertEqual([m.name for m in k.arg.metadata], ["rand"])
+
+  @unittest.skip("metadata not reaching kernel schedule")
+  def test_exclude_const_metadata(self):
+    a = Tensor.arange(4)
+    b = Tensor.full((4,), -1, dtype=dtypes.int).contiguous()
+    sched = a.schedule_linear(b)
+    self.assertEqual([m.name for m in sched.src[0].arg.metadata], ["arange"])
+    self.assertEqual([m.name for m in sched.src[1].arg.metadata], ["contiguous"])
+
+  def test_matmul(self):
+    x = Tensor.rand(3)
+    W = Tensor.rand(3, 3)
+    out = x.matmul(W)
+    self.assertEqual(out.uop.metadata[0].name, "matmul")
+    si = out.schedule_linear().src[-1]
+    self.assertEqual(len(si.arg.metadata), 1)
+    self.assertEqual(si.arg.metadata[0].name, "matmul")
+
+  def test_relu(self):
+    x = Tensor.rand(3)
+    out = x.relu()
+    self.assertEqual(out.uop.metadata[0].name, "relu")
+    si = out.schedule_linear().src[-1]
+    self.assertEqual(len(si.arg.metadata), 1)
+    self.assertEqual(si.arg.metadata[0].name, "relu")
+
+  @unittest.skip("assign metadata no longer captured")
+  def test_assign(self):
+    x = Tensor.empty(10, 10).realize()
+    x.assign(Tensor.ones(10, 10).contiguous())
+    si = x.schedule_linear().src[-1]
+    self.assertEqual(len(si.arg.metadata), 1)
+    self.assertEqual(si.arg.metadata[0].name, "assign")
+
+  def test_complex(self):
+    x = Tensor.rand(3)
+    y = Tensor.rand(3)
+    out = x.relu() * y.sigmoid()
+    self.assertEqual(out.uop.metadata[0].name, "__mul__")
+    self.assertEqual(out.uop.src[0].metadata[0].name, "relu")
+    self.assertEqual(out.uop.src[1].metadata[0].name, "sigmoid")
+    si = out.schedule_linear().src[-1]
+    self.assertEqual(len(si.arg.metadata), 3)
+    self.assertEqual(set(m.name for m in si.arg.metadata), {"relu", "sigmoid", "__mul__"})
+
+  @unittest.skip("flaky")
+  def test_complex_backward(self):
+    x = Tensor.rand(3).realize()
+    y = Tensor.rand(3).realize()
+    out = (x.relu() * y.sigmoid()).sum()
+    self.assertEqual(out.uop.metadata[0].name, "sum")
+    out.backward()
+    self.assertEqual(x.grad.uop.metadata[0].name, "relu")
+    #self.assertTrue(x.grad.uop.metadata[0].backward)  # TODO: backward flag is False
+    self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid")
+    #self.assertTrue(y.grad.uop.metadata[0].backward)  # TODO: backward flag is False
+    si = out.schedule_linear(x.grad, y.grad).src[-1]
+    #self.assertEqual(len(si.arg.metadata), 3, f"failed with {si.arg.metadata}")
+    # skip numpy, this is schedule cache
+    self.assertSetEqual(set(m.name for m in si.arg.metadata if m.name != "numpy"), {"sigmoid", "relu"})
+    #bw = [m for m in si.metadata if m.backward]
+    #self.assertEqual(len(bw), 1)
+    #self.assertEqual(bw[0].name, "sigmoid")
+
+  def test_tracemeta_0(self):
+    with Context(TRACEMETA=0):
+      x = Tensor.rand(3)
+      y = Tensor.rand(3)
+      out = (x.relu() * y.sigmoid()).sum()
+      self.assertIsNone(out.uop.metadata)
+      self.assertIsNone(out.uop.src[0].metadata)
+      si = out.schedule_linear().src[-1]
+      self.assertEqual(si.arg.metadata, ())
+
+  def _has_metadata(self, h, name):
+    linears = []
+    capturing.append(type("", (), {"add_linear": lambda _, linear, var_vals: linears.append(linear)})())
+    try: h.realize()
+    finally: capturing.clear()
+    calls = [call for linear in linears for call in linear.src]
+    return any(m.name == name for call in calls for m in call.arg.metadata)
+
+  def test_metadata_survives_realize_pending_assign(self):
+    shared = Tensor.rand(4)
+    c = Tensor.zeros(8).contiguous().realize()
+    c[:4].assign(shared)
+    self.assertTrue(self._has_metadata(c[:4].relu(), "relu"))
+
+  @unittest.expectedFailure
+  def test_metadata_lost_realize_pending_assign(self):
+    shared = Tensor.rand(4)
+    c = Tensor.zeros(8).contiguous().realize()
+    c[:4].assign(shared)
+    self.assertTrue(self._has_metadata((c[:4] + shared).relu(), "relu"))
+
+class TestTraceMetaShutdown(unittest.TestCase):
+  def test_tracemeta_del_no_shutdown_error(self):
+    import subprocess, os
+    result = subprocess.run(['python3', '-c', 'from tinygrad import Tensor\n'
+                             'x=Tensor.eye(3); (x@x).sum().backward()'],
+                            env={**os.environ, "TRACEMETA": "2"}, capture_output=True)
+    self.assertEqual(result.returncode, 0)
+    self.assertNotIn(b"Exception", result.stderr)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_tensor_uop_mixin.py

@@ -0,0 +1,433 @@
+import math, unittest
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite
+
+_strip_unique_pm = PatternMatcher([
+  (UPat((Ops.UNIQUE, Ops.LUNIQUE), name="u"), lambda u: u.replace(arg=0) if u.arg != 0 else None),
+])
+def _strip_unique(u: UOp) -> UOp: return graph_rewrite(u, _strip_unique_pm)
+
+def _t(*shape):
+  return Tensor.arange(math.prod(shape)).reshape(*shape)
+
+# Tensor().func().uop should be the same as UOp.func()
+def _check(tc: unittest.TestCase, t: Tensor, fn):
+  tc.assertIs(fn(t).uop, fn(t.uop), f"\ntensor.uop = {fn(t).uop}\nuop = {fn(t.uop)}")
+
+class TestTensorUOpBinop(unittest.TestCase):
+  # Tensor's binop upcasts mixed dtypes via least_upper_dtype + explicit CAST; UOp should match.
+  def test_mul_float_int(self):
+    t = _t(3).float()
+    self.assertIs(_strip_unique((t * Tensor.arange(3)).uop), _strip_unique(t.uop * UOp.arange(3)))
+  def test_mul_bool_int(self):
+    t = _t(3)
+    self.assertIs(_strip_unique((t.eq(1) * Tensor.arange(3)).uop), _strip_unique(t.uop.eq(1) * UOp.arange(3)))
+  # Tensor's ufix picks float dtype when scalar is float and self is int; UOp should match.
+  def test_add_scalar_float_on_int(self): _check(self, _t(3), lambda x: x + 1.5)
+  # div: Tensor.div (default case) delegates to ElementwiseMixin.div; trees must match for Tensor and UOp.
+  def test_div_tensor_by_tensor(self):
+    a, b = _t(4).float(), _t(4).float() + 1
+    self.assertIs(_strip_unique((a/b).uop), _strip_unique(a.uop/b.uop))
+  def test_div_int_by_int(self):                 _check(self, _t(4), lambda x: x / 3)
+  def test_div_sum_by_sum(self):                 _check(self, _t(4).float(), lambda x: x.sum() / (x + 1).sum())
+  def test_div_broadcast_tensor_by_tensor(self):
+    a, b = _t(3, 4).float(), _t(4).float() + 1
+    self.assertIs(_strip_unique((a/b).uop), _strip_unique(a.uop/b.uop))
+  # isclose used `self == other` which is Python identity on UOp (not elementwise); now uses .eq().
+  def test_isclose(self):
+    t = _t(4).float()
+    self.assertIs(_strip_unique(t.isclose(t).uop), _strip_unique(t.uop.isclose(t.uop)))
+  # __floordiv__/mod/fmod and div(rounding_mode=...) dispatch on dtype in mixin
+  def test_floordiv_int(self):   _check(self, _t(4), lambda x: x // 3)
+  def test_floordiv_float(self): _check(self, _t(4).float() + 1.5, lambda x: x // 2.0)
+  def test_rfloordiv_int(self):  _check(self, _t(4)+1, lambda x: 7 // x)
+  def test_mod_int(self):        _check(self, _t(4), lambda x: x % 3)
+  def test_mod_float(self):      _check(self, _t(4).float() + 1.5, lambda x: x % 2.0)
+  def test_div_trunc_int(self):  _check(self, _t(4), lambda x: x.div(3, rounding_mode="trunc"))
+  def test_div_trunc_float(self):_check(self, _t(4).float() + 1.5, lambda x: x.div(2.0, rounding_mode="trunc"))
+  def test_fmod_int(self):       _check(self, _t(4), lambda x: x.fmod(3))
+  def test_fmod_float(self):     _check(self, _t(4).float() + 1.5, lambda x: x.fmod(2.0))
+  def test_floordiv_bool(self):  _check(self, _t(4).cast(dtypes.bool), lambda x: x // True)
+  def test_mod_bool(self):       _check(self, _t(4).cast(dtypes.bool), lambda x: x % True)
+  def test_fmod_bool(self):      _check(self, _t(4).cast(dtypes.bool), lambda x: x.fmod(True))
+
+class TestTensorUOpClone(unittest.TestCase):
+  def test_clone(self):
+    t = _t(3, 4).float()
+    self.assertIs(_strip_unique(t.clone().uop), _strip_unique(t.uop.clone()))
+  def test_clone_deviceless_const(self):
+    u = UOp.const(dtypes.float, 2.0)
+    self.assertIs(_strip_unique(Tensor(u).clone().uop), _strip_unique(u.clone()))
+
+class TestTensorUOpGetitem(unittest.TestCase):
+  # ---- pure slice patterns ----
+  def test_slice_full(self):           _check(self, _t(4), lambda x: x[slice(None)])
+  def test_slice_positive(self):       _check(self, _t(8), lambda x: x[1:5])
+  def test_slice_open_start(self):     _check(self, _t(8), lambda x: x[:5])
+  def test_slice_open_stop(self):      _check(self, _t(8), lambda x: x[3:])
+  def test_slice_negative_start(self): _check(self, _t(8), lambda x: x[-3:])
+  def test_slice_negative_stop(self):  _check(self, _t(8), lambda x: x[:-2])
+  def test_slice_both_negative(self):  _check(self, _t(8), lambda x: x[-5:-1])
+
+  # ---- slice with stride ----
+  def test_slice_stride(self):                  _check(self, _t(6), lambda x: x[::2])
+  def test_slice_start_stop_stride(self):       _check(self, _t(6), lambda x: x[1:5:2])
+  def test_slice_reverse(self):                 _check(self, _t(6), lambda x: x[::-1])
+  def test_slice_singleton_negative_step(self): _check(self, _t(8), lambda x: x[3:2:-1])
+
+  # ---- empty / out-of-bounds slice ----
+  def test_slice_empty(self):    _check(self, _t(6), lambda x: x[3:1])
+  def test_slice_oob_stop(self): _check(self, _t(6), lambda x: x[0:100])
+
+  # ---- single int (reduces a dim) ----
+  def test_int_positive(self): _check(self, _t(8), lambda x: x[3])
+  def test_int_negative(self): _check(self, _t(8), lambda x: x[-1])
+
+  # ---- ellipsis ----
+  def test_ellipsis_only(self):       _check(self, _t(2, 3, 4), lambda x: x[...])
+  def test_ellipsis_then_int(self):   _check(self, _t(2, 3, 4), lambda x: x[..., -1])
+  def test_ellipsis_then_slice(self): _check(self, _t(2, 3, 4), lambda x: x[..., 1:3])
+  def test_ellipsis_then_none(self):  _check(self, _t(2, 3), lambda x: x[..., None])
+
+  # ---- None (unsqueeze) ----
+  def test_none_front(self):    _check(self, _t(4), lambda x: x[None])
+  def test_none_back(self):     _check(self, _t(4), lambda x: x[:, None])
+  def test_none_middle(self):   _check(self, _t(2, 3), lambda x: x[:, None, :])
+  def test_multiple_none(self): _check(self, _t(2, 3), lambda x: x[None, :, None])
+
+  # ---- mixed multi-dim ----
+  def test_int_then_slice(self):    _check(self, _t(2, 3), lambda x: x[1, :])
+  def test_multi_int(self):         _check(self, _t(2, 3, 4), lambda x: x[1, 2])
+  def test_mixed_slice_int(self):   _check(self, _t(2, 3, 4), lambda x: x[0:2, -1, 1:3])
+  def test_mixed_slice_slice(self): _check(self, _t(3, 4, 5), lambda x: x[1:3, :, 0:2])
+  def test_high_rank_combo(self):   _check(self, _t(4, 5, 6), lambda x: x[1:3, :, -1, None])
+
+class TestTensorUOpCumalu(unittest.TestCase):
+  def test_cumsum_1d(self):       _check(self, _t(5), lambda x: x.cumsum())
+  def test_cumsum_2d(self):       _check(self, _t(3, 4), lambda x: x.cumsum(1))
+  def test_cumsum_non_last(self): _check(self, _t(3, 4), lambda x: x.cumsum(0))
+  def test_cumsum_large(self):    _check(self, _t(600), lambda x: x.cumsum())  # exercises _split_cumalu
+  def test_cumprod(self):         _check(self, _t(4), lambda x: x.cumprod(0))
+
+class TestTensorUOpCumMinMax(unittest.TestCase):
+  def _check_pair(self, t, fn):
+    vt, it = fn(t)
+    vu, iu = fn(t.uop)
+    self.assertIs(_strip_unique(vt.uop), _strip_unique(vu))
+    self.assertIs(_strip_unique(it.uop), _strip_unique(iu))
+  def test_cummax_1d(self):    self._check_pair(_t(5), lambda x: x.cummax(0))
+  def test_cummax_2d(self):    self._check_pair(_t(3, 4), lambda x: x.cummax(1))
+  def test_cummax_0d(self):    self._check_pair(_t(1).reshape(()), lambda x: x.cummax(0))
+  def test_cummin_1d(self):    self._check_pair(_t(5), lambda x: x.cummin(0))
+  def test_cummin_2d(self):    self._check_pair(_t(3, 4), lambda x: x.cummin(1))
+
+class TestTensorUOpArgMinMax(unittest.TestCase):
+  def _check_stripped(self, t, fn): self.assertIs(_strip_unique(fn(t).uop), _strip_unique(fn(t.uop)))
+  def test_argmax(self):       self._check_stripped(_t(3, 4), lambda x: x.argmax(axis=1))
+  def test_argmax_flat(self):  self._check_stripped(_t(3, 4), lambda x: x.argmax())
+  def test_argmin(self):       self._check_stripped(_t(3, 4), lambda x: x.argmin(axis=0))
+
+class TestTensorUOpSequential(unittest.TestCase):
+  def test_sequential(self): _check(self, _t(4), lambda x: x.sequential([lambda y: y * 2, lambda y: y + 1]))
+
+class TestTensorUOpOneHot(unittest.TestCase):
+  def test_one_hot(self):
+    t = _t(5)
+    self.assertIs(_strip_unique(t.one_hot(5).uop), _strip_unique(t.uop.one_hot(5)))
+
+class TestTensorUOpSort(unittest.TestCase):
+  def _check(self, t, **kw):
+    tv, ti = t.sort(**kw)
+    uv, ui = t.uop.sort(**kw)
+    self.assertIs(_strip_unique(tv.uop), _strip_unique(uv))
+    self.assertIs(_strip_unique(ti.uop), _strip_unique(ui))
+  def test_sort_1d(self):         self._check(Tensor([0.5, 0.1, 0.3]).float())
+  def test_sort_descending(self): self._check(Tensor([0.5, 0.1, 0.3]).float(), descending=True)
+  def test_sort_2d(self):         self._check(_t(2, 4).float())
+  def test_sort_single(self):     self._check(Tensor([1.0]).float())
+  def test_argsort(self):
+    t = Tensor([0.5, 0.1, 0.3]).float()
+    self.assertIs(_strip_unique(t.argsort().uop), _strip_unique(t.uop.argsort()))
+  def test_topk(self):
+    t = _t(2, 4).float()
+    tv, ti = t.topk(2)
+    uv, ui = t.uop.topk(2)
+    self.assertIs(_strip_unique(tv.uop), _strip_unique(uv))
+    self.assertIs(_strip_unique(ti.uop), _strip_unique(ui))
+
+class TestTensorUOpAllclose(unittest.TestCase):
+  def test_allclose(self):
+    a, b = _t(4).float(), _t(4).float()
+    self.assertIs(_strip_unique(a.allclose(b).uop), _strip_unique(a.uop.allclose(b.uop)))
+
+class TestTensorUOpBitcast(unittest.TestCase):
+  def test_bitcast_same_dtype(self): _check(self, _t(4).float(), lambda x: x.bitcast(dtypes.float32))
+
+class TestTensorUOpRand(unittest.TestCase):
+  def test_random_bits(self):
+    k = UOp.empty((2,), dtype=dtypes.uint32)
+    c = UOp.zeros(2, dtype=dtypes.uint32)
+    for num in (1, 4, 7, 1024):
+      self.assertIs(_strip_unique(Tensor.random_bits(Tensor(k), Tensor(c), num).uop),
+                    _strip_unique(UOp.random_bits(k, c, num)))
+  def test_bits_to_rand_float32(self):
+    bits_uop = UOp.empty((8,), dtype=dtypes.uint32)
+    for shape in ((8,), (2, 4), (5,)):
+      self.assertIs(_strip_unique(Tensor._bits_to_rand(Tensor(bits_uop), shape, dtypes.float32).uop),
+                    _strip_unique(UOp._bits_to_rand(bits_uop, shape, dtypes.float32)))
+
+class TestTensorUOpGather(unittest.TestCase):
+  def _check(self, t, dim, idx):
+    self.assertIs(_strip_unique(t.gather(dim, idx).uop), _strip_unique(t.uop.gather(dim, idx.uop)))
+  def test_gather_1d(self):  self._check(_t(5), 0, Tensor([2, 1, 0, 1, 2], dtype=dtypes.int32))
+  def test_gather_dim0(self): self._check(_t(3, 4), 0, Tensor([[0, 1, 2, 0], [1, 2, 0, 1], [2, 0, 1, 2]], dtype=dtypes.int32))
+  def test_gather_dim1(self): self._check(_t(3, 4), 1, Tensor([[0, 1, 2, 3], [1, 2, 3, 0], [2, 3, 0, 1]], dtype=dtypes.int32))
+
+class TestTensorUOpInterpolate(unittest.TestCase):
+  def _check(self, t, mode):
+    self.assertIs(_strip_unique(t.interpolate(size=(2, 2), mode=mode).uop),
+                  _strip_unique(t.uop.interpolate(size=(2, 2), mode=mode)))
+  def test_interpolate_nearest(self):       self._check(_t(1, 1, 4, 4).float(), "nearest")
+  def test_interpolate_nearest_exact(self): self._check(_t(1, 1, 4, 4).float(), "nearest-exact")
+  def test_interpolate_linear(self):        self._check(_t(1, 1, 4, 4).float(), "linear")
+
+class TestTensorUOpLoss(unittest.TestCase):
+  def test_cross_entropy(self):
+    t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(t.cross_entropy(Y).uop), _strip_unique(t.uop.cross_entropy(Y.uop)))
+  def test_sparse_categorical_crossentropy(self):
+    t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(t.sparse_categorical_crossentropy(Y).uop), _strip_unique(t.uop.sparse_categorical_crossentropy(Y.uop)))
+  def test_sparse_categorical_crossentropy_ignore_index(self):
+    t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(t.sparse_categorical_crossentropy(Y, ignore_index=0).uop),
+                  _strip_unique(t.uop.sparse_categorical_crossentropy(Y.uop, ignore_index=0)))
+  def test_nll_loss(self):
+    t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(t.nll_loss(Y).uop), _strip_unique(t.uop.nll_loss(Y.uop)))
+  def test_nll_loss_weight(self):
+    t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float()
+    self.assertIs(_strip_unique(t.nll_loss(Y, weight=w).uop), _strip_unique(t.uop.nll_loss(Y.uop, weight=w.uop)))
+  def test_nll_loss_ignore_index(self):
+    t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(t.nll_loss(Y, ignore_index=1).uop), _strip_unique(t.uop.nll_loss(Y.uop, ignore_index=1)))
+  def test_nll_loss_none_reduction(self):
+    t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(t.nll_loss(Y, reduction="none").uop), _strip_unique(t.uop.nll_loss(Y.uop, reduction="none")))
+  def test_nll_loss_weight_ignore_index(self):
+    t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float()
+    self.assertIs(_strip_unique(t.nll_loss(Y, weight=w, ignore_index=1).uop),
+                  _strip_unique(t.uop.nll_loss(Y.uop, weight=w.uop, ignore_index=1)))
+
+class TestTensorUOpScatter(unittest.TestCase):
+  def test_scatter(self):
+    x, idx, src = _t(3, 4).float(), Tensor([[0, 1, 2, 0]], dtype=dtypes.int32), _t(1, 4).float()
+    self.assertIs(_strip_unique(x.scatter(0, idx, src).uop), _strip_unique(x.uop.scatter(0, idx.uop, src.uop)))
+  def test_scatter_scalar_src(self):
+    x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(x.scatter(1, idx, 3.14).uop), _strip_unique(x.uop.scatter(1, idx.uop, 3.14)))
+  # inf cannot be cast to int — this regresses if scalar src is routed through index.dtype first
+  def test_scatter_inf_src(self):
+    x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(x.scatter(1, idx, float("inf")).uop),
+                  _strip_unique(x.uop.scatter(1, idx.uop, float("inf"))))
+  def test_scatter_add(self):
+    x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(x.scatter(1, idx, 3.14, reduce="add").uop),
+                  _strip_unique(x.uop.scatter(1, idx.uop, 3.14, reduce="add")))
+  def test_scatter_multiply(self):
+    x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(x.scatter(1, idx, 3.14, reduce="multiply").uop),
+                  _strip_unique(x.uop.scatter(1, idx.uop, 3.14, reduce="multiply")))
+  # tensor src with reduce hits the "elif reduce: raise" branch in both Tensor and UOp paths
+  def test_scatter_tensor_src_with_reduce_raises(self):
+    x, idx, src = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32), _t(1, 2).float()
+    with self.assertRaises(TypeError): x.scatter(1, idx, src, reduce="add")
+    with self.assertRaises(TypeError): x.uop.scatter(1, idx.uop, src.uop, reduce="add")
+
+class TestTensorUOpScatterReduce(unittest.TestCase):
+  def _check(self, x, idx, src, **kw):
+    self.assertIs(_strip_unique(x.scatter_reduce(0, idx, src, **kw).uop),
+                  _strip_unique(x.uop.scatter_reduce(0, idx.uop, src.uop, **kw)))
+  def test_sum(self):  self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="sum")
+  def test_prod(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="prod")
+  def test_mean(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="mean")
+  def test_amax(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="amax")
+  def test_amin(self): self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="amin")
+  def test_mean_exclude_self(self):
+    self._check(_t(3, 4).float(), Tensor([[0, 1, 0, 1]]*3, dtype=dtypes.int32), Tensor.ones(3, 4).float(), reduce="mean", include_self=False)
+
+class TestTensorUOpPool(unittest.TestCase):
+  def test_avg_pool2d(self):                _check(self, _t(1, 1, 5, 5).float(), lambda x: x.avg_pool2d())
+  def test_avg_pool2d_padding(self):        _check(self, _t(1, 1, 5, 5).float(), lambda x: x.avg_pool2d(padding=1))
+  def test_avg_pool2d_ceil(self):           _check(self, _t(1, 1, 5, 5).float(), lambda x: x.avg_pool2d(ceil_mode=True))
+  def test_avg_pool2d_no_count_pad(self):   _check(self, _t(1, 1, 5, 5).float(), lambda x: x.avg_pool2d(padding=1, count_include_pad=False))
+  def test_max_pool2d(self):                _check(self, _t(1, 1, 5, 5).float(), lambda x: x.max_pool2d())
+  def test_max_pool2d_padding(self):        _check(self, _t(1, 1, 5, 5).float(), lambda x: x.max_pool2d(padding=1))
+  def test_max_pool2d_ceil(self):           _check(self, _t(1, 1, 5, 5).float(), lambda x: x.max_pool2d(ceil_mode=True))
+  def test_max_pool2d_return_indices(self):
+    t = _t(1, 1, 5, 5).float()
+    vt, it = t.max_pool2d(return_indices=True)
+    vu, iu = t.uop.max_pool2d(return_indices=True)
+    self.assertIs(_strip_unique(vt.uop), _strip_unique(vu))
+    self.assertIs(_strip_unique(it.uop), _strip_unique(iu))
+  def test_max_unpool2d(self):
+    t = _t(1, 1, 4, 4).float()
+    out, idx = t.max_pool2d(return_indices=True)
+    self.assertIs(_strip_unique(out.max_unpool2d(idx).uop), _strip_unique(out.uop.max_unpool2d(idx.uop)))
+
+class TestTensorUOpCat(unittest.TestCase):
+  def test_cat_dim0(self):     _check(self, _t(2, 3), lambda x: x.cat(x, dim=0))
+  def test_cat_dim1(self):     _check(self, _t(2, 3), lambda x: x.cat(x, dim=1))
+  def test_cat_3tensors(self): _check(self, _t(2, 3), lambda x: x.cat(x, x, dim=0))
+  def test_cat_neg_dim(self):  _check(self, _t(2, 3, 4), lambda x: x.cat(x, dim=-1))
+
+class TestTensorUOpPad(unittest.TestCase):
+  def test_pad_flat(self):               _check(self, _t(4, 5), lambda x: x.pad((1, 2, 0, 3)))
+  def test_pad_flat_negative(self):      _check(self, _t(4, 5), lambda x: x.pad((1, -1, 0, 2), value=-1.0))
+  def test_pad_grouped_none(self):       _check(self, _t(4, 5), lambda x: x.pad((None, (0, 3))))
+  def test_pad_circular(self):           _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="circular"))
+  def test_pad_circular_zero_after(self):_check(self, _t(4, 5), lambda x: x.pad(((1, 0), (2, 0)), mode="circular"))
+  def test_pad_reflect(self):            _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="reflect"))
+  def test_pad_reflect_negative(self):   _check(self, _t(4, 5), lambda x: x.pad(((1, -1), (0, 2)), mode="reflect"))
+  def test_pad_replicate(self):          _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="replicate"))
+  def test_pad_replicate_negative(self): _check(self, _t(4, 5), lambda x: x.pad(((1, -1), (0, 2)), mode="replicate"))
+
+class TestTensorUOpStack(unittest.TestCase):
+  def test_stack_dim0(self):     _check(self, _t(2, 3), lambda x: x.stack(x, dim=0))
+  def test_stack_dim1(self):     _check(self, _t(2, 3), lambda x: x.stack(x, dim=1))
+  def test_stack_3tensors(self): _check(self, _t(2, 3), lambda x: x.stack(x, x, dim=0))
+  def test_stack_new_last(self): _check(self, _t(2, 3), lambda x: x.stack(x, dim=-1))
+
+class TestTensorUOpConv2d(unittest.TestCase):
+  def test_conv2d_basic(self):
+    w = _t(1, 1, 2, 2).float()
+    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
+  def test_conv2d_padded(self):
+    w = _t(1, 1, 2, 2).float()
+    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, padding=1))
+  def test_conv2d_negative_padding(self):
+    w = _t(1, 1, 3, 3).float()
+    _check(self, _t(1, 1, 5, 5).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, padding=(-1,-1,-1,-1)))
+  def test_conv2d_multichannel_bias(self):
+    w, b = _t(4, 2, 3, 3).float(), _t(4).float()
+    _check(self, _t(2, 2, 5, 5).float(), lambda x: x.conv2d(*(y if isinstance(x, Tensor) else y.uop for y in (w, b))))
+  def test_conv2d_stride_dilation(self):
+    w = _t(2, 2, 2, 2).float()
+    _check(self, _t(1, 2, 6, 6).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, stride=2, dilation=2))
+  def test_conv2d_groups(self):
+    w = _t(4, 1, 2, 2).float()
+    _check(self, _t(1, 4, 4, 4).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, groups=4))
+  def test_conv2d_3d(self):
+    w = _t(1, 1, 2, 2, 2).float()
+    _check(self, _t(1, 1, 3, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
+  def test_conv_transpose2d_basic(self):
+    w = _t(1, 1, 2, 2).float()
+    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv_transpose2d(w if isinstance(x, Tensor) else w.uop))
+  def test_conv_transpose2d_stride(self):
+    w = _t(1, 1, 2, 2).float()
+    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv_transpose2d(w if isinstance(x, Tensor) else w.uop, stride=2))
+
+class TestTensorUOpEinsum(unittest.TestCase):
+  def test_einsum_dot(self):       _check(self, _t(2, 3), lambda x: type(x).einsum("ij,ij->", x, x))
+  def test_einsum_transpose(self): _check(self, _t(2, 3), lambda x: type(x).einsum("ij->ji", x))
+
+class TestTensorUOpSoftmax(unittest.TestCase):
+  def test_softmax_default(self):     _check(self, _t(2, 3).float(), lambda x: x.softmax())
+  def test_softmax_axis0(self):       _check(self, _t(2, 3).float(), lambda x: x.softmax(axis=0))
+  def test_log_softmax_default(self): _check(self, _t(2, 3).float(), lambda x: x.log_softmax())
+  def test_log_softmax_axis0(self):   _check(self, _t(2, 3).float(), lambda x: x.log_softmax(axis=0))
+
+class TestTensorUOpQR(unittest.TestCase):
+  def _check(self, t):
+    qt, rt = t.qr()
+    qu, ru = t.uop.qr()
+    self.assertIs(_strip_unique(qt.uop), _strip_unique(qu))
+    self.assertIs(_strip_unique(rt.uop), _strip_unique(ru))
+  def test_qr_square(self):   self._check(_t(3, 3).float())
+  def test_qr_tall(self):     self._check(_t(4, 3).float())
+  def test_qr_wide(self):     self._check(_t(3, 4).float())
+  def test_qr_zero_col(self): self._check(Tensor([[0.0, 1.0], [0.0, 2.0]]))
+  def test_qr_batched(self):  self._check(_t(2, 3, 3).float())
+
+class TestTensorUOpSVD(unittest.TestCase):
+  def _check(self, t, **kw):
+    ut, st, vt = t.svd(**kw)
+    uu, su, vu = t.uop.svd(**kw)
+    self.assertIs(_strip_unique(ut.uop), _strip_unique(uu))
+    self.assertIs(_strip_unique(st.uop), _strip_unique(su))
+    self.assertIs(_strip_unique(vt.uop), _strip_unique(vu))
+  def test_svd_square(self):    self._check(_t(2, 2).float())
+  def test_svd_tall(self):      self._check(_t(3, 2).float())
+  def test_svd_wide(self):      self._check(_t(2, 3).float())
+  def test_svd_odd_num(self):   self._check(_t(3, 3).float())  # exercises odd-num runoff path
+  def test_svd_batched(self):   self._check(_t(2, 2, 2).float())
+  def test_svd_nonfull(self):   self._check(_t(3, 2).float(), full_matrices=False)
+
+# UOp.empty / UOp.empty_like are the canonical buffer allocators; Tensor.empty / Tensor.empty_like just forward.
+class TestUOpEmpty(unittest.TestCase):
+  def test_empty_dtype_string(self):
+    self.assertEqual(UOp.empty((3, 4), dtype="float32").dtype, dtypes.float32)
+
+  def test_empty_like_dtype_override(self):
+    u = Tensor.ones(3, 4).uop.empty_like(dtype=dtypes.int8)
+    self.assertEqual((u.shape, u.dtype), ((3, 4), dtypes.int8))
+    self.assertTrue(u.has_buffer_identity())
+
+  def test_empty_like_sharded_to_single_device(self):
+    # regression: sharded source, override to single device must yield full logical shape with no axis
+    t = Tensor.ones(8, 4).shard(("NULL:0", "NULL:1"), axis=0)
+    for dev in ("NULL:2", ("NULL:2",)):  # singleton tuple also canonicalizes to single device
+      u = t.uop.empty_like(device=dev, dtype=dtypes.int32)
+      self.assertEqual((u.shape, u.device, u.dtype, u.axis), ((8, 4), "NULL:2", dtypes.int32, None))
+      self.assertTrue(u.has_buffer_identity())
+
+  def test_empty_direct_singleton_tuple_device(self):
+    # regression: direct UOp.empty with a singleton-tuple device + axis must not trip .multi()'s tuple assert
+    u = UOp.empty((4,), dtype=dtypes.float32, device=("NULL:0",), axis=0)
+    self.assertEqual((u.shape, u.device, u.axis), ((4,), "NULL", None))
+
+class TestTensorUOpCreation(unittest.TestCase):
+  def test_full(self):
+    self.assertIs(_strip_unique(Tensor.full((2, 3), 42).uop), _strip_unique(UOp.full((2, 3), 42)))
+  def test_full_kwargs(self):
+    self.assertIs(_strip_unique(Tensor.full((2, 3), 42, dtype=dtypes.int8, device="NULL").uop),
+                  _strip_unique(UOp.full((2, 3), 42, dtype=dtypes.int8, device="NULL")))
+  def test_full_symbolic_fill(self):
+    t = Tensor.full((2, 3), UOp.variable("x", 1, 10).bind(5))
+    self.assertEqual(t.shape, (2, 3))
+  def test_zeros(self):
+    self.assertIs(_strip_unique(Tensor.zeros(2, 3).uop), _strip_unique(UOp.zeros(2, 3)))
+  def test_ones(self):
+    self.assertIs(_strip_unique(Tensor.ones(2, 3).uop), _strip_unique(UOp.ones(2, 3)))
+  def test_invalids(self):
+    self.assertIs(_strip_unique(Tensor.invalids(2, 3, dtype=dtypes.int8).uop), _strip_unique(UOp.invalids(2, 3, dtype=dtypes.int8)))
+  def test_arange(self):
+    self.assertIs(_strip_unique(Tensor.arange(5).uop), _strip_unique(UOp.arange(5)))
+  def test_arange_empty(self):
+    self.assertIs(_strip_unique(Tensor.arange(5, 5).uop), _strip_unique(UOp.arange(5, 5)))
+  def test_arange_step(self):
+    self.assertIs(_strip_unique(Tensor.arange(5, 10, 2).uop), _strip_unique(UOp.arange(5, 10, 2)))
+  def test_linspace(self):
+    self.assertIs(_strip_unique(Tensor.linspace(0, 10, 5).uop), _strip_unique(UOp.linspace(0, 10, 5)))
+  def test_linspace_one_step(self):
+    self.assertIs(_strip_unique(Tensor.linspace(5, 10, 1).uop), _strip_unique(UOp.linspace(5, 10, 1)))
+  def test_eye(self):
+    self.assertIs(_strip_unique(Tensor.eye(3).uop), _strip_unique(UOp.eye(3)))
+  def test_eye_rect(self):
+    self.assertIs(_strip_unique(Tensor.eye(2, 4).uop), _strip_unique(UOp.eye(2, 4)))
+  def test_triu(self):
+    t = _t(3, 4)
+    self.assertIs(_strip_unique(t.triu().uop), _strip_unique(t.uop.triu()))
+  def test_triu_diagonal(self):
+    t = _t(3, 4)
+    self.assertIs(_strip_unique(t.triu(diagonal=1).uop), _strip_unique(t.uop.triu(diagonal=1)))
+  def test_tril(self):
+    t = _t(3, 4)
+    self.assertIs(_strip_unique(t.tril().uop), _strip_unique(t.uop.tril()))
+  def test_tril_diagonal(self):
+    t = _t(3, 4)
+    self.assertIs(_strip_unique(t.tril(diagonal=-1).uop), _strip_unique(t.uop.tril(diagonal=-1)))
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/null/test_tensor_uop_representation.py

@@ -0,0 +1,60 @@
+import unittest
+from tinygrad import Tensor
+from tinygrad.uop.ops import UPat, Ops, UOp
+
+# NOTE: unlike before base for a realized tensor is always a BUFFER
+realized_pattern = UPat(Ops.BUFFER)
+def is_pattern_uop(u:UOp, pat:UPat): assert pat.match(u, {}), f"{u}\nis not\n{pat}"
+def is_pattern(ten:Tensor, pat:UPat): is_pattern_uop(ten.uop, pat)
+
+class TestTensorMutates(unittest.TestCase):
+  @unittest.skip("this doesn't mutate anymore")
+  def test_mutate_add(self):
+    a = Tensor([1,2,3])
+    b = Tensor([4,5,6])
+    ret = a+b
+    pa = a.uop
+    pb = b.uop
+    pr = ret.uop
+    ret.schedule_linear()
+    self.assertIsNot(pa, a.uop)
+    self.assertIsNot(pb, b.uop)
+    self.assertIsNot(pr, ret.uop)
+    for t in [a,b,ret]: is_pattern_uop(t.uop.base, realized_pattern)
+
+  def test_reshape_is_same_parent(self):
+    a = Tensor([1,2,3])
+    b = Tensor([4,5,6])
+    c = a+b
+    d = (a+b).reshape(3,1)
+    d.realize()
+    is_pattern_uop(d.uop.base, realized_pattern)
+    is_pattern_uop(c.uop.base, realized_pattern)
+    is_pattern_uop(c.uop.base, realized_pattern)
+    assert d.uop is not d.uop.base
+
+  def test_reshape_is_same_child(self):
+    a = Tensor([1,2,3])
+    b = Tensor([4,5,6])
+    c = a+b
+    d = (a+b).reshape(3,1)
+    c.realize()
+    is_pattern_uop(c.uop.base, realized_pattern)
+    is_pattern_uop(d.uop.base, realized_pattern)
+
+class TestTensorUopRepresentation(unittest.TestCase):
+  def test_realized(self):
+    a = Tensor([1.,2,3]).realize()
+    print(a.uop)
+    is_pattern_uop(a.uop.base, realized_pattern)
+
+  def test_add_realized(self):
+    a = Tensor([1.,2,3]).realize()
+    b = Tensor([4.,5,6]).realize()
+    c = a+b
+    print(c.uop)
+    is_pattern(c, UPat(Ops.ADD))
+    for s in c.uop.src: is_pattern_uop(s.base, realized_pattern)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_tinyfs.py

@@ -0,0 +1,26 @@
+import unittest
+from tinygrad import Tensor
+
+class TestLoadStore(unittest.TestCase):
+  def test_load_shape(self):
+    t = Tensor(bytes(16)).fs_load(1024)
+    assert t.shape == (1024,), t.shape
+    t.schedule_linear()
+
+  def test_store_shape(self):
+    t = Tensor.zeros(1024).fs_store()
+    assert t.shape == (16,), t.shape
+    t.schedule_linear()
+
+  def test_load_large_shape(self):
+    t = Tensor(bytes(16)).fs_load(10_000_000)
+    assert t.shape == (10_000_000,), t.shape
+    t.schedule_linear()
+
+  def test_store_large_shape(self):
+    t = Tensor.zeros(10_000_000).fs_store()
+    assert t.shape == (16,), t.shape
+    t.schedule_linear()
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/null/test_tqdm.py

@@ -0,0 +1,332 @@
+import time, random, unittest, itertools
+from unittest.mock import patch
+from io import StringIO
+from collections import namedtuple
+from tqdm import tqdm
+from tinygrad.helpers import tqdm as tinytqdm, trange
+import numpy as np
+
+def _get_iter_per_second(raw:str) -> float:
+  # raw might have unit scale
+  if raw.endswith("k"): return float(raw[:-1])*1e3
+  if raw.endswith("M"): return float(raw[:-1])*1e6
+  return float(raw)
+
+# TODO: _get_iter_per_second in test_unit_scale might fail if lower bound is too small
+NCOLS_RANGE = [80, 240]
+
+class TestProgressBar(unittest.TestCase):
+  def _compare_bars(self, bar1, bar2):
+    prefix1, prog1, suffix1 = bar1.split("|")
+    prefix2, prog2, suffix2 = bar2.split("|")
+
+    self.assertEqual(len(bar1), len(bar2))
+    self.assertEqual(prefix1, prefix2)
+
+    def parse_timer(timer): return sum(int(x) * y for x, y in zip(timer.split(':')[::-1], (1, 60, 3600)))
+
+    if "?" not in suffix1 and "?" not in suffix2:
+      # allow for few sec diff in timers (removes flakiness)
+      timer1, rm1 = [parse_timer(timer) for timer in suffix1.split("[")[-1].split(",")[0].split("<")]
+      timer2, rm2 = [parse_timer(timer) for timer in suffix2.split("[")[-1].split(",")[0].split("<")]
+      np.testing.assert_allclose(timer1, timer2, atol=5, rtol=1e-2)
+      np.testing.assert_allclose(rm1, rm2, atol=5, rtol=1e-2)
+
+      # get suffix without timers
+      suffix1 = suffix1.split("[")[0] + suffix1.split(",")[1]
+      suffix2 = suffix2.split("[")[0] + suffix2.split(",")[1]
+      self.assertEqual(suffix1, suffix2)
+    else:
+      self.assertEqual(suffix1, suffix2)
+
+    diff = sum([c1 != c2 for c1, c2 in zip(prog1, prog2)])  # allow 1 char diff to be less flaky, but it should match
+    assert diff <= 1, f"{diff=}\n{prog1=}\n{prog2=}"
+
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_tqdm_output_iter(self, mock_terminal_size, mock_stderr):
+    for _ in range(10):
+      total, ncols = random.randint(5, 30), random.randint(*NCOLS_RANGE)
+      mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+      mock_stderr.truncate(0)
+
+      # compare bars at each iteration (only when tinytqdm bar has been updated)
+      for n in (bar := tinytqdm(range(total), desc="Test")):
+        if bar.i % bar.skip != 0: continue
+        tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+        iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
+        elapsed = n/iters_per_sec if n>0 else 0
+        tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
+        self._compare_bars(tinytqdm_output, tqdm_output)
+
+      # compare final bars
+      tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+      iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
+      elapsed = total/iters_per_sec if n>0 else 0
+      tqdm_output = tqdm.format_meter(n=total, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
+      self._compare_bars(tinytqdm_output, tqdm_output)
+
+  @unittest.skip("this is flaky")
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_unit_scale(self, mock_terminal_size, mock_stderr):
+    for unit_scale in [True, False]:
+      # NOTE: numpy comparison raises TypeError if exponent > 22
+      for exponent in range(1, 22, 3):
+        low, high = 10 ** exponent, 10 ** (exponent+1)
+        for _ in range(5):
+          total, ncols = random.randint(low, high), random.randint(*NCOLS_RANGE)
+          mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+          mock_stderr.truncate(0)
+
+          # compare bars at each iteration (only when tinytqdm bar has been updated)
+          # setting high rate to make sure it does not skip
+          for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=10**9):
+            tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+
+            if n:
+              iters_per_sec = _get_iter_per_second(tinytqdm_output.split("it/s")[-2].split(" ")[-1])
+              elapsed = n/iters_per_sec
+            else:
+              elapsed = 0
+            tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=unit_scale)
+            self._compare_bars(tinytqdm_output, tqdm_output)
+            if n > 3: break
+
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_unit_scale_exact(self, mock_terminal_size, mock_stderr):
+    unit_scale = True
+    ncols = 80
+    mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+    mock_stderr.truncate(0)
+
+    total = 10
+    with patch('time.perf_counter', side_effect=[0]+list(range(100))):  # one more 0 for the init call
+      # compare bars at each iteration (only when tinytqdm bar has been updated)
+      for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=10**9):
+        tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+        elapsed = n
+        tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=unit_scale)
+        self._compare_bars(tinytqdm_output, tqdm_output)
+        if n > 5: break
+
+    total = 10
+    k=0.001000001
+    # regression test for
+    # E   AssertionError: ' 1.00/10.0  1000it/s]' != ' 1.00/10.0  1.00kit/s]'
+    # E   -  1.00/10.0  1000it/s]
+    # E   ?                ^
+    # E   +  1.00/10.0  1.00kit/s]
+    # E   ?              +  ^
+    with patch('time.perf_counter', side_effect=[0, *[i*k for i in range(100)]]):  # one more 0 for the init call
+      # compare bars at each iteration (only when tinytqdm bar has been updated)
+      for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=10**9):
+        tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+        elapsed = n*k
+        tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=unit_scale)
+        self._compare_bars(tinytqdm_output, tqdm_output)
+        if n > 5: break
+
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_si_boundary(self, mock_terminal_size, mock_stderr):
+    """Test SI formatting at boundaries (e.g., 999.5 -> 1.00k, not 1000)"""
+    ncols = 80
+    mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+
+    # Test rates at the boundary: 999 stays as "999", 999.5+ becomes "1.00k"
+    for rate in [999, 999.4, 999.5, 1000, 1001]:
+      mock_stderr.truncate(0)
+      mock_stderr.seek(0)
+      elapsed = 1.0 / rate
+      # Need 3 perf_counter calls: init st, init update, final update
+      with patch('time.perf_counter', side_effect=[0, 0, elapsed]):
+        bar = tinytqdm(desc="Test", total=1, unit_scale=True, rate=10**9)
+        bar.update(1, close=True)
+      tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+      tqdm_output = tqdm.format_meter(n=1, total=1, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=True)
+      self._compare_bars(tinytqdm_output, tqdm_output)
+
+  @unittest.skip("this is flaky")
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_set_description(self, mock_terminal_size, mock_stderr):
+    for _ in range(10):
+      total, ncols = random.randint(5, 30), random.randint(*NCOLS_RANGE)
+      mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+      mock_stderr.truncate(0)
+
+      expected_prefix = "Test"
+      # compare bars at each iteration (only when tinytqdm bar has been updated)
+      for i,n in enumerate(bar := tinytqdm(range(total), desc="Test")):
+        if bar.i % bar.skip != 0: continue
+        tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+        iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
+        elapsed = n/iters_per_sec if n>0 else 0
+        tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix=expected_prefix)
+        expected_prefix = desc = f"Test {i}" if i % 2 == 0 else ""
+        bar.set_description(desc)
+        self._compare_bars(tinytqdm_output, tqdm_output)
+
+      # compare final bars
+      tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+      iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
+      elapsed = total/iters_per_sec if n>0 else 0
+      tqdm_output = tqdm.format_meter(n=total, total=total, elapsed=elapsed, ncols=ncols, prefix=expected_prefix)
+      self._compare_bars(tinytqdm_output, tqdm_output)
+
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_trange_output_iter(self, mock_terminal_size, mock_stderr):
+    for _ in range(5):
+      total, ncols = random.randint(5, 30), random.randint(*NCOLS_RANGE)
+      mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+      mock_stderr.truncate(0)
+
+      # compare bars at each iteration (only when tinytqdm bar has been updated)
+      for n in (bar := trange(total, desc="Test")):
+        if bar.i % bar.skip != 0: continue
+        tiny_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+        iters_per_sec = float(tiny_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
+        elapsed = n/iters_per_sec if n>0 else 0
+        tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
+        self._compare_bars(tiny_output, tqdm_output)
+
+      # compare final bars
+      tiny_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+      iters_per_sec = float(tiny_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
+      elapsed = total/iters_per_sec if n>0 else 0
+      tqdm_output = tqdm.format_meter(n=total, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
+      self._compare_bars(tiny_output, tqdm_output)
+
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_tqdm_output_custom(self, mock_terminal_size, mock_stderr):
+    for _ in range(10):
+      total, ncols = random.randint(10000, 1000000), random.randint(*NCOLS_RANGE)
+      mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+      mock_stderr.truncate(0)
+
+      # compare bars at each iteration (only when tinytqdm bar has been updated)
+      bar = tinytqdm(total=total, desc="Test")
+      n = 0
+      while n < total:
+        incr = (total // 100) + random.randint(0, 1000)
+        if n + incr > total: incr = total - n
+        bar.update(incr, close=n+incr==total)
+        n += incr
+        if bar.i % bar.skip != 0: continue
+
+        tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+        iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
+        elapsed = n/iters_per_sec if n>0 else 0
+        tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
+        self._compare_bars(tinytqdm_output, tqdm_output)
+
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_tqdm_output_custom_0_total(self, mock_terminal_size, mock_stderr):
+    for _ in range(10):
+      total, ncols = random.randint(10000, 100000), random.randint(*NCOLS_RANGE)
+      mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+      mock_stderr.truncate(0)
+
+      # compare bars at each iteration (only when tinytqdm bar has been updated)
+      bar = tinytqdm(total=0, desc="Test")
+      n = 0
+      while n < total:
+        incr = (total // 10) + random.randint(0, 100)
+        if n + incr > total: incr = total - n
+        bar.update(incr, close=n+incr==total)
+        n += incr
+        if bar.i % bar.skip != 0: continue
+
+        tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+        iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
+        elapsed = n/iters_per_sec if n>0 else 0
+        tqdm_output = tqdm.format_meter(n=n, total=0, elapsed=elapsed, ncols=ncols, prefix="Test")
+        self.assertEqual(tinytqdm_output, tqdm_output)
+
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_tqdm_output_custom_nolen_total(self, mock_terminal_size, mock_stderr):
+    for unit_scale in [True, False]:
+      for _ in range(5):
+        gen = itertools.count(0)
+        ncols = random.randint(*NCOLS_RANGE)
+        mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+        mock_stderr.truncate(0)
+
+        # compare bars at each iteration (only when tinytqdm bar has been updated)
+        # setting high rate to make sure it does not skip
+        for n,g in enumerate(tinytqdm(gen, desc="Test", unit_scale=unit_scale, rate=10**9)):
+          assert g == n
+          tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+          if n:
+            iters_per_sec = _get_iter_per_second(tinytqdm_output.split("it/s")[-2].split(" ")[-1])
+            elapsed = n/iters_per_sec
+          else:
+            elapsed = 0
+          tqdm_output = tqdm.format_meter(n=n, total=0, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=unit_scale)
+          self.assertEqual(tinytqdm_output, tqdm_output)
+          if n > 5: break
+
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_tqdm_write(self, mock_terminal_size, mock_stderr):
+    for _ in range(5):
+      ncols, tqdm_fp = random.randint(*NCOLS_RANGE), StringIO()
+      mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+      mock_stderr.truncate(0)
+      tqdm_fp.truncate(0)
+      for i in tinytqdm(range(10)):
+        time.sleep(0.01)
+        tinytqdm.write(str(i))
+        tqdm.write(str(i), file=tqdm_fp)
+        tinytqdm_out, tqdm_out = mock_stderr.getvalue(), tqdm_fp.getvalue()
+        self.assertEqual(tinytqdm_out.split("\r\033[K")[-1], tqdm_out.split(f"{i-1}\n")[-1])
+      self.assertEqual(tinytqdm_out, tinytqdm_out)
+
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_tqdm_context_manager(self, mock_terminal_size, mock_stderr):
+    for _ in range(10):
+      total, ncols = random.randint(5, 30), random.randint(*NCOLS_RANGE)
+      mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+      mock_stderr.truncate(0)
+
+      with tinytqdm(desc="Test", total=total) as bar:
+        for _ in range(total):
+          bar.update(1)
+
+      tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+      iters_per_sec = float(tinytqdm_output.split("it/s")[-2].split(" ")[-1])
+      elapsed = total/iters_per_sec
+      tqdm_output = tqdm.format_meter(n=total, total=total, elapsed=elapsed, ncols=ncols, prefix="Test")
+      self._compare_bars(tinytqdm_output, tqdm_output)
+
+  def test_tqdm_perf(self):
+    st = time.perf_counter()
+    for _ in tqdm(range(100)): pass
+    tqdm_time = time.perf_counter() - st
+
+    st = time.perf_counter()
+    for _ in tinytqdm(range(100)): pass
+    tinytqdm_time = time.perf_counter() - st
+
+    assert tinytqdm_time < 5 * tqdm_time
+
+  def test_tqdm_perf_high_iter(self):
+    st = time.perf_counter()
+    for _ in tqdm(range(10^7)): pass
+    tqdm_time = time.perf_counter() - st
+
+    st = time.perf_counter()
+    for _ in tinytqdm(range(10^7)): pass
+    tinytqdm_time = time.perf_counter() - st
+
+    assert tinytqdm_time < 20 * tqdm_time
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_transcendental_helpers.py

@@ -0,0 +1,109 @@
+import unittest, math
+import numpy as np
+from tinygrad import dtypes
+from tinygrad.uop.ops import UOp
+from tinygrad.uop.decompositions import TRANSCENDENTAL_DTYPES, payne_hanek_reduction, cody_waite_reduction
+from tinygrad.uop.decompositions import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if
+from test.helpers import eval_uop
+
+class TestTranscendentalFunctions(unittest.TestCase):
+  def test_payne_hanek_reduction(self):
+    # TODO: Test constant input when constant folding is fixed (or maybe test both variants)
+    # Load input value from a buffer to prevent constant folding
+    input_buf = UOp.param(1, dtypes.double.ptr())
+    loaded_value = input_buf.index(UOp.const(dtypes.int, 0))
+    def eval_payne_hanek_reduction(v:float) -> tuple[float, int]:
+      return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value))
+
+    r, q = eval_payne_hanek_reduction(12 * math.pi + 0.1)
+    np.testing.assert_allclose(r, 0.1 - math.pi / 2)
+    np.testing.assert_equal(q, 1)
+
+    r, q = eval_payne_hanek_reduction(12 * math.pi)
+    np.testing.assert_allclose(r, 0.0, atol=1e-8)
+    np.testing.assert_equal(q, 4)
+
+    r, q = eval_payne_hanek_reduction(12 * math.pi - 0.1)
+    np.testing.assert_allclose(r, -0.1)
+    np.testing.assert_equal(q, 4)
+
+  def test_cody_waite_reduction(self):
+    r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(dtypes.float64, 12 * math.pi + 0.1)))
+    np.testing.assert_allclose(r, 0.1)
+    np.testing.assert_equal(q, 12)
+
+  def test_frexp(self):
+    for x in (1, -1):
+      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, x)))
+      np.testing.assert_equal(mantissa, 0.5)
+      np.testing.assert_equal(exponent, 1)
+
+    for x in (2, -2):
+      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 2.0)))
+      np.testing.assert_equal(mantissa, 0.5)
+      np.testing.assert_equal(exponent, 2)
+
+    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 5.0)))
+    np.testing.assert_equal(mantissa, 0.625)
+    np.testing.assert_equal(exponent, 3)
+
+    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 1000.0)))
+    np.testing.assert_allclose(mantissa, 0.9765625)
+    np.testing.assert_equal(exponent, 10)
+
+  def test_rintk(self):
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 0.0))), 0)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.0))), 5)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.5))), 6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.999))), 6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.0))), -5)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.5))), -6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.999))), -6)
+
+  def test_pow2if(self):
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 0), dtypes.float)), 1.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 1), dtypes.float)), 2.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 2), dtypes.float)), 4.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 10), dtypes.float)), 1024.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 63), dtypes.float)), 2**63)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -1), dtypes.float)), 0.5)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -2), dtypes.float)), 0.25)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -10), dtypes.float)), 2**-10)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -63), dtypes.float)), 2**-63)
+
+class TestTranscendentalVectorizedFunctions(unittest.TestCase):
+  # given a scalar and vectorized input, check that the fxn outputs have the same
+  # scalar_dtypes, args, ops, and vcount (only for vectorized input)
+
+  def _check_uop_vcount(self, u:tuple|UOp, vcount:int):
+    # check all UOps in u are vectorized with vcount
+    if isinstance(u, UOp):
+      assert u.dtype.vcount == vcount, f'expected {vcount=} but got {u.dtype.vcount=} for UOp\n{u=}'
+    [self._check_uop_vcount(x, vcount) for x in (u if isinstance(u, tuple) else u.src)]
+
+  def _check_uops_match(self, u1:tuple|UOp, u2:tuple|UOp):
+    # check all UOps in u1, u2 have the same scalar_dtype, args, ops
+    if isinstance(u1, UOp) and isinstance(u2, UOp):
+      assert u1.dtype.scalar() == u2.dtype.scalar(), f'expected {u1.dtype.scalar()=} but got {u2.dtype.scalar()=} for UOps\n{u1=}\n{u2}'
+      assert u1.arg == u2.arg or (math.isnan(u1.arg) and math.isnan(u2.arg)), f'expected {u1.arg=} but got {u2.arg=} for UOps\n{u1=}\n{u2}'
+      assert u1.op == u2.op, f'expected {u1.op=} but got {u2.op=} for UOps\n{u1=}\n{u2}'
+    [self._check_uops_match(x1, x2) for x1, x2 in zip((u1 if isinstance(u1, tuple) else u1.src), (u2 if isinstance(u2, tuple) else u2.src))]
+
+  def _test_vectorized(self, fxn, scalar_dtypes=TRANSCENDENTAL_DTYPES, vals=[-2,1.3,194], vcounts=[1,4,19]):
+    for scalar_dtype in scalar_dtypes:
+      for val in vals:
+        for vcount in vcounts:
+          in_scalar, in_vec = UOp.const(scalar_dtype, val), UOp.const(scalar_dtype.vec(vcount), val)
+          out_scalar, out_vec = fxn(in_scalar), fxn(in_vec)
+          self._check_uops_match(out_scalar, out_vec)
+          self._check_uop_vcount(out_vec, vcount)
+
+  def test_xpow(self): return self._test_vectorized(lambda x: xpow(x, x))
+  def test_xexp2(self): return self._test_vectorized(xexp2)
+  def test_xlog2(self): return self._test_vectorized(xlog2)
+  def test_payne_hanek_reduction(self): return self._test_vectorized(payne_hanek_reduction)
+  def test_cody_waite_reduction(self): return self._test_vectorized(cody_waite_reduction)
+  def test_trig_poly(self): return self._test_vectorized(lambda x: trig_poly(x, [0.0], [1.0]))
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_uop_graph.py

@@ -0,0 +1,979 @@
+import unittest, pytest
+from tinygrad import dtypes, Variable
+from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import DEBUG, Context
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp, AxisType
+from tinygrad.uop.symbolic import sym
+from tinygrad.codegen.late.expander import expander
+from test.helpers import to_uops_list
+
+simple_pm = PatternMatcher([
+  (UPat.cvar('x', dtypes.int), lambda x: UOp.const(dtypes.float, 1.0) + UOp.const(dtypes.float, 2.0)),
+  (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(dtypes.float, x.arg+y.arg)),
+  (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(dtypes.float, x.arg*y.arg*z.arg)),
+  ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.arg+c2.arg)),
+])
+
+def const_values(u:UOp):
+  if u.op is Ops.CONST: return (u.arg,)*u.dtype.count
+  if u.op is Ops.STACK: return tuple(x.arg for x in u.src)
+  raise AssertionError(f"expected const-like UOp, got {u.op}")
+
+class TestGraphRewriteConst(unittest.TestCase):
+  def test_gep_const(self):
+    v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
+    v2 = v1.gep(1)
+    ret = graph_rewrite(v2, sym)
+    self.assertEqual(ret.dtype, dtypes.int)
+    self.assertEqual(ret.arg, 1)
+
+  def test_gep_const_single(self):
+    v1 = UOp.const(dtypes.int.vec(3), 4)
+    v2 = v1.gep(1)
+    ret = graph_rewrite(v2, sym)
+    self.assertEqual(ret.dtype, dtypes.int)
+    self.assertEqual(ret.arg, 4)
+
+  def test_add_const(self):
+    v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
+    v2 = UOp.const(dtypes.int.vec(3), (5,6,7))
+    ret = graph_rewrite(v1+v2, sym)
+    self.assertEqual(ret.op, Ops.STACK)
+    self.assertEqual(ret.dtype, dtypes.int.vec(3))
+    self.assertEqual(const_values(ret), (5,7,9))
+
+  def test_add_const_lose_v(self):
+    v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
+    v2 = UOp.const(dtypes.int.vec(3), (2,1,0))
+    ret = graph_rewrite(v1+v2, sym)
+    self.assertEqual(ret.op, Ops.CONST)
+    self.assertEqual(ret.dtype, dtypes.int.vec(3))
+    self.assertEqual(ret.arg, 2)
+
+def xfail_broken_const_wraparound(fn):
+  fn = pytest.mark.xfail(reason="const folding does not properly implement modular arithmetic")(fn)
+  return unittest.expectedFailure(fn)
+class TestModularWraparound(unittest.TestCase):
+  def _test(self, uop:UOp, expected:int):
+    results = to_uops_list([uop])
+    self.assertEqual(len(results), 2)  # +1 for SINK
+    self.assertEqual(results[0].op, Ops.CONST)
+    self.assertEqual(results[0].dtype, uop.dtype)
+    self.assertEqual(results[0].arg, expected)
+
+  @xfail_broken_const_wraparound
+  def test_cast(self):
+    t = self._test
+    t(UOp.const(dtypes.uint, 0xABCD17D6).cast(dtypes.uint8), 0xD6)
+    t(UOp.const(dtypes.uint, 0xABCD17D6).cast(dtypes.uint8).cast(dtypes.uint), 0xD6)
+
+  @xfail_broken_const_wraparound
+  def test_mul(self):
+    t = self._test
+    t(UOp.const(dtypes.uint, 0xABCD17D6) * 0xAABBCCDD, 1147018174)
+    t(UOp.const(dtypes.int, 0xABCD17D6) * 10, -1241321892)
+
+  @xfail_broken_const_wraparound
+  def test_div(self):
+    t = self._test
+    t(UOp.const(dtypes.uint, 0xABCD17D6) * 0xAABBCCDD // 11, 104274379)
+    t(UOp.const(dtypes.int, 0xABCD17D6) * 10 // 11, -112847444)
+
+  @xfail_broken_const_wraparound
+  def test_neg(self):
+    t = self._test
+    t(-UOp.const(dtypes.uint8, 1), 0xFF)
+    t(-UOp.const(dtypes.uint16, 1), 0xFFFF)
+    t(-UOp.const(dtypes.uint32, 1), 0xFFFFFFFF)
+    t(-UOp.const(dtypes.uint64, 1), 0xFFFFFFFFFFFFFFFF)
+
+  @xfail_broken_const_wraparound
+  def test_neg_min_int(self):
+    t = self._test
+    t(-UOp.const(dtypes.int8, -2**7), -2**7)
+    t(-UOp.const(dtypes.int16, -2**15), -2**15)
+    t(-UOp.const(dtypes.int32, -2**31), -2**31)
+    t(-UOp.const(dtypes.int64, -2**63), -2**63)
+
+  @xfail_broken_const_wraparound
+  def test_payne_hanek_reduction_bug(self):
+    t = self._test
+    a = (UOp.const(dtypes.uint, 43748177600).cast(dtypes.uint) | 36).cast(dtypes.ulong)
+    b = 2536655455 * a + 4294967296 * UOp.const(dtypes.ulong, 25366554550)
+    c = (b + 2261737165) // 4611686018427387904
+    t(c, 0)
+
+class TestGraphRewrite(unittest.TestCase):
+  def test_dedup(self):
+    v1 = UOp(Ops.DEFINE_VAR, dtypes.float)
+    v2 = UOp(Ops.DEFINE_VAR, dtypes.float)
+    nout = graph_rewrite(v1+v2, PatternMatcher([]))
+    self.assertIs(nout.src[0], nout.src[1])
+
+  # NOTE: this shows why we can't have a UOp in arg
+  @unittest.expectedFailure
+  def test_no_dedup_args(self):
+    a1 = UOp(Ops.DEFINE_VAR, dtypes.int, (), ("a1", UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 11)))
+    a2 = UOp(Ops.DEFINE_VAR, dtypes.int, (), ("a2", UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 11)))
+    sink = a1.sink(a2)
+    define_vars = [x for x in graph_rewrite(sink, PatternMatcher([])).toposort() if x.op is Ops.DEFINE_VAR]
+    self.assertEqual(len(define_vars), 1)
+
+  def test_simple(self):
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    nout = graph_rewrite(c1+c2, simple_pm)
+    self.assertEqual(nout.op, Ops.CONST)
+    self.assertEqual(nout.arg, 3.0)
+
+  def test_depth_2_late(self):
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
+    nout = graph_rewrite(c1*c2*(c3+c3), simple_pm)
+    self.assertEqual(nout.op, Ops.CONST)
+    self.assertEqual(nout.arg, 12.0)
+
+  def test_double(self):
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
+    nout = graph_rewrite(c1+c2+c3, simple_pm)
+    self.assertEqual(nout.op, Ops.CONST)
+    self.assertEqual(nout.arg, 6.0)
+
+  def test_triple(self):
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
+    c4 = UOp.const(dtypes.float, 4.0)
+    nout = graph_rewrite(c1+c2+c3+c4, simple_pm)
+    self.assertEqual(nout.op, Ops.CONST)
+    self.assertEqual(nout.arg, 10.0)
+
+  def test_diamond(self):
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
+    nout = graph_rewrite((c1+c2)+(c1+c3), simple_pm)
+    self.assertEqual(nout.op, Ops.CONST)
+    self.assertEqual(nout.arg, 7.0)
+
+  def test_magic_4(self):
+    c1 = UOp.const(dtypes.int, 4.0)
+    nout = graph_rewrite(c1, simple_pm)
+    self.assertEqual(nout.op, Ops.CONST)
+    self.assertEqual(nout.arg, 3.0)
+
+  def test_depth_2_fold(self):
+    v = UOp(Ops.DEFINE_VAR, dtypes.float)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    nout = graph_rewrite(v+c1+c2, simple_pm)
+    self.assertEqual(nout.op, Ops.ADD)
+    self.assertEqual(nout.src[0].op, Ops.DEFINE_VAR)
+    self.assertEqual(nout.src[1].op, Ops.CONST)
+    self.assertEqual(nout.src[1].arg, 3.0)
+
+  def test_commutative_work(self):
+    a = UOp.variable('a', 0, 1)
+    b = UOp.variable('b', 0, 1)
+    self.assertIs((a+b).simplify(), (b+a).simplify())
+
+  def test_consts_go_last_right_away(self):
+    a = UOp.variable('a', 0, 1)
+    tst = (2+a).simplify()
+    self.assertIs(tst.src[0], a)
+    self.assertIs(tst.src[1], a.const_like(2))
+
+  def test_consts_go_last(self):
+    a = UOp.variable('a', 0, 1)
+    b = UOp.variable('b', 0, 1)
+    c = UOp.variable('c', 0, 1)
+    d = UOp.variable('d', 0, 1)
+    outs = [2+a, 2+a+d+3+b+c+4, UOp(Ops.ADD, a.dtype, src=(a.const_like(2), a)), (4+d)+c+(2+a)+b]
+    for out in outs:
+      sink = graph_rewrite(out, sym)
+      print(sink.render())
+      self.assertEqual(sink.op, Ops.ADD)
+      self.assertEqual(sink.src[1].op, Ops.CONST)
+      self.assertEqual(len([x for x in sink.toposort() if x.op is Ops.CONST]), 1)
+
+class TestUOpGraph(unittest.TestCase):
+  def test_add_constant_fold(self):
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    out = UOp(Ops.ADD, dtypes.float, (c1, c2))
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.arg, 3.0)
+
+  def test_where_same_fold(self):
+    v = UOp.variable('tmp', 0, 1)
+    c0 = UOp.const(dtypes.weakint, 0)
+    vc = UOp(Ops.CMPNE, dtypes.bool, (v, c0))
+    c1 = UOp.const(dtypes.float, 1.0)
+    out = UOp(Ops.WHERE, dtypes.float, (vc, c1, c1))
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.arg, 1.0)
+
+  def test_where_const_fold(self):
+    bf = UOp.const(dtypes.bool, False)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    out = UOp(Ops.WHERE, dtypes.float, (bf, c1, c2))
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.arg, 2.0)
+
+  def test_const_cast(self):
+    bf = UOp.const(dtypes.bool, False)
+    out = UOp(Ops.CAST, dtypes.int, (bf,))
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.arg, 0)
+
+  def test_const_bitcast(self):
+    bf = UOp.const(dtypes.float, 1.0)
+    out = UOp(Ops.BITCAST, dtypes.uint32, (bf,))
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.arg, 0x3F800000)
+
+  @unittest.expectedFailure
+  def test_const_shape_change_bitcast(self):
+    bf = UOp.const(dtypes.uint8, 0x3F)
+    out = UOp(Ops.BITCAST, dtypes.half, (bf,))
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+
+  @unittest.skip("this test isn't valid uops")
+  def test_noop_vectorize_fold(self):
+    d0 = UOp.param(0, dtypes.float.ptr())
+    idx = UOp.const(dtypes.int, 0)
+    ld = UOp(Ops.LOAD, dtypes.float.vec(2), (d0, idx))
+    vec = UOp(Ops.STACK, dtypes.float.vec(2), (ld,))
+    x = UOp(Ops.GEP, dtypes.float, (vec, ), arg=0)
+    alu = UOp(Ops.SQRT, dtypes.float, (x, ))
+    out = UOp(Ops.STORE, dtypes.void, (d0, idx, alu))
+    uops = to_uops_list([out])
+    self.assertEqual(len([x for x in uops if x.op is Ops.STACK]), 0)
+
+  @unittest.skip("this test isn't valid uops")
+  def test_gep_vec_fold(self):
+    d0 = UOp.param(0, dtypes.float.ptr())
+    d1 = UOp.param(1, dtypes.float.ptr())
+    d2 = UOp.param(2, dtypes.float.ptr())
+    idx = UOp.const(dtypes.int, 0)
+    def _test_vec(geps, count=4):
+      vec = UOp(Ops.STACK, dtypes.float.vec(count), geps)
+      out = UOp(Ops.STORE, dtypes.void, (d0.index(idx), vec))
+      uops = to_uops_list([out])
+      if DEBUG >= 4:
+        from tinygrad import Device
+        print(Device[Device.DEFAULT].renderer.render(uops))
+      return uops[-2].src[-1]  # -2 to skip SINK
+
+    # possible
+    val = UOp(Ops.LOAD, dtypes.float.vec(4), (d1.index(idx),))
+    xyzw = tuple(UOp(Ops.GEP, dtypes.float, (val,), (i,)) for i in range(4))
+    self.assertIs(_test_vec(xyzw).op, Ops.LOAD)
+
+    # unaligned
+    val = UOp(Ops.LOAD, dtypes.float.vec(4), (d1.index(idx),))
+    wzyx = tuple(UOp(Ops.GEP, dtypes.float, (val,), (i,)) for i in reversed(range(4)))
+    self.assertIs(_test_vec(wzyx).op, Ops.STACK)
+
+    # different_size
+    val = UOp(Ops.LOAD, dtypes.float.vec(2), (d1.index(idx),))
+    xy = tuple(UOp(Ops.GEP, dtypes.float, (val, ), (i,)) for i in range(2))
+    self.assertIs(_test_vec(xy+xy).op, Ops.STACK)
+    val = UOp(Ops.LOAD, dtypes.float.vec(4), (d1.index(idx),))
+    xy = tuple(UOp(Ops.GEP, dtypes.float, (val, ), (i,)) for i in range(2))
+    self.assertIs(_test_vec(xy, count=2).op, Ops.STACK)
+
+    # different vals
+    val1 = UOp(Ops.LOAD, dtypes.float.vec(2), (d1.index(idx),))
+    val2 = UOp(Ops.LOAD, dtypes.float.vec(2), (d2.index(idx),))
+    xy1 = tuple(UOp(Ops.GEP, dtypes.float, (val1, ), (i,)) for i in range(2))
+    xy2 = tuple(UOp(Ops.GEP, dtypes.float, (val2, ), (i,)) for i in range(2))
+    self.assertIs(_test_vec(xy1+xy2).op, Ops.STACK)
+
+  def test_gep_vec_const_fold(self):
+    for vec_size in [2, 4, 8]:
+      consts = [UOp.const(dtypes.float, float(i)) for i in range(vec_size)]
+      vec = UOp(Ops.STACK, dtypes.float.vec(vec_size), tuple(consts))
+      with Context(SPEC=0):
+        uops = to_uops_list([UOp(Ops.GEP, dtypes.float, (vec,), (i,)) for i in range(vec_size)])
+        for uop, const in zip(uops, consts):
+          self.assertEqual(uop, const)
+
+  @unittest.skip("no longer testable standalone")
+  def test_wmma_vectorize_fold(self):
+    for i in [2, 4, 8]:
+      vec = UOp(Ops.STACK, dtypes.half.vec(i), tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
+      var = UOp(Ops.DEFINE_VAR, dtypes.half.vec(i))
+      acc = UOp.variable('acc', 0, 1, dtypes.half.vec(i))
+      wmma = UOp(Ops.WMMA, dtypes.half.vec(i), (vec, var, acc))
+      uops = to_uops_list([wmma])
+      self.assertEqual(uops[0], acc)
+      self.assertEqual(len(uops), 2)  # +1 for SINK
+
+    for i in [2, 4, 8]:
+      var = UOp(Ops.DEFINE_VAR, dtypes.half.vec(i))
+      vec = UOp(Ops.STACK, dtypes.half.vec(i), tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
+      acc = UOp.variable('acc', 0, 1, dtypes.half.vec(i))
+      wmma = UOp(Ops.WMMA, dtypes.half.vec(i), (var, vec, acc))
+      uops = to_uops_list([wmma])
+      self.assertEqual(uops[0], acc)
+      self.assertEqual(len(uops), 2)  # +1 for SINK
+
+  @unittest.skip("wmma is wrong here, it needs an arg")
+  def test_wmma_vectorize_no_fold(self):
+    for i in [4, 8]:
+      vec = UOp(Ops.STACK, dtypes.half.vec(i),
+                tuple(UOp.const(dtypes.half, 0.0) for _ in range(i//2)) +
+                tuple(UOp(Ops.DEFINE_VAR, dtypes.half, arg=(f'tmp{j}', UOp.const(dtypes.half, 0), UOp.const(dtypes.half, 1))) for j in range(i//2)))
+      var = UOp(Ops.DEFINE_VAR, dtypes.half.vec(i), arg=(f'tmp{i}', UOp.const(dtypes.half, 0), UOp.const(dtypes.half, 1)))
+      acc = UOp(Ops.DEFINE_VAR, dtypes.half.vec(i), arg=('acc', UOp.const(dtypes.half, 0), UOp.const(dtypes.half, 1)))
+      wmma = UOp(Ops.WMMA, dtypes.half.vec(i), (vec, var, acc))
+      uops = to_uops_list([wmma])
+      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
+
+    for i in [4, 8]:
+      var = UOp(Ops.DEFINE_VAR, dtypes.half.vec(i), arg=(f'tmp{i}', UOp.const(dtypes.half, 0), UOp.const(dtypes.half, 1)))
+      vec = UOp(Ops.STACK, dtypes.half.vec(i),
+                tuple(UOp.const(dtypes.half, 0.0) for _ in range(i//2)) +
+                tuple(UOp(Ops.DEFINE_VAR, dtypes.half, arg=(f'tmp{j}', UOp.const(dtypes.half, 0), UOp.const(dtypes.half, 1))) for j in range(i//2)))
+      acc = UOp(Ops.DEFINE_VAR, dtypes.half.vec(i), arg=('acc', UOp.const(dtypes.half, 0), UOp.const(dtypes.half, 1)))
+      wmma = UOp(Ops.WMMA, dtypes.half.vec(i), (var, vec, acc))
+      uops = to_uops_list([wmma])
+      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
+
+    for i in [2, 4, 8]:
+      vec = UOp(Ops.STACK, dtypes.half.vec(i),
+                tuple(UOp.const(dtypes.half, 1.0 if j == 0 else 0.0) for j in range(i)))
+      var = UOp(Ops.DEFINE_VAR, dtypes.half.vec(i), arg=(f'tmp{i}', UOp.const(dtypes.half, 0), UOp.const(dtypes.half, 1)))
+      acc = UOp(Ops.DEFINE_VAR, dtypes.half.vec(i), arg=('acc', UOp.const(dtypes.half, 0), UOp.const(dtypes.half, 1)))
+      wmma = UOp(Ops.WMMA, dtypes.half.vec(i), (vec, var, acc))
+      uops = to_uops_list([wmma])
+      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
+
+    for i in [2, 4, 8]:
+      var = UOp(Ops.DEFINE_VAR, dtypes.half.vec(i), arg=(f'tmp{i}', UOp.const(dtypes.half, 0), UOp.const(dtypes.half, 1)))
+      vec = UOp(Ops.STACK, dtypes.half.vec(i),
+                tuple(UOp.const(dtypes.half, 1.0 if j == 0 else 0.0) for j in range(i)))
+      acc = UOp(Ops.DEFINE_VAR, dtypes.half.vec(i), arg=('acc', UOp.const(dtypes.half, 0), UOp.const(dtypes.half, 1)))
+      wmma = UOp(Ops.WMMA, dtypes.half.vec(i), (var, vec, acc))
+      uops = to_uops_list([wmma])
+      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
+
+  def test_cast_alu_fold(self):
+    d0 = UOp.param(0, dtypes.bool.ptr())
+    d1 = UOp.param(1, dtypes.int.ptr())
+    idx = UOp.const(dtypes.int, 0)
+    ld = d1.index(idx)
+    alu = (ld<1).cast(dtypes.bool)
+    out = UOp(Ops.STORE, dtypes.void, (d0.index(idx), alu))
+    uops = to_uops_list([out])
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
+
+  def test_double_cast_fold(self):
+    d0 = UOp.param(0, dtypes.float.ptr())
+    d1 = UOp.param(1, dtypes.int.ptr())
+    idx = UOp.const(dtypes.int, 0)
+    ld = d1.index(idx)
+    alu = ld.cast(dtypes.float).cast(dtypes.float)
+    out = UOp(Ops.STORE, dtypes.void, (d0.index(idx), alu))
+    uops = to_uops_list([out])
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
+
+  def test_depth_2_const_fold(self):
+    v = UOp.variable("tmp", 0, 1, dtypes.int)
+    c2 = UOp.const(dtypes.int, 2)
+    c4 = UOp.const(dtypes.int, 4)
+    vc = UOp(Ops.ADD, dtypes.int, (v, c2))
+    out = UOp(Ops.ADD, dtypes.int, (vc, c4))
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 4)  # +1 for SINK
+    out = uops[-2]  # -2 to skip SINK
+    self.assertEqual(out.op, Ops.ADD)
+    self.assertEqual(out.src[1].op, Ops.CONST)
+    self.assertEqual(out.src[1].arg, 6)
+
+  def test_bitcast_to_same_dtype_fold(self):
+    for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
+      d0 = UOp.param(0, dt.ptr())
+      v = d0.index(UOp.const(dtypes.int, 0))
+      uops = to_uops_list([v.bitcast(dt)])
+      self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}")
+
+  def test_sub_with_cast_folds(self):
+    a = Variable("a", 0, 5)
+    uops = to_uops_list([a.cast(dtypes.int)+(-a).cast(dtypes.int)])
+    assert uops[0] == UOp.const(dtypes.int, 0)
+    assert uops[-1].op == Ops.SINK
+
+  def test_where_on_gated_load_fold(self):
+    ridx0 = UOp.range(100, 0)
+    d0 = UOp.param(0, dtypes.long.ptr())
+    ld = d0.index(ridx0.valid(ridx0<50))
+    w = (ridx0<50).where(ld, 5)
+    out = UOp.param(1, dtypes.long.ptr())
+    uops = to_uops_list([out.index(ridx0).store(w)])
+    for u in uops:
+      assert u.op is not Ops.WHERE
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg==5
+
+  def test_where_on_gated_load_folds_swapped_branches(self):
+    ridx0 = UOp.range(100, 0)
+    d0 = UOp.param(0, dtypes.long.ptr())
+    ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
+    w = (ridx0<50).where(5, ld)
+    uops = to_uops_list([w])
+    for u in uops:
+      assert u.op is not Ops.WHERE
+      if u.op is Ops.LOAD: assert u.src[1].arg==5
+
+  def test_where_on_gated_load_with_cast(self):
+    ridx0 = UOp.range(100, 0)
+    d0 = UOp.param(0, dtypes.int.ptr())
+    gate_idx = ridx0.valid((ridx0<50))
+    ld = d0.index(gate_idx).cast(dtypes.float)
+    w = (ridx0<50).where(ld, 5.0)
+    out = UOp.param(1, dtypes.float.ptr())
+    uops = to_uops_list([out.index(ridx0).store(w)])
+    for u in uops:
+      assert u.op is not Ops.WHERE
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg == 5
+
+  def test_where_on_casted_gated_load_extra_cond(self):
+    ridx0 = UOp.range(100, 0)
+    d0 = UOp.param(0, dtypes.float.ptr())
+    ld = d0.index(ridx0.valid(ridx0<50))
+    w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(dtypes.float, 0)).cast(dtypes.half)
+    out = UOp.param(1, dtypes.half.ptr())
+    uops = to_uops_list([out.index(ridx0).store(w)])
+    for u in uops:
+      assert u.op is not Ops.WHERE
+
+  def test_where_on_casted_gated_load_extra_cond_swapped(self):
+    ridx0 = UOp.range(100, 0)
+    d0 = UOp.param(0, dtypes.float.ptr())
+    ld = d0.index(ridx0.valid(ridx0<50))
+    w = ((ridx0<50) & (ridx0>30)).where(UOp.const(dtypes.float, 0), ld).cast(dtypes.half)
+    out = UOp.param(1, dtypes.half.ptr())
+    uops = to_uops_list([out.index(ridx0).store(w)])
+    for u in uops:
+      assert u.op is not Ops.WHERE
+
+  def test_where_in_store_becomes_gate(self):
+    ridx0 = UOp.range(100, 0)
+    d0 = UOp.param(0, dtypes.long.ptr())
+    idx = d0.index(ridx0)
+    ld = idx.load()
+    val = (ridx0<50).where(5, ld)
+    st = idx.store(val).end(ridx0)
+    uops = to_uops_list([st])
+    for u in uops:
+      assert u.op is not Ops.WHERE
+      if u.op is Ops.STORE: assert u.src[1].arg==5
+
+  def test_load_idx_becomes_int(self):
+    # mnist indexing with split reduceop
+    # Make sure we are not doign math on the loaded index, which would promote it to long
+    c0 = UOp.param(0, dtypes.uchar.ptr(128000))
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
+    c3 = UOp.param(1, dtypes.int.ptr(512))
+    c4 = c3.index(c1)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
+    c7 = UOp.param(2, dtypes.uchar.ptr(60000))
+    c8 = c7.index(c6)
+    c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
+    uops = to_uops_list([c10])
+    for u in uops:
+      self.assertNotEqual(u.dtype, dtypes.long)
+
+  def test_load_idx_no_math_on_loaded(self):
+    # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
+    c0 = UOp.param(0, dtypes.uchar.ptr(128000))
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
+    c3 = UOp.param(1, dtypes.int.ptr(512))
+    c4 = c3.index(c1)  # c4 is a load
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
+    c7 = UOp.param(2, dtypes.uchar.ptr(60000))
+    c8 = c7.index(c6)
+    # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
+    loaded_idx = c4.cast(dtypes.weakint)
+    comparison = (loaded_idx + c5) < UOp.const(dtypes.weakint, 60000)
+    c9 = comparison.where(c8.cast(dtypes.uint).cast(dtypes.uchar), 0).reduce(c5, arg=Ops.ADD)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
+    uops = to_uops_list([c10])
+    for u in uops:
+      self.assertNotEqual(u.dtype, dtypes.long)
+
+  def test_fold_gated_load(self):
+    glbl0 = UOp.param(0, dtypes.int.ptr())
+    glbl1 = UOp.param(1, dtypes.int.ptr())
+    glbl2 = UOp.param(2, dtypes.int.ptr())
+    idx = UOp.const(dtypes.int, 0)
+    ld0 = glbl1.index(UOp.invalid())
+    ld1 = glbl2.index(idx.valid(UOp.const(dtypes.bool, True)))
+    uops = to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0.index(idx), ld1+ld0))])
+    ld0 = uops[-2].src[-1]  # -2 to skip SINK
+    # the gate and invalid value are deleted from ld1
+    self.assertEqual(ld0, UOp.load(glbl2.index(idx, ptr=True), dtype=dtypes.int))
+
+  def test_fold_gated_load_local(self):
+    glbl0 = UOp.param(0, dtypes.int.ptr())
+    smem = UOp(Ops.DEFINE_LOCAL, dtypes.int.ptr(size=18, addrspace=AddrSpace.LOCAL), (), "temp")
+    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 16),), "lidx0")
+    st = UOp(Ops.STORE, dtypes.void, (smem.index(lidx, ptr=True), glbl0.index(lidx, ptr=True).load()))
+    barrier = UOp(Ops.BARRIER, dtypes.void, (st, ))
+    ld0 = smem.after(barrier).index(UOp.invalid())
+    ld1 = smem.after(barrier).index((lidx+2).valid(UOp.const(dtypes.bool, True)))
+    uops = to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0.index(lidx), ld1+ld0))])
+
+    ld0 = uops[-2].src[-1]  # -2 to skip SINK
+    # the gate and invalid value are deleted from ld1
+    self.assertEqual(ld0.src[0], smem.after(barrier).index(lidx+2, ptr=True))
+
+  def test_fold_gated_store(self):
+    glbl = UOp.param(0, dtypes.int.ptr())
+    idx0 = UOp.const(dtypes.int, 0)
+    idx1 = UOp.const(dtypes.int, 0)
+    val = UOp.const(dtypes.int, 42)
+    st0 = glbl.index(UOp.invalid(), ptr=True).store(val)
+    st1 = glbl.index(idx0.valid(UOp.const(dtypes.bool, True)), ptr=True).store(val)
+    uops = to_uops_list([st0, st1])
+    # only the second store happens
+    self.assertEqual(len(uops), 7)  # +1 for SINK, +1 for PARAM shape sentinel
+    self.assertEqual(uops[-2], glbl.index(idx1, ptr=True).store(val))  # -2 to skip SINK
+
+  @unittest.skip("this is a uop type error")
+  def test_asserts_bad_gate(self):
+    glbl0 = UOp.param(0, dtypes.int.ptr())
+    idx = UOp.const(dtypes.int, 0)
+    bad_gate = UOp.const(dtypes.int, 1)
+    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0, idx, UOp.const(dtypes.int, 42), bad_gate))])
+
+  def test_after_end(self):
+    r = UOp.range(10, 0)
+
+    c = r + 1
+    self.assertIn(r, c.ranges)
+
+    e = UOp.const(dtypes.int, 1).end(r)
+    self.assertNotIn(r, e.ranges)
+
+    a = c.after(e)
+    self.assertNotIn(r, a.ranges)
+
+@track_rewrites()
+def expander_rewrite(sink): return graph_rewrite(sink, sym + expander)
+
+class TestExpander(unittest.TestCase):
+  def test_expand_add_broadcast(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((1,4),))
+    sink = expander_rewrite(e1+3)
+    assert sink.op is Ops.UNROLL and len(const_values(sink.src[0])) == 4
+    self.assertTupleEqual(const_values(sink.src[0]), (3,4,5,6))
+
+  def test_contract_simple(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((1,4),))
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((1,4),))
+    sink = expander_rewrite(con)
+    self.assertEqual(sink.op, Ops.STACK)
+    self.assertTupleEqual(const_values(sink), (0,1,2,3))
+
+  def test_contract_axis_1(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(16), tuple(x for x in range(16))),), ((1,4),(2,4)))
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((1,4),))
+    sink = expander_rewrite(con)
+    vals = const_values(sink.src[0])
+    assert sink.op is Ops.UNROLL and len(vals) == 16 and sink.arg == ((2,4),)
+    assert sink.src[0].op is Ops.STACK
+    self.assertTupleEqual(vals[0:4], (0,4,8,12))
+    self.assertTupleEqual(vals[12:], (3,7,11,15))
+
+  def test_contract_axis_2(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(16), tuple(x for x in range(16))),), ((1,4),(2,4)))
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((2,4),))
+    sink = expander_rewrite(con)
+    vals = const_values(sink.src[0])
+    assert sink.op is Ops.UNROLL and len(vals) == 16 and sink.arg == ((1,4),)
+    assert sink.src[0].op is Ops.STACK
+    self.assertTupleEqual(vals[0:4], (0,1,2,3))
+    self.assertTupleEqual(vals[12:], (12,13,14,15))
+
+  def test_contract_axis_2_big(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(16), tuple(x for x in range(16))),), ((1,2),(2,2),(3,2),(4,2)))
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(2), (e1,), ((2,2),))
+    sink = expander_rewrite(con)
+    assert sink.op is Ops.UNROLL and sink.arg == ((1, 2), (3, 2), (4, 2))
+    vals = const_values(sink.src[0])
+    self.assertTupleEqual(vals[0:2], (0,4))
+    self.assertTupleEqual(vals[12:14], (10,14))
+
+  def test_contract_multi_axis(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(16), tuple(x for x in range(16))),), ((1,2),(2,2),(3,2),(4,2)))
+    sink = expander_rewrite(UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((3, 2), (2, 2))))
+    assert sink.op is Ops.UNROLL and sink.arg == ((1, 2), (4, 2))
+    self.assertTupleEqual(const_values(sink.src[0])[0:4], (0, 4, 2, 6))
+    sink = expander_rewrite(UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((2, 2), (3, 2))))
+    assert sink.op is Ops.UNROLL and sink.arg == ((1, 2), (4, 2))
+    self.assertTupleEqual(const_values(sink.src[0])[0:4], (0, 2, 4, 6))
+
+  def test_contract_mid(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(8), tuple(x for x in range(8))),), ((1,2),(2,2),(3,2)))
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(2), (e1,), ((2,2),))
+    sink = expander_rewrite(con)
+    assert sink.op is Ops.UNROLL and sink.arg == ((1,2),(3,2))
+    assert sink.src[0].op is Ops.STACK and len(const_values(sink.src[0])) == 8
+    self.assertTupleEqual(const_values(sink.src[0]), (0,2,1,3,4,6,5,7))
+
+  def test_contract_no_expand(self):
+    e1 = UOp.variable("i", 0, 10, dtype=dtypes.int)
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(2), (e1,), ((2,2),))
+    sink = expander_rewrite(con)
+    assert sink.op is Ops.STACK and len(sink.src) == 2
+    assert sink.src[0] == sink.src[1]
+
+  def test_contract_half_expand(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((1,4),))
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(8), (e1,), ((1,4), (2,2)))
+    sink = expander_rewrite(con)
+    vals = const_values(sink)
+    assert sink.op is Ops.STACK and len(vals) == 8
+    assert vals[0] == vals[1]
+    assert vals[0] != vals[2]
+    assert vals[6] == vals[7]
+
+  def test_expand_same_axis(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((1,4),))
+    e2 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(4*x for x in range(4))),), ((1,4),))
+    sink = expander_rewrite(e1+e2)
+    self.assertEqual(sink.op, Ops.UNROLL)
+    self.assertEqual(sink.src[0].op, Ops.STACK)
+    self.assertTupleEqual(const_values(sink.src[0]), (0,5,10,15))
+
+  def test_expand_different_axis(self, flip=False):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(4*x for x in range(4))),), ((1,4),))
+    e2 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((2,4),))
+    sink = expander_rewrite((e2+e1) if flip else (e1+e2))
+    vals = const_values(sink.src[0])
+    assert sink.op is Ops.UNROLL and len(vals) == 16
+    assert sink.arg == ((1, 4), (2, 4))
+    self.assertTupleEqual(vals, (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15))
+
+  def test_expand_different_axis_flip(self): self.test_expand_different_axis(True)
+
+  @unittest.skip("no longer supported")
+  def test_reduce_known_axis(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
+    sink = (3*e1).reduce(e1, arg=Ops.ADD)
+    sink = expander_rewrite(sink)
+    assert sink.op is Ops.CONST
+    self.assertEqual(sink.arg, 3*(0+1+2+3))
+
+  @unittest.skip("no longer supported")
+  def test_reduce_const(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
+    sink = UOp.const(dtypes.int, 3).reduce(e1, arg=Ops.ADD)
+    sink = expander_rewrite(sink)
+    assert sink.op is Ops.CONST
+    self.assertEqual(sink.arg, 3*4)
+
+  @unittest.skip("no longer supported")
+  def test_double_expand(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((2,4),))
+    e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, 4+x) for x in range(4)), ((2,4),))
+    e = UOp(Ops.UNROLL, dtypes.int, (e1, e2), ((1,2),))
+    sink = expander_rewrite(e)
+    assert sink.op is Ops.UNROLL and len(sink.src) == 8
+    assert sink.arg == ((1, 2), (2, 4))
+    self.assertListEqual([x.arg for x in sink.src], [0,1,2,3,4,5,6,7])
+
+  @unittest.skip("no longer supported")
+  def test_double_expand_reverse(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
+    e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, 4+x) for x in range(4)), ((1,4),))
+    e = UOp(Ops.UNROLL, dtypes.int, (e1, e2), ((2,2),))
+    sink = expander_rewrite(e)
+    assert sink.op is Ops.UNROLL and len(sink.src) == 8
+    assert sink.arg == ((1, 4), (2, 2))
+    self.assertListEqual([x.arg for x in sink.src], [0, 4, 1, 5, 2, 6, 3, 7])
+
+  @unittest.skip("no longer supported")
+  def test_double_expand_middle(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,2),(3,2)))
+    e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, 4+x) for x in range(4)), ((1,2),(3,2)))
+    e = UOp(Ops.UNROLL, dtypes.int, (e1, e2), ((2,2),))
+    sink = expander_rewrite(e)
+    assert sink.op is Ops.UNROLL and len(sink.src) == 8
+    assert sink.arg == ((1, 2), (2, 2), (3, 2))
+    self.assertListEqual([x.arg for x in sink.src], [0, 1, 4, 5, 2, 3, 6, 7])
+
+  # does this need to work?
+  @unittest.expectedFailure
+  @unittest.skip
+  def test_reduce_different_axis(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
+    e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((2,4),))
+    sink = e1.reduce(e2, arg=Ops.ADD)
+    sink = expander_rewrite(sink)
+    print(sink)
+
+class TestReduceCollapse(unittest.TestCase):
+  def test_multi_range_reduce_add(self):
+    """Test that (x + y).reduce(r1, r2) distributes over multiple ranges"""
+    from tinygrad.codegen.simplify import pm_reduce_collapse
+    # Create two ranges
+    r1 = UOp.range(3, 0)
+    r2 = UOp.range(4, 1)
+    # Create x + y where x and y depend on different ranges
+    x = r1.cast(dtypes.float)
+    y = r2.cast(dtypes.float)
+    # (x + y).reduce(r1, r2) should be rewritten
+    red = (x + y).reduce(r1, r2, arg=Ops.ADD)
+    self.assertEqual(len(red.src), 3)  # value + 2 ranges
+    result = graph_rewrite(red, pm_reduce_collapse, name='test')
+    # Should become add of two separate reduces
+    self.assertEqual(result.op, Ops.ADD)
+
+class TestMovementOps(unittest.TestCase):
+  def test_pm_mops_partial_reshape_index_removes_reshape(self):
+    from tinygrad.schedule.rangeify import pm_mops
+    src = UOp.param(0, dtypes.float, shape=(32, 4))
+    r0, r1 = UOp.range(4, 0), UOp.range(8, 1)
+    result = graph_rewrite(src.reshape((4, 8, 4)).index(r0, r1), pm_mops, name="test")
+    self.assertEqual(result.op, Ops.INDEX)
+    self.assertIs(result.src[0], src)
+    self.assertEqual(result.shape, (4,))
+    self.assertNotIn(Ops.RESHAPE, [u.op for u in result.toposort()])
+
+  def test_pm_mops_partial_reshape_index_suffix_mismatch_does_nothing(self):
+    from tinygrad.schedule.rangeify import pm_mops
+    src = UOp.param(0, dtypes.float, shape=(2, 6))
+    result = graph_rewrite(src.reshape((2, 3, 2)).index(UOp.range(2, 0)), pm_mops, name="test")
+    self.assertEqual(result.op, Ops.INDEX)
+    self.assertEqual(result.src[0].op, Ops.RESHAPE)
+
+class TestLoadStoreFolding(unittest.TestCase):
+  def test_gated_load_gep_preserves_alt(self):
+    """Test that LOAD(GEP, alt) preserves alt value after rewrite"""
+    from tinygrad.codegen.late.devectorizer import load_store_folding
+    buf = UOp.param(0, dtypes.float.vec(4).ptr())
+    idx = UOp.const(dtypes.int, 0)
+    gate = UOp.const(dtypes.bool, True)
+    gated_index = buf.index(idx.valid(gate))
+    gep = gated_index.gep(0)
+    alt = UOp.const(dtypes.float, 42.0)
+    gated_load = gep.load(alt)
+    self.assertEqual(len(gated_load.src), 2)  # GEP + alt
+    result = graph_rewrite(gated_load, load_store_folding, name='test')
+    # After rewrite, should still have alt value preserved
+    self.assertEqual(result.op, Ops.GEP)
+    inner_load = result.src[0]
+    self.assertEqual(inner_load.op, Ops.LOAD)
+    self.assertEqual(len(inner_load.src), 2)  # INDEX + alt
+
+  def test_gated_load_ptrcat_preserves_alt(self):
+    """Test that LOAD(PTRCAT, alt) preserves alt value after rewrite"""
+    from tinygrad.codegen.late.devectorizer import load_store_folding
+    buf1 = UOp.param(0, dtypes.float.ptr())
+    buf2 = UOp.param(1, dtypes.float.ptr())
+    idx = UOp.const(dtypes.int, 0)
+    idx1 = buf1.index(idx)
+    idx2 = buf2.index(idx)
+    ptrcat = UOp(Ops.PTRCAT, dtypes.float.ptr().vec(2), (idx1, idx2))
+    alt = UOp.const(dtypes.float.vec(2), 42.0)
+    gated_load = ptrcat.load(alt)
+    self.assertEqual(len(gated_load.src), 2)  # PTRCAT + alt
+    result = graph_rewrite(gated_load, load_store_folding, name='test')
+    # After rewrite, should be CAT of LOADs, each preserving alt
+    self.assertEqual(result.op, Ops.VCAT)
+    for inner_load in result.src:
+      self.assertEqual(inner_load.op, Ops.LOAD)
+      self.assertEqual(len(inner_load.src), 2)  # INDEX + alt
+    self.assertEqual(inner_load.src[1].arg, 42.0)  # alt value preserved
+
+class TestConstBufferize(unittest.TestCase):
+  def test_const_bufferize_with_ranges(self):
+    """Test that CONST.BUFFERIZE with ranges is folded correctly.
+
+    BUFFERIZE can have ranges as additional sources beyond the value.
+    The pattern at rangeify.py uses allow_any_len=True because
+    CONST doesn't depend on ranges (constant is same value everywhere).
+    """
+    from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
+    c = UOp.const(dtypes.float, 42.0)
+    r1 = UOp.range(3, 0)
+    bufferize_with_range = UOp(Ops.STAGE, dtypes.float, (c, r1), arg=BufferizeOpts(device="CPU"))
+    self.assertEqual(len(bufferize_with_range.src), 2)  # const + 1 range
+
+    result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test')
+    # BUFFERIZE should be removed, result is const broadcast to shape
+    self.assertNotEqual(result.op, Ops.STAGE)
+    const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
+    self.assertIn(42.0, const_vals)
+
+  def test_const_bufferize_with_multiple_ranges(self):
+    """Test CONST.BUFFERIZE with multiple ranges is also folded."""
+    from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
+    c = UOp.const(dtypes.float, 3.14)
+    r1 = UOp.range(3, 0)
+    r2 = UOp.range(4, 1)
+    bufferize_with_ranges = UOp(Ops.STAGE, dtypes.float, (c, r1, r2), arg=BufferizeOpts(device="CPU"))
+    self.assertEqual(len(bufferize_with_ranges.src), 3)  # const + 2 ranges
+
+    result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test')
+    # BUFFERIZE should be removed
+    self.assertNotEqual(result.op, Ops.STAGE)
+    const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
+    self.assertIn(3.14, const_vals)
+
+class TestUOpTags(unittest.TestCase):
+  def test_inc_by_one(self):
+    g = UOp.const(dtypes.int, 1) + UOp.const(dtypes.int, 1)
+    assert g.ssimplify() == 2
+    pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.replace(arg=x.arg+1, tag=1) if x.tag is None else None)])
+    pm_strip_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
+    g = graph_rewrite(g, pm_plus_1)
+    assert g.ssimplify() == 4
+    g = graph_rewrite(g, pm_plus_1)
+    assert g.ssimplify() == 4
+    g = graph_rewrite(g, pm_strip_tags)
+    assert g.ssimplify() == 4
+    g = graph_rewrite(g, pm_plus_1)
+    assert g.ssimplify() == 6
+
+class TestUOpGetItem(unittest.TestCase):
+  def _placeholder(self, shape, dtype=dtypes.half):
+    return UOp.placeholder(shape, dtype, slot=0, addrspace=AddrSpace.LOCAL)
+
+  # full slices (no shrink)
+  def test_full_slice(self):
+    p = self._placeholder((64, 64))
+    self.assertEqual(p[:, :].shape, (64, 64))
+  def test_full_slice_explicit(self):
+    p = self._placeholder((64, 64))
+    self.assertEqual(p[0:64, 0:64].shape, (64, 64))
+
+  # partial slices (shrink)
+  def test_shrink_cols(self):
+    p = self._placeholder((64, 80))
+    self.assertEqual(p[:, :64].shape, (64, 64))
+  def test_shrink_rows(self):
+    p = self._placeholder((80, 64))
+    self.assertEqual(p[:64, :].shape, (64, 64))
+  def test_shrink_both(self):
+    p = self._placeholder((80, 80))
+    self.assertEqual(p[:64, :64].shape, (64, 64))
+  def test_shrink_start(self):
+    p = self._placeholder((64, 64))
+    self.assertEqual(p[8:, :].shape, (56, 64))
+  def test_shrink_start_and_end(self):
+    p = self._placeholder((64, 64))
+    self.assertEqual(p[8:56, 4:60].shape, (48, 56))
+
+  # mixed slice and index
+  def test_index_and_slice(self):
+    p = self._placeholder((64, 80))
+    r = UOp.range(64, 100)
+    result = p[r, :64]
+    self.assertEqual(result.shape, (64,))
+  def test_slice_and_index(self):
+    p = self._placeholder((80, 64))
+    r = UOp.range(64, 100)
+    result = p[:64, r]
+    self.assertEqual(result.shape, (64,))
+  def test_shrink_then_index(self):
+    p = self._placeholder((64, 80))
+    s = p[:, :64]
+    r = UOp.range(64, 100)
+    result = s[r]
+    self.assertEqual(result.shape, (64,))
+
+  # integer index (no slice)
+  def test_int_index(self):
+    p = self._placeholder((64, 64))
+    result = p[0]
+    self.assertEqual(result.shape, (64,))
+
+  # ellipsis
+  def test_ellipsis_all_slices(self):
+    p = self._placeholder((64, 80))
+    self.assertEqual(p[..., :64].shape, (64, 64))
+  def test_ellipsis_with_int(self):
+    p = self._placeholder((64, 80))
+    r = UOp.range(64, 100)
+    result = p[..., r]
+    self.assertEqual(result.op, Ops.INDEX)
+  def test_ellipsis_only(self):
+    p = self._placeholder((64, 64))
+    self.assertEqual(p[...].shape, (64, 64))
+
+  # all slices should not create a bare INDEX
+  def test_all_slices_no_index(self):
+    p = self._placeholder((64, 80))
+    result = p[:, :64]
+    self.assertNotEqual(result.op, Ops.INDEX)
+  def test_all_full_slices_no_index(self):
+    p = self._placeholder((64, 64))
+    result = p[:, :]
+    self.assertNotEqual(result.op, Ops.INDEX)
+
+class TestUOpBroadcast(unittest.TestCase):
+  def test_broadcast_row(self):
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    b = UOp.const(dtypes.float, 2, shape=(4, 1))
+    c = a + b
+    self.assertEqual(c.shape, (4, 8))
+    self.assertEqual(c.op, Ops.ADD)
+
+  def test_broadcast_col(self):
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    b = UOp.const(dtypes.float, 2, shape=(1, 8))
+    c = a + b
+    self.assertEqual(c.shape, (4, 8))
+    self.assertEqual(c.op, Ops.ADD)
+
+  def test_broadcast_lower_dim(self):
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    b = UOp.const(dtypes.float, 2, shape=(8,))
+    c = a * b
+    self.assertEqual(c.shape, (4, 8))
+    self.assertEqual(c.op, Ops.MUL)
+
+  def test_broadcast_scalar(self):
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    c = a * 2
+    self.assertEqual(c.shape, (4, 8))
+    self.assertEqual(c.op, Ops.MUL)
+
+  def test_broadcast_symbolic_same_shape(self):
+    t = Variable("t", 1, 10)
+    a = UOp.const(dtypes.float, 1, shape=(1, 1, t))
+    b = UOp.const(dtypes.float, 2, shape=(1, 1, t))
+    c = a + b
+    self.assertEqual(c.op, Ops.ADD)
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)

tinygrad_repo/test/null/test_uop_repr.py

@@ -0,0 +1,38 @@
+import unittest
+from tinygrad import UOp, dtypes
+
+class TestUOpRepr(unittest.TestCase):
+  def test_simple_const(self):
+    a = UOp.const(dtypes.int, 42)
+    self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.int, arg=42, src=())")
+  def test_different_consts(self):
+    a, b = UOp.const(dtypes.int, 42), UOp.const(dtypes.int, 3)
+    expected = (
+      "UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
+      "  UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" +
+      "  UOp(Ops.CONST, dtypes.int, arg=3, src=()),))"
+    )
+    self.assertEqual(repr(a+b), expected)
+  def test_walrus_operator_indentation(self):
+    # The reference should have the same indentation as the definition
+    a = UOp.const(dtypes.int, 42)
+    expected = (
+      "UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" +
+      "  x0,))"
+    )
+    self.assertEqual(repr(a+a), expected)
+  def test_nested_walrus_indentation(self):
+    # Ensure indentation is consistent at multiple levels
+    b = (a:=UOp.const(dtypes.int, 1)) + a
+    expected = (
+      "UOp(Ops.MUL, dtypes.int, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
+      "    x1:=UOp(Ops.CONST, dtypes.int, arg=1, src=()),\n" +
+      "    x1,)),\n" +
+      "  x0,))"
+    )
+    self.assertEqual(repr(b*b), expected)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_uop_resolve.py

@@ -0,0 +1,128 @@
+import unittest
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import UOp, resolve
+
+class TestUOpResolve(unittest.TestCase):
+  def test_simple_int(self):
+    u = UOp.const(dtypes.int, 4)
+    self.assertEqual(int(u), 4)
+
+  def test_int_add(self):
+    u = UOp.const(dtypes.int, 4) + 7
+    self.assertEqual(int(u), 11)
+
+  def test_lt(self):
+    u = UOp.const(dtypes.int, 4) < 7
+    self.assertTrue(u)
+
+  def test_rfloordiv(self):
+    u = 8 // UOp.const(dtypes.int, 4)
+    self.assertEqual(int(u), 2)
+
+  def test_rtruediv(self):
+    u = 9 / UOp.const(dtypes.float, 4)
+    self.assertEqual(float(u), 2.25)
+
+  def test_leq(self):
+    u = UOp.const(dtypes.int, 4) <= 4
+    self.assertTrue(u)
+
+  def test_ne(self):
+    u = UOp.const(dtypes.int, 4) != 7
+    self.assertTrue(u)
+
+  def test_ne_f(self):
+    u = UOp.const(dtypes.int, 4) != 4
+    self.assertFalse(u)
+
+  def test_ngt(self):
+    u = UOp.const(dtypes.int, 4) > 7
+    self.assertFalse(u)
+
+  def test_ssimplify(self):
+    self.assertEqual((8 % UOp.const(dtypes.int, 4)).ssimplify(), 0)
+    self.assertEqual((8 * UOp.const(dtypes.int, 4)).ssimplify(), 32)
+
+  def test_ambiguous_less_than(self):
+    u = UOp.variable("i", 1, 10)
+    self.assertTrue(resolve(u < 4))
+    self.assertFalse(resolve(u < 4, False))
+    self.assertTrue(resolve(u < 11, False))
+    self.assertFalse(resolve(u < -1, False))
+    self.assertFalse(resolve(u < -1, True))
+
+  def test_float_direct(self):
+    u = UOp.const(dtypes.float, 4.5) + 7
+    self.assertEqual(float(u), 11.5)
+
+  def test_var_cmp_t(self):
+    u = UOp.variable("i", 1, 10) < 20
+    self.assertTrue(u)
+
+  def test_var_cmp_t2(self):
+    u = UOp.variable("i", 1, 10)//2 < 20
+    self.assertTrue(u)
+
+  def test_var_cmp_f(self):
+    u = UOp.variable("i", 1, 10) < 1
+    self.assertFalse(u)
+
+  def test_var_cmp_f2(self):
+    u = UOp.variable("i", 1, 10) > 11
+    self.assertFalse(u)
+
+  def test_or_true(self):
+    u = UOp.variable("b", False, True, dtypes.bool) | True
+    self.assertTrue(u)
+
+  def test_or_false(self):
+    with self.assertRaises(ValueError):
+      u = UOp.variable("b", False, True, dtypes.bool) | False
+      self.assertTrue(u)
+
+  def test_and_false(self):
+    u = UOp.variable("b", False, True, dtypes.bool) & False
+    self.assertFalse(u)
+
+  def test_max(self):
+    x = UOp.variable("x", 1, 10)
+    y = UOp.variable("y", 5, 10)
+    u = x.maximum(y)
+    self.assertTrue(u < 20)
+    self.assertFalse(u < 3)
+
+  def test_x_lt_x(self):
+    x = UOp.variable("i", 1, 10)
+    self.assertFalse(x < x)
+
+  def test_x_lt_xp1(self):
+    x = UOp.variable("i", 1, 10)
+    u = x < (x+1)
+    # TODO: improve
+    with self.assertRaises(ValueError):
+      bool(u)
+
+  def test_and_true(self):
+    u = UOp.variable("b", False, True, dtypes.bool) & True
+    with self.assertRaises(ValueError):
+      bool(u)
+
+  def test_var_cmp_range(self):
+    v = UOp.variable("i", 1, 10)
+    u = (v > 4) | (v < 6)
+    # TODO: improve
+    with self.assertRaises(ValueError):
+      bool(u)
+
+  def test_var_cmp_assert(self):
+    with self.assertRaises(ValueError):
+      u = UOp.variable("i", 1, 10) < 5
+      self.assertFalse(u)
+
+  def test_plus_ordering_lt(self):
+    i = UOp.variable("i", 1, 10)
+    j = UOp.variable("j", 1, 10)
+    self.assertFalse((i+j) < (j+i))
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_uop_vmin_vmax.py

@@ -0,0 +1,406 @@
+import unittest, math
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.dtype import dtypes, Invalid
+
+class TestVminVmaxProperties(unittest.TestCase):
+  def test_vmin_vmax_constant(self):
+    # vmin and vmax for a constant
+    uop = UOp.const(dtypes.int32, 42)
+    self.assertEqual(uop.vmin, 42)
+    self.assertEqual(uop.vmax, 42)
+
+  def test_vmin_vmax_cmpne(self):
+    uop = UOp.const(dtypes.int32, 42)
+    def test_bool(u, x):
+      self.assertEqual(u.vmin, x)
+      self.assertEqual(u.vmax, x)
+    test_bool(uop != 42, False)
+    test_bool(uop != 43, True)
+    test_bool(uop != 41, True)
+
+  def test_vmin_vmax_addition_with_variable(self):
+    # vmin and vmax for addition with a variable
+    x = UOp.variable('x', 10, 20)
+    uop = x + 5
+    self.assertEqual(uop.vmin, 15)
+    self.assertEqual(uop.vmax, 25)
+
+  def test_vmin_vmax_subtraction_with_variable(self):
+    x = UOp.variable('x', 10, 20)
+    uop = x - 5
+    self.assertEqual(uop.vmin, 5)
+    self.assertEqual(uop.vmax, 15)
+    uop = 5 - x
+    self.assertEqual(uop.vmin, -15)
+    self.assertEqual(uop.vmax, -5)
+
+  def test_vmin_vmax_and_with_variable(self):
+    x = UOp.variable('x', 10, 20)
+    uop = x & 5
+    self.assertEqual(uop.vmin, 0)
+    self.assertEqual(uop.vmax, 5)
+
+    uop = x & 15
+    self.assertEqual(uop.vmin, 0)
+    self.assertEqual(uop.vmax, 15)
+
+    # TODO: this can be improved
+    uop = x & 32
+    self.assertEqual(uop.vmin, 0)
+    self.assertEqual(uop.vmax, 20) # shoud be 0
+
+  def test_vmin_vmax_and_with_negative_variable(self):
+    # when mask doesn't have sign bit set, result is always non-negative
+    x = UOp.variable('x', -100, 100, dtypes.int32)
+    # 511 = 0x1FF, doesn't have sign bit set for int32
+    uop = x & 511
+    self.assertEqual(uop.vmin, 0)
+    self.assertEqual(uop.vmax, 511)
+
+    # 0x7FFFFFFF is max positive int32, doesn't have sign bit
+    uop = x & 0x7FFFFFFF
+    self.assertEqual(uop.vmin, 0)
+    self.assertEqual(uop.vmax, 0x7FFFFFFF)
+
+    # negative mask: x & -1 could be anything since -1 has all bits set
+    uop = x & -1
+    self.assertEqual(uop.vmin, dtypes.int32.min)
+    self.assertEqual(uop.vmax, dtypes.int32.max)
+
+  def test_vmin_vmax_multiplication_with_variable(self):
+    # vmin and vmax for multiplication with a variable
+    x = UOp.variable('x', -3, 4)
+    uop = x * 2
+    self.assertEqual(uop.vmin, -6)
+    self.assertEqual(uop.vmax, 8)
+
+  def test_vmin_vmax_variable_inside_special(self):
+    uop = UOp(Ops.SPECIAL, dtypes.int, arg='gidx0', src=(UOp(Ops.DEFINE_VAR, dtypes.int, arg=('i', 1, 10)),))
+    self.assertEqual(uop.vmin, 0)
+    self.assertEqual(uop.vmax, 9)
+
+  def test_vmin_vmax_multiplication_0_inf(self):
+    # vmin and vmax for multiplication with a variable
+    x = UOp.const(dtypes.float, 0.0)
+    y = UOp.load(UOp.param(0, dtypes.float.ptr(1)), UOp.const(dtypes.int, 0), dtype=dtypes.float)
+    uop = x * y
+    # TODO: these should be 0, but definitely should not be nan
+    self.assertEqual(uop.vmin, -math.inf)
+    self.assertEqual(uop.vmax, math.inf)
+
+  def test_vmin_vmax_with_negative_multiplication(self):
+    # vmin and vmax when multiplying by a negative number
+    x = UOp.variable('x', 2, 5)
+    uop = x * -3
+    self.assertEqual(uop.vmin, -15)
+    self.assertEqual(uop.vmax, -6)
+
+  def test_vmin_vmax_with_negative_multiplication2(self):
+    # vmin and vmax when multiplying by a negative number
+    x = UOp.variable('x', -2, 5)
+    uop = x * -3
+    self.assertEqual(uop.vmin, -15)
+    self.assertEqual(uop.vmax, 6)
+
+  def test_vmin_vmax_nested_min_max(self):
+    # vmin and vmax with nested min/max operations
+    x = UOp.variable('x', 0, 10)
+    uop = x.maximum(5).minimum(8)
+    self.assertEqual(uop.vmin, 5)
+    self.assertEqual(uop.vmax, 8)
+
+  def test_vmin_vmax_where(self):
+    x = UOp.variable('x', 0, 10)
+    y = UOp.variable('y', 1, 11)
+    z = UOp.variable('z', 2, 12)
+    uop = (x<5).where(y, z)
+    self.assertEqual(uop.vmin, 1)
+    self.assertEqual(uop.vmax, 12)
+
+  def test_vmin_vmax_shl(self):
+    x = UOp.variable('x', 0, 10) << 5
+    self.assertEqual(x.vmin, 0)
+    self.assertEqual(x.vmax, 10 << 5)
+
+  def test_vmin_vmax_shr(self):
+    x = UOp.variable('x', 0, 10) >> 2
+    self.assertEqual(x.vmin, 0)
+    self.assertEqual(x.vmax, 10 >> 2)
+
+  def test_vmin_vmax_xor_neg1(self):
+    x = UOp.variable('x', 3, 7)
+    uop = x ^ -1
+    self.assertEqual(uop.vmin, ~7)
+    self.assertEqual(uop.vmax, ~3)
+    # negative range
+    y = UOp.variable('y', -10, -3)
+    uop2 = y ^ -1
+    self.assertEqual(uop2.vmin, ~(-3))
+    self.assertEqual(uop2.vmax, ~(-10))
+    # range spanning zero
+    z = UOp.variable('z', -5, 6)
+    uop3 = z ^ -1
+    self.assertEqual(uop3.vmin, ~6)
+    self.assertEqual(uop3.vmax, ~(-5))
+
+  def test_vmin_vmax_cast(self):
+    x = UOp.variable('x', -10, 10, dtypes.int)
+    x_float = x.cast(dtypes.float)
+    self.assertEqual(x_float.vmin, -10)
+    self.assertEqual(x_float.vmax, 10)
+    x_bool = x.cast(dtypes.bool)
+    self.assertEqual(x_bool.vmin, False)
+    self.assertEqual(x_bool.vmax, True)
+    x_uint = x.cast(dtypes.uint)
+    self.assertEqual(x_uint.vmin, dtypes.uint.min)
+    self.assertEqual(x_uint.vmax, dtypes.uint.max)
+
+  def test_vmin_vmax_invalid(self):
+    i = UOp.invalid()
+    self.assertNotEqual(i.vmin, i.vmax)
+
+  def test_vmin_vmax_invalid_vconst(self):
+    x = UOp.const(dtypes.weakint.vec(4), (0, 4, Invalid, Invalid))
+    self.assertLess(x.vmin, 0)
+    self.assertGreater(x.vmax, 4)
+
+class TestVminVmaxDivMod(unittest.TestCase):
+  def test_vmin_vmax_division_positive(self):
+    # vmin and vmax for division of a variable by a positive constant
+    x = UOp.variable('x', 10, 20)
+    uop = x // 2
+    self.assertEqual(uop.vmin, 5)
+    self.assertEqual(uop.vmax, 10)
+
+  def test_vmin_vmax_division_negative(self):
+    # floor division of a variable by a negative constant
+    x = UOp.variable('x', 10, 20)
+    uop = x // -2
+    self.assertEqual(uop.vmin, -10)
+    self.assertEqual(uop.vmax, -5)
+    uop = x // -3
+    self.assertEqual(uop.vmin, -7)
+    self.assertEqual(uop.vmax, -4)
+
+    x = UOp.variable('x', -20, -10)
+    uop = x // -2
+    self.assertEqual(uop.vmin, 5)
+    self.assertEqual(uop.vmax, 10)
+    uop = x // -3
+    self.assertEqual(uop.vmin, 3)
+    self.assertEqual(uop.vmax, 6)
+
+  def test_vmin_vmax_floordiv_floormod(self):
+    x = UOp.variable('x', -7, 7)
+    floordiv = x.alu(Ops.FLOORDIV, x.const_like(3))
+    self.assertEqual(floordiv.vmin, -3)
+    self.assertEqual(floordiv.vmax, 2)
+    floormod = x.alu(Ops.FLOORMOD, x.const_like(3))
+    self.assertEqual(floormod.vmin, 0)
+    self.assertEqual(floormod.vmax, 2)
+    # negative const divisor: floormod range is [c+1, 0]
+    floormod_neg = x.alu(Ops.FLOORMOD, x.const_like(-3))
+    self.assertEqual(floormod_neg.vmin, -2)
+    self.assertEqual(floormod_neg.vmax, 0)
+
+    # cross 0
+    x = UOp.variable('x', -10, 10)
+    uop = x // -2
+    self.assertEqual(uop.vmin, -5)
+    self.assertEqual(uop.vmax, 5)
+    uop = x // -3
+    self.assertEqual(uop.vmin, -4)
+    self.assertEqual(uop.vmax, 3)
+
+  def test_vmin_vmax_floordiv_floormod_empty_range(self):
+    # empty numerator range (vmin > vmax, e.g. RANGE with end=0) short-circuits to (0, 0)
+    rng = UOp.range(0, 0)
+    self.assertEqual(rng.vmin, 0)
+    self.assertEqual(rng.vmax, -1)
+    self.assertEqual((rng // 4).vmin, 0)
+    self.assertEqual((rng // 4).vmax, 0)
+    self.assertEqual((rng % 4).vmin, 0)
+    self.assertEqual((rng % 4).vmax, 0)
+
+  def test_vmin_vmax_div_symbolic(self):
+    x = UOp.variable('x', 1, 10)
+    y = UOp.variable('y', 3, 5)
+    self.assertEqual((x//y).vmin, 0)
+    self.assertEqual((x//y).vmax, 3)
+    self.assertEqual(((-x)//y).vmin, -4)
+    self.assertEqual(((-x)//y).vmax, -1)
+    self.assertEqual((x//(-y)).vmin, -4)
+    self.assertEqual((x//(-y)).vmax, -1)
+    self.assertEqual(((-x)//(-y)).vmin, 0)
+    self.assertEqual(((-x)//(-y)).vmax, 3)
+
+    self.assertEqual((100//y).vmin, 20)
+    self.assertEqual((100//y).vmax, 33)
+    self.assertEqual(((-100)//y).vmin, -34)
+    self.assertEqual(((-100)//y).vmax, -20)
+    self.assertEqual((100//(-y)).vmin, -34)
+    self.assertEqual((100//(-y)).vmax, -20)
+    self.assertEqual(((-100)//(-y)).vmin, 20)
+    self.assertEqual(((-100)//(-y)).vmax, 33)
+
+  def test_vmin_vmax_mod_positive(self):
+    # floor mod with positive divisor: result in [0, c-1] regardless of dividend sign
+    positive = UOp.variable('positive', 10, 20)
+    uop = positive % 3
+    self.assertEqual(uop.vmin, 0)
+    self.assertEqual(uop.vmax, 2)
+
+    negative = UOp.variable('negative', -20, -10)
+    uop = negative % 3
+    self.assertEqual(uop.vmin, 0)
+    self.assertEqual(uop.vmax, 2)
+
+    mixed = UOp.variable('mixed', -20, 20)
+    uop = mixed % 3
+    self.assertEqual(uop.vmin, 0)
+    self.assertEqual(uop.vmax, 2)
+
+  def test_vmin_vmax_mod_negative(self):
+    # floor mod with negative divisor: result in [c+1, 0] regardless of dividend sign
+    positive = UOp.variable('positive', 10, 20)
+    uop = positive % -3
+    self.assertEqual(uop.vmin, -2)
+    self.assertEqual(uop.vmax, 0)
+
+    negative = UOp.variable('negative', -20, -10)
+    uop = negative % -3
+    self.assertEqual(uop.vmin, -2)
+    self.assertEqual(uop.vmax, 0)
+
+    mixed = UOp.variable('mixed', -20, 20)
+    uop = mixed % -3
+    self.assertEqual(uop.vmin, -2)
+    self.assertEqual(uop.vmax, 0)
+
+class TestVminVmaxVConst(unittest.TestCase):
+  def test_vmin_vmax_vconst_single_element(self):
+    # vmin and vmax for a single-element vector constant
+    uop = UOp.const(dtypes.int32.vec(1), (42,))
+    self.assertEqual(uop.vmin, 42)
+    self.assertEqual(uop.vmax, 42)
+
+  def test_vmin_vmax_vconst_multiple_elements(self):
+    # vmin and vmax for a multi-element vector constant
+    uop = UOp.const(dtypes.int32.vec(4), (10, 20, -5, 7))
+    self.assertEqual(uop.vmin, -5)
+    self.assertEqual(uop.vmax, 20)
+
+  def test_vmin_vmax_vconst_all_equal(self):
+    # vmin and vmax for a vector where all elements are equal
+    uop = UOp.const(dtypes.int32.vec(3), (7, 7, 7))
+    self.assertEqual(uop.vmin, 7)
+    self.assertEqual(uop.vmax, 7)
+
+  def test_vmin_vmax_vconst_with_negative_values(self):
+    # vmin and vmax for a vector constant containing negative values
+    uop = UOp.const(dtypes.int32.vec(4), (-10, -20, -5, -15))
+    self.assertEqual(uop.vmin, -20)
+    self.assertEqual(uop.vmax, -5)
+
+  def test_vmin_vmax_vconst_with_floats(self):
+    # vmin and vmax for a vector constant of float values
+    uop = UOp.const(dtypes.float32.vec(3), (1.5, -3.2, 0.0))
+    self.assertEqual(uop.vmin, -3.2)
+    self.assertEqual(uop.vmax, 1.5)
+
+  def test_vmin_vmax_vconst_with_bools(self):
+    # vmin and vmax for a vector constant of bool values
+    uop = UOp.const(dtypes.bool.vec(3), (True, False, False))
+    self.assertIs(uop.vmin, False)
+    self.assertIs(uop.vmax, True)
+
+  def test_vmin_vmax_vector_with_gep(self):
+    # vmin and vmax for a vector constant of bool values
+    d1 = UOp.param(1, dtypes.int.ptr())
+    idx = UOp.const(dtypes.int, 0)
+    val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx).cast(dtypes.int.vec(2).ptr()),))
+    uop = (val // 32).gep(0)
+    self.assertEqual(uop.vmin, -67108864)
+    self.assertEqual(uop.vmax, 67108863)
+
+class TestConstFactor(unittest.TestCase):
+  def test_const_factor_constant(self):
+    # const_factor for a constant
+    uop = UOp.const(dtypes.int32, 42)
+    self.assertEqual(uop.const_factor(), 42)
+
+  def test_const_factor_addition(self):
+    # const_factor for an addition of constants
+    uop = UOp.const(dtypes.int32, 30) + UOp.const(dtypes.int32, 12)
+    self.assertEqual(uop.const_factor(), 6)  # GCD(30, 12) = 6
+
+  def test_const_factor_multiplication(self):
+    # const_factor for a multiplication of constants
+    uop = UOp.const(dtypes.int32, 5) * UOp.const(dtypes.int32, 7)
+    self.assertEqual(uop.const_factor(), 5)  # For multiplication, it's one of the factors
+
+  def test_const_factor_with_variable(self):
+    # const_factor for an expression involving a variable
+    x = UOp.variable('x', 10, 20)
+    uop = x * 3
+    self.assertEqual(uop.const_factor(), 3)
+
+  def test_const_factor_division(self):
+    # const_factor for an expression with division
+    x = UOp.variable('x', 10, 20)
+    uop = x // 4
+    self.assertEqual(uop.const_factor(), 1)  # Division reduces the const_factor to 1
+
+  def test_const_factor_multiplication_of_var_and_const(self):
+    # const_factor for multiplication of a variable and a constant
+    x = UOp.variable('x', 6, 18)
+    uop = x * 4
+    self.assertEqual(uop.const_factor(), 4)  # Constant factor 4
+
+  @unittest.skip("broken")
+  def test_const_factor_multiplication_of_consts_and_vars(self):
+    # Multiplying constants and variables
+    x = UOp.variable('x', 10, 20)
+    uop = (x * 3) * 5
+    self.assertEqual(uop.const_factor(), 15)  # Constant multipliers are combined (3 * 5 = 15)
+
+class TestDivides(unittest.TestCase):
+  def test_divides_constant_exact(self):
+    # Divides a constant by an exact divisor
+    uop = UOp.const(dtypes.int32, 42)
+    result = uop.divides(7)
+    self.assertIsNotNone(result)
+    self.assertEqual(result.const_factor(), 6)  # 42 / 7 = 6
+
+  def test_divides_constant_inexact(self):
+    # Try to divide a constant by a non-exact divisor
+    uop = UOp.const(dtypes.int32, 42)
+    result = uop.divides(5)
+    self.assertIsNone(result)  # 42 is not divisible by 5
+
+  @unittest.skip("broken")
+  def test_divides_variable_and_constant(self):
+    # Multiplying a variable by a constant, then dividing by the same constant
+    x = UOp.variable('x', 10, 20)
+    uop = x * 6
+    result = uop.divides(6)
+    self.assertIsNotNone(result)
+    self.assertEqual(result, x)  # (x * 6) / 6 = x
+
+  def test_divides_complex_expression(self):
+    # Dividing a more complex expression
+    x = UOp.variable('x', 10, 20)
+    uop = (x * 6) + 18
+    result = uop.divides(6)
+    self.assertIsNotNone(result)
+    self.assertEqual(result.const_factor(), 1)  # (x + 3), const_factor is 1
+
+  def test_divides_with_inexact_factors(self):
+    # Multiplying by a constant but dividing by a non-exact divisor
+    x = UOp.variable('x', 15, 45)
+    uop = x * 4
+    result = uop.divides(3)
+    self.assertIsNone(result)  # Cannot divide by 3, since 4 is not divisible by 3
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_uops.py

@@ -0,0 +1,370 @@
+# uops tests that pass on NULL backend (no copyout needed)
+import unittest
+import numpy as np
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import Timing, Context, cdiv
+from tinygrad.dtype import dtypes, ConstFloat  # noqa: F401
+from tinygrad.device import Device
+from tinygrad.uop.ops import Ops, UOp, UPat, exec_alu
+from tinygrad.uop.spec import spec_shared
+from tinygrad.uop.symbolic import sym
+from test.helpers import eval_uop, to_uops_list
+
+class TestSafeCast(unittest.TestCase):
+  def test_cast_folds(self):
+    a = UOp.variable("a", 1, 10, dtype=dtypes.int32)
+    self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int32).simplify(), a)
+    self.assertEqual(a.cast(dtypes.double).cast(dtypes.int32).simplify(), a)
+    a = UOp.variable("a", 1, 10, dtype=dtypes.uint8)
+    self.assertEqual(a.cast(dtypes.int64).cast(dtypes.uint8).simplify(), a)
+    self.assertEqual(a.cast(dtypes.uint32).cast(dtypes.uint8).simplify(), a)
+
+  def test_remove_intermediate_cast(self):
+    a = UOp.variable("a", 0., 100., dtype=dtypes.half)
+    self.assertEqual(a.cast(dtypes.double).cast(dtypes.float).simplify(), a.cast(dtypes.float))
+    a = UOp.variable("a", 1, 10, dtype=dtypes.int32)
+    # TODO: double preserves certain int dtypes
+    self.assertEqual(a.cast(dtypes.double).cast(dtypes.float).simplify(), a.cast(dtypes.float))
+    self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int16).simplify(), a.cast(dtypes.int16))
+    a = UOp.variable("a", 1, 10, dtype=dtypes.uint8)
+    self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int32).simplify(), a.cast(dtypes.int32))
+
+  def test_safe_cast_using_bounds(self):
+    a = UOp.variable("a", 1, 10, dtype=dtypes.uint64)
+    self.assertEqual(a.cast(dtypes.int16).cast(dtypes.int).simplify(), a.cast(dtypes.int))
+    a = UOp.variable("a", -10, 10, dtype=dtypes.int32)
+    self.assertEqual(a.cast(dtypes.int8).cast(dtypes.int64).simplify(), a.cast(dtypes.int64))
+    self.assertEqual(a.cast(dtypes.int8).cast(dtypes.float).simplify(), a.cast(dtypes.float))
+
+class TestExecALU(unittest.TestCase):
+  def test_sqrt(self):
+    self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0)
+
+  def test_div(self):
+    self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (8, 2)), 4)
+    self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (7, 3)), 2)
+    self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (7, -3)), -2)
+    self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (-50, 6)), -8)
+
+  def test_floordiv(self):
+    self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (8, 2)), 4)
+    self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (7, 3)), 2)
+    self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (7, -3)), -3)
+    self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (-7, 3)), -3)
+    self.assertEqual(exec_alu(Ops.FLOORDIV, dtypes.int8, (-50, 6)), -9)
+
+  def test_floormod(self):
+    self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (8, 2)), 0)
+    self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (7, 3)), 1)
+    self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (7, -3)), -2)
+    self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (-7, 3)), 2)
+    self.assertEqual(exec_alu(Ops.FLOORMOD, dtypes.int8, (-50, 6)), 4)
+
+    np.testing.assert_allclose(exec_alu(Ops.MUL, dtypes.float32, (7.0, exec_alu(Ops.RECIPROCAL, dtypes.float32, (3.0,)))), 2+(1.0/3.0))
+    np.testing.assert_allclose(exec_alu(Ops.MUL, dtypes.float32, (7.0, exec_alu(Ops.RECIPROCAL, dtypes.float32, (-3.0,)))), -2-(1.0/3.0))
+
+  def test_recip(self):
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (8,)), 1/8)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (7,)), 1/7)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (-3,)), 1/-3)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (-50,)), 1/-50)
+
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, ((32+521+3),)), 1/(32+521+3))
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, ((34**2),)), 1/(34**2))
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (10,)), 1/10)
+
+  def test_bool_cmplt(self):
+    self.assertEqual(exec_alu(Ops.CMPLT, dtypes.bool, (False, False)), False)
+    self.assertEqual(exec_alu(Ops.CMPLT, dtypes.bool, (False, True)), True)
+    self.assertEqual(exec_alu(Ops.CMPLT, dtypes.bool, (True, False)), False)
+    self.assertEqual(exec_alu(Ops.CMPLT, dtypes.bool, (True, True)), False)
+
+  def test_bool_cmpne(self):
+    self.assertEqual(exec_alu(Ops.CMPNE, dtypes.bool, (False, False)), False)
+    self.assertEqual(exec_alu(Ops.CMPNE, dtypes.bool, (False, True)), True)
+    self.assertEqual(exec_alu(Ops.CMPNE, dtypes.bool, (True, False)), True)
+    self.assertEqual(exec_alu(Ops.CMPNE, dtypes.bool, (True, True)), False)
+
+  def test_bool_where(self):
+    self.assertEqual(exec_alu(Ops.WHERE, dtypes.bool, (False, False, False)), False)
+    self.assertEqual(exec_alu(Ops.WHERE, dtypes.int, (False, 2, 4)), 4)
+    np.testing.assert_allclose(exec_alu(Ops.WHERE, dtypes.float, (False, 2.2, 4.5)), 4.5)
+
+  def test_overflow(self):
+    self.assertEqual(exec_alu(Ops.ADD, dtypes.uint8, (250, 250)), 244)
+    self.assertEqual(exec_alu(Ops.ADD, dtypes.uint8, (256, 0)), 0)
+    self.assertEqual(exec_alu(Ops.ADD, dtypes.uint8, (0, -1)), 255)
+    self.assertEqual(exec_alu(Ops.ADD, dtypes.uint8, (0, -1000)), 24)
+
+    self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (127, 0)), 127)
+    self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (-128, 0)), -128)
+    self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (-100, -100)), 56)
+    self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (-1000, -0)), 24)
+    self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (-130, -0)), 126)
+
+    self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (1, 1)), 2)
+    self.assertEqual(exec_alu(Ops.ADD, dtypes.int8, (-128, 0)), -128)
+
+    # test no truncate
+    self.assertEqual(exec_alu(Ops.ADD, dtypes.uint8, (250, 250), truncate_output=False), 500)
+
+class TestGatedStoreRewrite(unittest.TestCase):
+  def test_tiny_gate_store(self):
+    gmem = UOp.param(0, dtypes.float.ptr())
+    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
+    gate = gidx0<UOp.const(dtypes.int, 1)
+    idx = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem, (gidx0 * UOp.const(dtypes.int, 2)).valid(gate)))
+    val = UOp.const(dtypes.float, 42.0)
+    store = UOp(Ops.STORE, dtypes.void, (idx, val))
+    uops = to_uops_list([store])
+    if_uop = next(u for u in uops if u.op is Ops.IF)
+    endif = next(u for u in uops if u.op is Ops.ENDIF)
+    assert endif.src[0] is if_uop
+    gated_uops = tuple(uops[uops.index(if_uop)+1:uops.index(endif)])
+    self.assertEqual(len(gated_uops), 1)
+    self.assertIs(gated_uops[-1].op, Ops.STORE)
+    self.assertEqual(len(gated_uops[-1].src), 2)
+
+  def test_gate_some_stores(self):
+    gmem0 = UOp.param(0, dtypes.float.ptr())
+    gmem1 = UOp.param(1, dtypes.float.ptr())
+    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
+    idx = gidx0 * UOp.const(dtypes.int, 2)
+    idx0 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem0, idx.valid(gidx0<UOp.const(dtypes.int, 1))))
+    idx1 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem1, idx))
+    val = UOp.const(dtypes.float, 42.0)
+    stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
+    uops = to_uops_list(stores)
+    if_uop = next(u for u in uops if u.op is Ops.IF)
+    endif = next(u for u in uops if u.op is Ops.ENDIF)
+    assert endif.src[0] is if_uop
+    gated_uops = tuple(uops[uops.index(if_uop)+1:uops.index(endif)])
+    self.assertEqual(len(gated_uops), 1)
+    self.assertIs(gated_uops[-1].op, Ops.STORE)
+    self.assertEqual(len(gated_uops[-1].src), 2)
+
+  # scaled down version of TestLinearizerDumb.test_unmerged_ifs
+  @unittest.skip("we don't merge ifs anymore")
+  def test_merge_ifs_alt(self):
+    gmem0 = UOp.param(0, dtypes.float.ptr())
+    gmem1 = UOp.param(1, dtypes.float.ptr())
+    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
+    idx = gidx0*UOp.const(dtypes.int, 2)
+    gate = gidx0<UOp.const(dtypes.int, 1)
+    idx0 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem0, idx.valid(gate)))
+    idx1 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem1, idx.valid(gate)))
+    val = UOp.const(dtypes.float, 42.0)
+    stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
+    uops = to_uops_list(stores)
+    ifs = [u for u in uops if u.op is Ops.IF]
+    endifs = [u for u in uops if u.op is Ops.ENDIF]
+    self.assertEqual(len(ifs), 1)
+    self.assertEqual(len(endifs), 1)
+    gated_uops = tuple(uops[uops.index(ifs[0])+1:uops.index(endifs[0])])
+    self.assertEqual(len(gated_uops), 2)
+    for x in gated_uops: self.assertIs(x.op, Ops.STORE)
+    for x in gated_uops: self.assertEqual(len(x.src), 2)
+
+@unittest.skipIf(Device.DEFAULT == "METAL", "compiler bug")
+@unittest.skipUnless(Ops.SHR in Device[Device.DEFAULT].renderer.code_for_op, "fast_idiv requires SHR")
+class TestFastIdiv(unittest.TestCase):
+  def test_division_power_of_two(self):
+    for dt in (dtypes.int32, dtypes.uint32):
+      g = UOp.param(0, dt.ptr())
+      c = UOp.const(dt, 2)
+      l = g.index(c)
+      a = UOp(Ops.CDIV, dt, (l, c))
+      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
+      Device[Device.DEFAULT].renderer.render(uops)
+      ops = [x.op for x in uops]
+      self.assertIn(Ops.SHR, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
+      self.assertNotIn(Ops.CDIV, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
+
+  def test_floormod_power_of_two(self):
+    # FLOORMOD by a power of two lowers to AND (correct floor mod for any sign in two's complement)
+    for dt in (dtypes.int32, dtypes.uint32):
+      g = UOp.param(0, dt.ptr())
+      c = UOp.const(dt, 8)
+      a = UOp(Ops.FLOORMOD, dt, (g.index(c), c))
+      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
+      ops = [x.op for x in uops]
+      self.assertIn(Ops.AND, ops, f"For dtype={dt} FLOORMOD by pow2 did not simplify to AND")
+      self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORMOD by pow2 left a MOD")
+      self.assertNotIn(Ops.FLOORMOD, ops, f"For dtype={dt} FLOORMOD survived past late rewrite")
+
+  def test_floordiv_power_of_two_uint(self):
+    # uint FLOORDIV by a power of two lowers to a shift, leaving no IDIV/FLOORDIV in the kernel
+    for dt in (dtypes.uint32, dtypes.uint64):
+      g = UOp.param(0, dt.ptr())
+      c = UOp.const(dt, 2)
+      a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
+      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
+      ops = [x.op for x in uops]
+      self.assertIn(Ops.SHR, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
+      self.assertNotIn(Ops.CDIV, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
+      self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")
+
+  @Context(DISABLE_FAST_IDIV=0)
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
+  def test_fast_idiv_and_mod(self):
+    g = UOp.param(0, dtypes.uint32.ptr())
+    c = UOp.const(dtypes.uint, 3)
+    l = g.index(c)
+    a = UOp(Ops.CDIV, dtypes.uint, (l, c))
+    uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
+    Device[Device.DEFAULT].renderer.render(uops)
+    ops = [x.op for x in uops]
+    self.assertIn(Ops.SHR, ops)
+    self.assertNotIn(Ops.CDIV, ops)
+
+    b = UOp(Ops.CMOD, dtypes.uint, (l, c))
+    uops = to_uops_list([b], ren=Device[Device.DEFAULT].renderer)
+    Device[Device.DEFAULT].renderer.render(uops)
+    ops = [x.op for x in uops]
+    self.assertIn(Ops.SHR, ops)
+    self.assertNotIn(Ops.CMOD, ops)
+
+  @Context(DISABLE_FAST_IDIV=0)
+  def test_fast_idiv_bounded_numerator_zero(self):
+    x = UOp.variable("x", 0, 1, dtype=dtypes.int32)
+    for val in range(2):
+      self.assertEqual(eval_uop(x.alu(Ops.CDIV, x.const_like(3)), vals=(val,)), cdiv(val, 3))
+
+  @Context(DISABLE_FAST_IDIV=0)
+  def test_fast_idiv_remove_powers_of_two(self):
+    ridx = UOp.range(2**20, 0)
+    uops = to_uops_list([ridx//(7*64)], ren=Device[Device.DEFAULT].renderer)
+    ops = [x.op for x in uops]
+    # this requires shifting out the powers of two before doing fast_idiv
+    # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
+    self.assertNotIn(Ops.CAST, ops)
+
+  @unittest.expectedFailure
+  def test_fast_idiv_overflow(self):
+    # This will be possible with a slightly different method for fast_idiv
+    g = UOp.param(0, dtypes.uint32.ptr())
+    c = UOp.const(dtypes.uint, 7)
+    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
+    a = UOp(Ops.CDIV, dtypes.uint, (l, c))
+    uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
+    Device[Device.DEFAULT].renderer.render(uops)
+    ops = [x.op for x in uops]
+    self.assertIn(Ops.SHR, ops)
+    self.assertNotIn(Ops.CDIV, ops)
+
+  def test_disable_fast_idiv(self):
+    g = UOp.param(0, dtypes.uint32.ptr())
+    c = UOp.const(dtypes.uint, 3)
+    l = g.index(c)
+    a = UOp(Ops.CDIV, dtypes.uint, (l, c))
+    with Context(DISABLE_FAST_IDIV=1):
+      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
+    ops = [x.op for x in uops]
+    self.assertNotIn(Ops.SHR, ops)
+    self.assertIn(Ops.CDIV, ops)
+
+class TestUOpMethod(unittest.TestCase):
+  @unittest.skip("uops lt no longer ordered")
+  def test_compare_alu_same_src_different_arg(self):
+    a = UOp.const(dtypes.float, 2.0)
+    b = UOp.const(dtypes.float, 3.0)
+
+    add = UOp(Ops.ADD, dtypes.float, (a, b))
+    mul = UOp(Ops.MUL, dtypes.float, (a, b))
+    assert (add < mul) or (mul < add), "add and mul with same src should have an order"
+
+  def test_uop_variables(self):
+    a = UOp.variable("a", 1, 10)
+    uop_var = Tensor(a.bind(1))
+    st_var = Tensor.empty((2, 10))[:, :a.bind(1)]
+    _, var_vals = (uop_var+st_var).linear_with_vars()
+    self.assertEqual(len(var_vals), 1)
+    self.assertEqual(list(var_vals)[0], a.expr)
+
+  def test_const_factor(self):
+    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 8),), 'gidx0')
+    self.assertEqual(UOp.const(dtypes.int, 17).const_factor(), 17)
+    self.assertEqual(gidx0.const_factor(), 1)
+    self.assertEqual((gidx0*3).const_factor(), 3)
+    self.assertEqual((gidx0*3+6).const_factor(), 3)
+    self.assertEqual((gidx0*3+1).const_factor(), 1)
+
+  def test_replace(self):
+    x = UOp.param(0, dtypes.int.ptr())
+    self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int.ptr()).arg).arg.slot, 1)
+    with self.assertRaises(AssertionError): x.replace(field="a")
+
+  def test_const_zero_neg_zero_different(self):
+    # -0.0 and 0.0 must be different UOps (for IEEE754 correctness, e.g. 1/-0.0 = -inf)
+    pos_zero = UOp.const(dtypes.float, 0.0)
+    neg_zero = UOp.const(dtypes.float, -0.0)
+    self.assertIsNot(pos_zero, neg_zero)
+    self.assertNotEqual(hash(pos_zero.arg), hash(neg_zero.arg))
+
+  def test_const_nan_same(self):
+    # nan constants should be deduplicated
+    nan1 = UOp.const(dtypes.float, float('nan'))
+    nan2 = UOp.const(dtypes.float, float('nan'))
+    self.assertIs(nan1, nan2)
+
+class TestUOpStr(unittest.TestCase):
+  def test_uop_str(self):
+    a = UOp.const(dtypes.float, 2.0) + UOp.const(dtypes.float, 3.0)
+    for _ in range(20): a = a + a
+    assert len(str(a)) < 10_000, "exponential string growth"
+    assert str(eval(str(a))) == str(a)
+
+  def test_vectorized_str(self):
+    vec = UOp(Ops.STACK, dtypes.int.vec(4), tuple(UOp.const(dtypes.int, x) for x in range(4)))
+    assert str(eval(str(vec))) == str(vec)
+
+  def test_device_arg(self):
+    device = UOp(Ops.DEVICE, arg="CL")
+    assert str(eval(str(device))) == str(device)
+
+  def test_reduceop_arg(self):
+    sum_uop = Tensor.empty(32, 32).sum().uop
+    assert str(eval(str(sum_uop))) == str(sum_uop)
+
+class TestUPatHelpers(unittest.TestCase):
+  def test_location(self):
+    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "symbolic.py")
+    self.assertEqual(spec_shared.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
+    test_upat = UPat(Ops.CONST, dtypes.bool)
+    self.assertEqual(test_upat.location[0].replace("\\", "/").split("/")[-1], __file__.replace("\\", "/").split("/")[-1])
+    test_upat_named = test_upat.named("test_name")
+    self.assertEqual(test_upat.location[0], test_upat_named.location[0])
+    self.assertNotEqual(test_upat.location[1], test_upat_named.location[1])
+
+class TestUopsObject(unittest.TestCase):
+  def test_timing(self):
+    with Timing("create 10k uops:"): ret = [UOp(Ops.CONST, dtypes.int, arg=10000000+i) for i in range(10000)]
+    assert len(ret) == 10000
+
+  def test_nested(self):
+    a = UOp.new_buffer(Device.DEFAULT, 1, dtypes.char)
+    for _ in range(10_000): a = a+a
+    self.assertEqual(a.device, Device.DEFAULT)
+
+class TestUOpRender(unittest.TestCase):
+  def test_render_vectorize_empty(self):
+    u = UOp(Ops.STACK, dtype=dtypes.int.vec(0), src=())
+    self.assertEqual(u.render(simplify=False), "{}")
+  def test_render_vectorize_empty_simplified(self):
+    u = UOp(Ops.STACK, dtype=dtypes.int.vec(0), src=())
+    self.assertEqual(u.render(), "{}")
+  def test_render_vectorize_same(self):
+    u = UOp(Ops.STACK, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
+    self.assertEqual(u.render(simplify=False), "{0, ...}")
+  def test_render_vectorize_different(self):
+    u = UOp(Ops.STACK, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
+    self.assertEqual(u.render(simplify=False), "{0,1,2}")
+  def test_render_vectorize_same_simplified(self):
+    u = UOp(Ops.STACK, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
+    self.assertEqual(u.render(), "0")
+  def test_render_vectorize_different_simplified(self):
+    u = UOp(Ops.STACK, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
+    self.assertEqual(u.render(), "{0,1,2}")
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad_repo/test/null/test_uops_stats.py

@@ -0,0 +1,239 @@
+import unittest
+from tinygrad import Tensor
+from tinygrad.helpers import GlobalCounters, DEV
+from tinygrad.engine.realize import compile_linear, estimate_uop
+from tinygrad.codegen import to_program
+from tinygrad.renderer import Estimates
+from tinygrad.uop.ops import Ops, UOp
+from tinygrad.dtype import dtypes
+from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
+from tinygrad.device import Device
+from tinygrad.renderer.ptx import PTXRenderer
+from test.helpers import replace_opts
+
+def flops_mem(uops, ignore_indexing=False):
+  est = Estimates.from_uops(uops, ignore_indexing)
+  return est.ops, est.lds
+
+# **************** new FlopCounter ****************
+
+def get_stats(x:Tensor):
+  est = estimate_uop(compile_linear(x.schedule_linear()).src[-1])
+  return est.ops, est.mem
+
+@unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu does extra load/store for packed types")
+class TestMemoryCount(unittest.TestCase):
+  def test_add(self):
+    a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
+    b = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
+    _, mem = get_stats(a+b)
+    self.assertEqual(mem, 1024*1024*3)  # 2 reads + 1 write
+
+  def test_add_const(self):
+    a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
+    _, mem = get_stats(a+3)
+    self.assertEqual(mem, 1024*1024*2)  # 1 read + 1 write
+
+  @unittest.skip("depends on subbuffer working")
+  def test_add_slice(self):
+    a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)[:512]
+    _, mem = get_stats(a+3)
+    self.assertEqual(mem, 512*1024*2)  # 1 read + 1 write
+
+  def test_expanded(self):
+    a = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
+    b = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
+    _, mem = get_stats(a+b)
+    self.assertEqual(mem, 1024*1024*2 + 1024)  # 1 full read + 1 lil read + 1 write
+
+  @unittest.skip("no longer supported")
+  def test_both_expanded(self):
+    # TODO: this probably should be a full write
+    a = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
+    b = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
+    _, mem = get_stats(a+b)
+    # rangeify is smart!
+    self.assertEqual(mem, 1024 + 2*1024)  # 2 lil reads + 1 lil write
+
+  def test_self_add(self):
+    a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
+    _, mem = get_stats(a+a)
+    self.assertEqual(mem, 1024*1024*2)  # 1 read + 1 write
+
+  def test_self_add_transposed(self):
+    a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
+    _, mem = get_stats(a+a.T)
+    self.assertEqual(mem, 1024*1024*2)  # 1 read + 1 write
+
+  def test_self_add_assign(self):
+    a = Tensor.empty(1024, 1024, dtype=dtypes.uint8).realize()
+    _, mem = get_stats(a.assign(a+a))
+    self.assertEqual(mem, 1024*1024*2)  # 1 read + 1 write
+
+  @unittest.skipIf(Device.DEFAULT == "CPU", "test copy to CPU from other device")
+  def test_copyout(self):
+    a = Tensor.empty(32, dtype=dtypes.uint8).to("CPU")
+    _, mem = get_stats(a)
+    self.assertEqual(mem, 32*1)
+    a = Tensor.empty(32, dtype=dtypes.uint32).to("CPU")
+    _, mem = get_stats(a)
+    self.assertEqual(mem, 32*4)
+
+# NOTE: this still isn't testing unroll using the acc
+@unittest.skipUnless(Device.DEFAULT == "PYTHON", "only run test on emulated tensor cores")
+class TestUOpsStatsMatmulHalf(unittest.TestCase):
+  def test_simple_matmul_half(self, N=16):
+    GlobalCounters.reset()
+    a, b = Tensor.empty(N, N, dtype=dtypes.half), Tensor.empty(N, N, dtype=dtypes.half)
+    c = a.matmul(b)
+    c.realize()
+    expected_ops = N ** 3 * 2
+    self.assertEqual(expected_ops, GlobalCounters.global_ops)
+
+  @unittest.skipIf(DEV.arch=="INTEL", "intel gets 524288 != 524352")
+  def test_bigger_matmul_half(self): self.test_simple_matmul_half(64)
+
+  def test_batched_matmul_half(self, N=16):
+    GlobalCounters.reset()
+    a, b = Tensor.empty(4, N, N, dtype=dtypes.half), Tensor.empty(1, N, N, dtype=dtypes.half)
+    c = a.matmul(b)
+    c.realize()
+    expected_ops = 4 * N ** 3 * 2
+    self.assertEqual(expected_ops, GlobalCounters.global_ops)
+
+class TestUOpsStats(unittest.TestCase):
+  def test_simple_add(self):
+    a = Tensor.empty(100,100)
+    b = Tensor.empty(100,100)
+    c = a+b
+    ops, mem = get_stats(c)
+    expected_ops = c.numel()
+    expected_mem = a.nbytes() + b.nbytes() + c.nbytes()
+    self.assertEqual(mem, expected_mem)
+    # NOTE; ops also include indexing ops
+    assert expected_ops <= ops and ops <= expected_ops * 2
+
+  def test_simple_add_sq(self):
+    a = Tensor.empty(100,100)
+    b = Tensor.empty(100,100)
+    c = (a+b)*(a+b)
+    ops, mem = get_stats(c)
+    expected_ops = c.numel()*2
+    expected_mem = a.nbytes() + b.nbytes() + c.nbytes()
+    self.assertEqual(mem, expected_mem)
+    # NOTE; ops also include indexing ops
+    assert expected_ops <= ops and ops <= expected_ops * 2
+
+  def test_simple_matmul(self, M=1024, N=1024, K=1024):
+    a = Tensor.empty(M,N)
+    b = Tensor.empty(N,K)
+    c = a@b
+    ops, mem = get_stats(c)
+    expected_ops = c.numel() * N * 2
+    required_mem = a.nbytes() + b.nbytes() + c.nbytes()
+    assert expected_ops <= ops and ops <= expected_ops * 1.2
+    # NOTE: it's hard to assert on the memory here, all depends on caching
+    assert required_mem <= mem
+
+  def test_simple_matmul_8192(self): self.test_simple_matmul(8192, 8192, 8192)
+
+  #MULACC should have the same stats as MUL + ADD
+  def test_mulacc(self):
+    globl = UOp.param(0, dtypes.int.ptr())
+    o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
+    o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
+    u1 = globl.index(o1)
+    u2 = globl.index(o2)
+    u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
+    u4 = UOp(Ops.MUL, dtypes.int, (u1,u2))
+    u5 = UOp(Ops.ADD, dtypes.int, (u4,u3))
+    uops = tuple(u5.toposort())
+
+    globl = UOp.param(0, dtypes.int.ptr())
+    o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
+    o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
+    u1 = globl.index(o1)
+    u2 = globl.index(o2)
+    u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
+    u4 = UOp(Ops.MULACC, dtypes.int, (u1,u2,u3))
+    uops_fma = tuple(u4.toposort())
+
+    self.assertEqual(flops_mem(uops), flops_mem(uops_fma))
+
+N = 64
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "wrong in PTX") # maybe?
+class TestStatsOptimized(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.ast_gemm = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
+    cls.ast_reduce = (Tensor.empty(N*N).sum()).schedule_linear().src[-1].src[0]
+
+  def check_gemm(self, p:UOp, extra_flops=0):
+    est = p.src[0].arg.estimates
+    print(p.arg.name, est.ops, est.mem, est.lds)
+    self.assertEqual(est.ops, 2*N*N*N + extra_flops)  # N**3 mulaccs
+    self.assertEqual(est.mem, 3*N*N*4) # 3 NxN mats with floats
+
+  def test_gemm(self):
+    p = to_program(replace_opts(self.ast_gemm, []), renderer=Device[Device.DEFAULT].renderer)
+    self.check_gemm(p)
+    self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4 + 4*N*N)
+
+  def test_gemm_tc_unroll(self):
+    try:
+      p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),
+                      renderer=Device[Device.DEFAULT].renderer)
+    except KernelOptError:
+      raise unittest.SkipTest("no tensor cores")
+    print(p.src[3].arg)
+    self.check_gemm(p)
+
+  # this is a good lesson about why UPCASTing is a good idea
+
+  def test_gemm_one_upcasted(self):
+    p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4)]), renderer=Device[Device.DEFAULT].renderer)
+    self.check_gemm(p)
+    self.assertEqual(p.src[0].arg.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
+
+  def test_gemm_upcasted(self):
+    p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)]),
+                    renderer=Device[Device.DEFAULT].renderer)
+    self.check_gemm(p)
+    self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
+
+  def test_gemm_upcasted_locals(self):
+    try:
+      p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 0, 4),
+                                                   Opt(OptOps.LOCAL, 1, 4)]), renderer=Device[Device.DEFAULT].renderer)
+    except KernelOptError:
+      raise unittest.SkipTest("no locals")
+    self.check_gemm(p)
+    self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
+
+  def test_gemm_group(self):
+    try:
+      p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.GROUP, 0, 4)]), renderer=Device[Device.DEFAULT].renderer)
+    except KernelOptError:
+      raise unittest.SkipTest("no locals")
+    SZ = N*N*4
+    # NOTE: these are sort of wrong. they aren't honoring the IF statement
+    self.check_gemm(p, extra_flops=SZ*4)
+    self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)
+
+  def test_reduce(self):
+    p = to_program(replace_opts(self.ast_reduce, []), renderer=Device[Device.DEFAULT].renderer)
+    est = p.src[0].arg.estimates
+    print(p.arg.name, est.ops, est.mem, est.lds)
+    self.assertEqual(est.ops, N*N)
+    self.assertEqual(est.mem, N*N*4 + 4)
+
+  def test_reduce_group(self):
+    try:
+      p = to_program(replace_opts(self.ast_reduce, [Opt(OptOps.GROUP, 0, 50)]), renderer=Device[Device.DEFAULT].renderer)
+    except KernelOptError:
+      raise unittest.SkipTest("no locals")
+    est = p.src[0].arg.estimates
+    print(p.arg.name, est.ops, est.mem, est.lds)
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)

tinygrad_repo/test/null/test_upat_compile.py

@@ -0,0 +1,54 @@
+import unittest
+from tinygrad.helpers import DEBUG, Context
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import UPat, track_rewrites, GroupOp, Ops
+from tinygrad.uop.upat import _get_code, upat_compile
+import dis
+
+@track_rewrites()
+def do_compile(up):
+  print("\n***** COMPILE", up)
+  match_code = _get_code(up, False)
+  match = upat_compile(up, lambda **kwargs: None)
+  print(match_code[0])
+  if DEBUG >= 2: dis.dis(match)
+  return match_code[0]
+
+@Context(SPEC=0)
+class TestUPatCompile(unittest.TestCase):
+  def test_double(self):
+    up = UPat.var("x") * UPat.cvar("c0") + UPat.var("x") * UPat.cvar("c1")
+    do_compile(up)
+
+  def test_single(self):
+    up = UPat.var("x") + UPat.var("y")
+    do_compile(up)
+
+  def test_xpx(self):
+    up = UPat.var("x") + UPat.var("x")
+    do_compile(up)
+
+  def test_xp0(self):
+    up = UPat.var("x") + 0
+    do_compile(up)
+
+  def test_bool(self):
+    up = UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool)
+    do_compile(up)
+
+  def test_single_c(self):
+    up = (UPat.var("x") + UPat.var("y")) * UPat.var("c")
+    do_compile(up)
+
+  def test_const_folding(self):
+    up = UPat(GroupOp.ALU-{Ops.THREEFRY}, name="a", src=UPat((Ops.CONST, Ops.STACK)))
+    do_compile(up)
+
+  @unittest.skip("fix this")
+  def test_range_named(self):
+    # this should be one src, but this should also still work
+    up = UPat(Ops.CAST, dtypes.float, UPat.var("x", dtypes.bfloat16))
+    do_compile(up)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/null/test_validate_oob.py

@@ -0,0 +1,179 @@
+import unittest
+from tinygrad import dtypes, Variable
+from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import Context
+from tinygrad.uop.ops import Ops, UOp, AxisType
+from test.helpers import to_uops_list
+
+class TestValidateOOB(unittest.TestCase):
+  """Test z3 validation of index bounds for different ALU ops and patterns."""
+
+  # basic index patterns
+  def test_const_index(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(16))
+      to_uops_list([buf.index(UOp.const(dtypes.int, 0), ptr=True).load(dtype=dtypes.int)])  # valid
+      to_uops_list([buf.index(UOp.const(dtypes.int, 15), ptr=True).load(dtype=dtypes.int)])  # valid (last element)
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(UOp.const(dtypes.int, 16), ptr=True).load(dtype=dtypes.int)])  # off by one
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(UOp.const(dtypes.int, 42), ptr=True).load(dtype=dtypes.int)])  # way out
+
+  def test_variable_index(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(16))
+      to_uops_list([buf.index(Variable("i", 0, 15), ptr=True).load(dtype=dtypes.int)])  # valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(Variable("i", 0, 20), ptr=True).load(dtype=dtypes.int)])  # oob
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(Variable("i", -5, 10), ptr=True).load(dtype=dtypes.int)])  # negative
+
+  def test_range_with_mask(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(16))
+      r = UOp.range(42, 0, AxisType.GLOBAL)
+      to_uops_list([buf.index(r.valid(r < 16), ptr=True).load(dtype=dtypes.int)])  # valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r.valid(r < 17), ptr=True).load(dtype=dtypes.int)])  # oob
+
+  def test_variable_with_mask(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(16))
+      v = Variable("v", -5, 80)
+      to_uops_list([buf.index(v.valid((v >= 0) & (v < 16)), ptr=True).load(dtype=dtypes.int)])  # valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(v.valid(v < 20), ptr=True).load(dtype=dtypes.int)])  # negative not masked
+
+  def test_gated_store(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(16))
+      v = Variable("v", 0, 20)
+      to_uops_list([buf.index(v.valid(v < 16), ptr=True).store(0)])  # valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(v.valid(v < 20), ptr=True).store(0)])  # oob
+
+  # ALU ops in index
+  def test_floordiv(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(16))
+      to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)])  # 0..16 oob
+
+  def test_mod(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(16))
+      r = UOp.range(100, 0, AxisType.GLOBAL)
+      to_uops_list([buf.index(r % 16, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r % 20, ptr=True).load(dtype=dtypes.int)])  # 0..19 oob
+
+  def test_shr(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(16))
+      to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)])  # 0..31 oob
+
+  def test_shl(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(64))
+      r = UOp.range(8, 0, AxisType.GLOBAL)
+      to_uops_list([buf.index(r << 2, ptr=True).load(dtype=dtypes.int)])  # 0..28 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r << 4, ptr=True).load(dtype=dtypes.int)])  # 0..112 oob
+
+  def test_and(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(16))
+      r = UOp.range(100, 0, AxisType.GLOBAL)
+      to_uops_list([buf.index(r & 15, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r & 31, ptr=True).load(dtype=dtypes.int)])  # 0..31 oob
+
+  def test_max(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(16))
+      to_uops_list([buf.index(Variable("v", -10, 15).maximum(0), ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0), ptr=True).load(dtype=dtypes.int)])  # 0..20 oob
+
+  def test_xor_in_mask(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(16))
+      r = UOp.range(32, 0, AxisType.GLOBAL)
+      to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16))), ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r.valid((r < 10) ^ (r >= 20)), ptr=True).load(dtype=dtypes.int)])  # 0..9,20..31 oob
+
+  # cast patterns
+  def test_float_cast_in_index(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(16))
+      r = UOp.range(20, 0)
+      i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
+      to_uops_list([buf.index(i.valid((i >= 0) & (i < 16)), ptr=True).load(dtype=dtypes.int)])
+
+  def test_bool_cast_in_mask(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int.ptr(1))
+      r = UOp.range(20, 0)
+      to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not()), ptr=True).load(dtype=dtypes.int)])  # only r=0 valid
+
+  # load result as index/mask
+  def test_load_as_index(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf0 = UOp.param(0, dtypes.int.ptr(16))
+      buf1 = UOp.param(1, dtypes.int.ptr(64))
+      r = UOp.range(42, 0, AxisType.GLOBAL)
+      ld0 = buf0.index(r.valid(r < 8), ptr=True).load(dtype=dtypes.int).cast(dtypes.weakint)
+      to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32)), ptr=True).load(dtype=dtypes.int)])  # valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64)), ptr=True).load(dtype=dtypes.int)])  # oob
+
+  def test_load_bool_as_mask(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf_bool = UOp.param(0, dtypes.bool.ptr(16))
+      buf_int = UOp.param(1, dtypes.int.ptr(8))
+      gidx = UOp(Ops.SPECIAL, dtypes.weakint, (UOp.const(dtypes.weakint, 16),), "gidx0")
+      ld_bool = buf_bool.index(gidx, ptr=True).load()
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf_int.index(gidx.valid(ld_bool), ptr=True).load()])  # gidx 0..15, buf_int size 8
+
+  # skipped tests (moved from test_uop_graph.py)
+  @unittest.skip("if not allowed in graph")
+  def test_in_bounds_access_gated_local(self):
+    with Context(CHECK_OOB=1):
+      # Define buffers
+      gbuf = UOp.param(0, dtypes.uint.ptr(400))
+      sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.uint.ptr(8, addrspace=AddrSpace.LOCAL), (), "temp0")
+
+      # Define indices, valids and barrier
+      gidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 416),), "gidx0")
+      lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "lidx0")
+
+      gate = (gidx<400) & (lidx<8)
+
+      local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(dtypes.uint, 1))
+
+      barrier = UOp(Ops.BARRIER, dtypes.void, (local_store,))
+      if_barrier = UOp(Ops.IF, dtypes.void, (gate, barrier))
+
+      # Load from local memory (after the IF/barrier)
+      local_load = UOp(Ops.LOAD, dtypes.uint, (sbuf.index(lidx, ptr=True), if_barrier))
+
+      # Store to global memory
+      global_store = UOp(Ops.STORE, dtypes.void, (gbuf.index(gidx), local_load))
+      to_uops_list([global_store])
+
+  @unittest.skip("Bool load is not supported yet")
+  def test_load_mask(self):
+    with Context(CHECK_OOB=1):
+      glbl0 = UOp.param(0, dtypes.int.ptr(16))
+      mask = UOp.param(0, dtypes.bool.ptr(16))
+      ridx = UOp.range(20, 0)
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask), ptr=True)))
+      to_uops_list([ld0])
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/null/test_winograd.py

@@ -0,0 +1,55 @@
+import unittest, sys
+from tinygrad import Tensor, GlobalCounters, dtypes, Context
+from tinygrad.helpers import WINO
+
+@unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
+class TestWinograd(unittest.TestCase):
+  def setUp(self):
+    self.old = WINO.value
+    WINO.value = 1
+  def tearDown(self):
+    WINO.value = self.old
+
+  def test_forward_kernels(self):
+    x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
+    out = Tensor.conv2d(x,w)
+    self.assertEqual(len(out.schedule_linear().src), 2)
+
+  def test_backward_kernels(self):
+    x,w = Tensor.empty(1,4,9,9).realize(), Tensor.empty(4,4,3,3).realize()
+    out = Tensor.conv2d(x,w, padding=1)
+    out.mean().backward()
+    backward_schedule = x.grad.schedule_linear(w.grad)
+    self.assertEqual(len(backward_schedule.src), 4)
+
+  @unittest.skip("this requires optimizations")
+  def test_counters(self):
+    IC, OC, X, Y = 4,4,9,9
+    x,w = Tensor.rand(1,IC,Y,X).realize(), Tensor.rand(OC,IC,3,3).realize()
+    GlobalCounters.reset()
+    with Context(WINO=1):
+      Tensor.conv2d(x,w).realize()
+    ops_wino, mem_wino = GlobalCounters.global_ops, GlobalCounters.global_mem
+    GlobalCounters.reset()
+    with Context(WINO=0):
+      Tensor.conv2d(x,w).realize()
+    ops_normal, mem_normal = GlobalCounters.global_ops, GlobalCounters.global_mem
+
+    ops_ratio, mem_ratio = ops_wino/ops_normal, mem_wino/mem_normal
+    print(f"ops: normal {ops_normal:9d} wino {ops_wino:9d} ratio {ops_ratio:.2f}")
+    print(f"mem: normal {mem_normal:9d} wino {mem_wino:9d} ratio {mem_ratio:.2f}")
+
+    # TODO: what's optimal on this?
+    self.assertLess(ops_ratio, 4.3)
+    self.assertLess(mem_ratio, 4)
+
+  def test_dtype(self):
+    IC, OC, X, Y = 4,4,9,9
+    x,w = Tensor.empty(1,IC,Y,X), Tensor.empty(OC,IC,3,3)
+    self.assertEqual(Tensor.conv2d(x,w).dtype, dtypes.default_float)
+
+    x,w = Tensor.empty(1,IC,Y,X,dtype=dtypes.half), Tensor.empty(OC,IC,3,3,dtype=dtypes.half)
+    self.assertEqual(Tensor.conv2d(x,w).dtype, dtypes.half)
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)