Merge branch 'master' into kernel_is_call

fixed wrong comments and simplified queue building

.github/workflows/benchmark.yml

@@ -41,6 +41,18 @@ jobs:
       run: |
         echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
         rm -f /tmp/pytest-db-ci*
+    # TODO: remove this step once all old caches are migrated
+    - name: Migrate old huggingface cache (symlinks break onnxruntime 1.24+)
+      run: |
+        cd ~/Library/Caches/tinygrad/downloads/models 2>/dev/null || exit 0
+        for old_dir in models--*; do
+          [ -d "$old_dir" ] || continue
+          repo_id=$(echo "$old_dir" | sed 's/models--//; s/--/\//g')
+          snapshot=$(ls -1 "$old_dir/snapshots" 2>/dev/null | head -1)
+          [ -n "$snapshot" ] || continue
+          mkdir -p "$repo_id"
+          cp -RLn "$old_dir/snapshots/$snapshot/"* "$repo_id/" 2>/dev/null || true
+        done
     - name: Run pytest -nauto
       run: |
         source /tmp/tinygrad_pytest_ci/bin/activate

.github/workflows/test.yml

@@ -29,9 +29,9 @@ jobs:
         deps: testing_unit
         llvm: 'true'
     - name: Speed Test
-      run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
     - name: Speed Test (BEAM=2)
-      run: BEAM=2 CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: BEAM=2 CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
 
   docs:
     name: Docs
@@ -207,11 +207,11 @@ jobs:
       run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - name: Test device flop counts
       run: |
-        DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul
+        DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStats.test_simple_matmul
 
   linter:
     name: Linters
@@ -569,7 +569,7 @@ jobs:
       with:
         key: dsp-minimal
         deps: testing_unit
-        pydeps: "onnx==1.18.0 onnxruntime"
+        pydeps: "onnx==1.18.0 onnxruntime ml_dtypes"
         llvm: "true"
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v3

examples/mlperf/dataloader.py

@@ -72,7 +72,7 @@ def loader_process(q_in, q_out, X:Tensor, seed):
       #storage_tensor._copyin(img_tensor.numpy())
 
       # faster
-      X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
+      X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()
 
       # ideal
       #X[idx].assign(img.tobytes())   # NOTE: this is slow!
@@ -264,8 +264,8 @@ def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tens
       x = random_brightness_augmentation(x)
       x = gaussian_noise(x)
 
-    X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = x.tobytes()
-    Y[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = y.tobytes()
+    X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = x.tobytes()
+    Y[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = y.tobytes()
 
     queue_out.put(idx)
   queue_out.put(None)
@@ -379,12 +379,12 @@ def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue
       clipped_match_idxs = np.clip(match_idxs, 0, None)
       clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]
 
-      boxes[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_boxes.tobytes()
-      labels[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_labels.tobytes()
-      matches[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = match_idxs.tobytes()
-      anchors[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = anchor.tobytes()
+      boxes[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_boxes.tobytes()
+      labels[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_labels.tobytes()
+      matches[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = match_idxs.tobytes()
+      anchors[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = anchor.tobytes()
 
-    imgs[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
+    imgs[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()
 
     queue_out.put(idx)
   queue_out.put(None)

examples/stable_diffusion.py

@@ -6,7 +6,6 @@ import argparse, time
 from collections import namedtuple
 from typing import Dict, Any
 
-from PIL import Image
 import numpy as np
 from tinygrad import Device, GlobalCounters, dtypes, Tensor, TinyJit
 from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten, profile_marker
@@ -336,6 +335,7 @@ if __name__ == "__main__":
   print(x.shape)
 
   profile_marker("save image")
+  from PIL import Image
   im = Image.fromarray(x.numpy())
   print(f"saving {args.out}")
   im.save(args.out)

examples/tinychat/tinychat-browser/compile.py

@@ -48,7 +48,7 @@ def prepare_browser_chunks(model):
         weight_metadata = metadata.get(name, default)
         weight_metadata["parts"][part_num] = {"file": i, "file_start_pos": cursor, "size": size}
         metadata[name] = weight_metadata
-        data = bytes(state_dict[name].uop.base.realized.as_buffer())
+        data = bytes(state_dict[name].uop.base.realized.as_memoryview())
         data = data if not offsets else data[offsets[0]:offsets[1]]
         writer.write(data)
         cursor += size

extra/amdpci/proclogs.py

@@ -19,8 +19,9 @@ amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
 amdev.AMDev = AMDFake
 from tinygrad.runtime.ops_amd import PCIIface
 
-def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = False):
+def parse_amdgpu_logs(log_content, register_names=None, register_objects=None, *, only_xcc0: bool = False):
   register_map = register_names or {}
+  register_objs = register_objects or {}
 
   def replace_register(match):
     reg = match.group(1)
@@ -37,6 +38,28 @@ def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = Fal
   # remove timing prefix
   processed_log = re.sub(r'^\[\s*\d+(?:\.\d+)?\]\s*', '', processed_log, flags=re.MULTILINE)
 
+  # decode register values into field dicts
+  def decode_value(match):
+    reg_name = match.group(1)
+    xcc_part = match.group(2)  # "xcc=0 " or ""
+    val_str = match.group(3)
+    val = int(val_str, 16)
+
+    reg_obj = register_objs.get(reg_name)
+    if reg_obj is not None and reg_obj.fields:
+      fields = reg_obj.decode(val)
+      # show raw for unaccounted bits
+      accounted = 0
+      for name, (start, end) in reg_obj.fields.items():
+        accounted |= (((1 << (end - start + 1)) - 1) << start)
+      unaccounted = val & ~accounted
+      parts = {k: v for k, v in fields.items() if v != 0}
+      if unaccounted: parts['_raw_unaccounted'] = hex(unaccounted)
+      return f"register {reg_name}, {xcc_part}with value {val_str} {parts}"
+    return match.group(0)
+
+  processed_log = re.sub(r'register (reg\w+), ((?:xcc=\d+ )?)with value (0x[0-9a-fA-F]+)', decode_value, processed_log)
+
   # keep only xcc=0 lines (but keep lines with no xcc at all)
   if only_xcc0:
     kept = []
@@ -50,16 +73,18 @@ def main():
   only_xcc0 = bool(getenv("ONLY_XCC0", 0))
 
   reg_names = {}
+  reg_objs = {}
   dev = PCIIface(None, 0)
   for x, y in dev.dev_impl.__dict__.items():
     if isinstance(y, AMRegister):
       for xcc, addr in y.addr.items():
         reg_names[addr] = f"{x}, xcc={xcc}"
+      reg_objs[x] = y
 
   with open(sys.argv[1], 'r') as f:
     log_content = f.read()
 
-  processed_log = parse_amdgpu_logs(log_content, reg_names, only_xcc0=only_xcc0)
+  processed_log = parse_amdgpu_logs(log_content, reg_names, reg_objs, only_xcc0=only_xcc0)
 
   with open(sys.argv[2], 'w') as f:
     f.write(processed_log)

extra/assembly/amd/emu.py

@@ -135,14 +135,6 @@ def _val_to_u32(val: UOp) -> UOp:
   if val.dtype in (dtypes.uint16, dtypes.int16): return val.cast(dtypes.uint32)
   return val.cast(dtypes.uint32)
 
-def _apply_clamp(val: UOp, clmp: int | UOp) -> UOp:
-  """Apply VOP3 clamp modifier: clamp float results to [0.0, 1.0] range."""
-  if isinstance(clmp, int) and clmp == 0: return val
-  if val.dtype not in (dtypes.float32, dtypes.half, dtypes.float64): return val
-  zero, one = UOp.const(val.dtype, 0.0), UOp.const(val.dtype, 1.0)
-  clamped = val.maximum(zero).minimum(one)
-  return clmp.ne(_c(0)).where(clamped, val) if isinstance(clmp, UOp) else clamped
-
 _pcode_fixes = {
   'V_DIV_FMAS_F32': ('D0.f32 = 2.0F ** 32 * fma(S0.f32, S1.f32, S2.f32)',
     'D0.f32 = (exponent(S2.f32) > 127) ? (2.0F ** 64 * fma(S0.f32, S1.f32, S2.f32)) : (2.0F ** -64 * fma(S0.f32, S1.f32, S2.f32))'),
@@ -354,7 +346,7 @@ class _Ctx:
     offset = reg.cast(dtypes.int) * _c(32, dtypes.int) + lane.cast(dtypes.int)
     return buf.index(offset, _lane_active(exec_mask, lane)).store(val.cast(dtypes.uint32))
 
-  def rsrc_dyn(self, off: UOp, lane: UOp | None, bits: int = 32, literal: UOp | None = None, is_f64: bool = False) -> UOp:
+  def rsrc_dyn(self, off: UOp, lane: UOp | None, bits: int = 32, literal: UOp | None = None, is_f64: bool = False, do_cast: bool = True) -> UOp:
     """Read source operand with dynamic offset. Handles SGPR/inline constants (<256), VGPR (>=256).
     If lane is None, only scalar access is supported (off must be < 256).
     is_f64: True for F64 operations where 64-bit literals go in high 32 bits."""
@@ -385,7 +377,7 @@ class _Ctx:
     else:
       scalar_val = sgpr_lo
       if literal is not None: scalar_val = off.eq(_c(255)).where(literal, scalar_val)
-      if bits == 16:  # Float constants: cast F32 to F16
+      if bits == 16 and do_cast:  # Float constants: cast F32 to F16
         scalar_val = is_float_const.where(scalar_val.bitcast(dtypes.float32).cast(dtypes.half).bitcast(dtypes.uint16).cast(dtypes.uint32), scalar_val)
 
     return is_vgpr.where(vgpr_val, scalar_val) if lane is not None else scalar_val
@@ -445,7 +437,7 @@ class _Ctx:
     return UOp.sink(*stores, *self.inc_pc())
 
   def compile_vop_pcode(self, op, srcs: dict[str, UOp], lane: UOp, vdst_reg: UOp, exec_mask: UOp,
-                        opsel_dst_hi: bool | UOp = False, sdst_reg: int | None = None, clmp: int | UOp = 0) -> UOp:
+                        opsel_dst_hi: bool | UOp = False, sdst_reg: int | None = None, clmp: int = 0) -> UOp:
     """Compile VOP instruction. Returns sink with stores and inc_pc."""
     pcode = get_pcode(op)
     vcc_reg = sdst_reg if sdst_reg is not None else VCC_LO.offset
@@ -454,6 +446,24 @@ class _Ctx:
                  'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0)})  # rounding mode: 0=RNE, RTZ constant
     _, assigns = parse_pcode(pcode, srcs)
 
+    # For integer ops with clamp, compute overflow using wide arithmetic
+    # NOTE: MUL_LO ops don't saturate - they always return the low bits
+    int_saturate = None
+    if clmp and any(p in op.name for p in ('_NC_U', '_MAD_U', '_NC_I', '_MAD_I')):
+      is_signed, is_16bit = '_I' in op.name and '_U' not in op.name, '16' in op.name
+      if not (is_16bit and is_signed):  # Skip 16-bit signed ops due to codegen issues
+        s0, s1, s2 = srcs.get('S0'), srcs.get('S1'), srcs.get('S2')
+        if s0 is not None and s1 is not None:
+          narrow_dt = dtypes.uint16 if is_16bit else (dtypes.int32 if is_signed else dtypes.uint32)
+          wide_dt = dtypes.int32 if is_16bit else dtypes.int64
+          narrow_max, narrow_min = (0xFFFF, 0) if is_16bit else ((0x7FFFFFFF, -0x80000000) if is_signed else (0xFFFFFFFF, 0))
+          def to_wide(x): return (x.bitcast(narrow_dt) if x.dtype.itemsize == narrow_dt.itemsize else x.cast(narrow_dt)).cast(wide_dt)
+          is_sub, is_mad = 'SUB' in op.name, 'MAD' in op.name
+          full = (to_wide(s0) * to_wide(s1) + to_wide(s2)) if is_mad and s2 is not None else \
+                 (to_wide(s1) - to_wide(s0)) if is_sub and 'SUBREV' in op.name else \
+                 (to_wide(s0) - to_wide(s1)) if is_sub else (to_wide(s0) + to_wide(s1))
+          int_saturate = full.clamp(narrow_min, narrow_max).cast(narrow_dt)
+
     raw_stores: list = []
     vcc_val, exec_val = None, None
     for dest, val in assigns:
@@ -468,7 +478,10 @@ class _Ctx:
                        val.cast(dtypes.uint32) if val.dtype in (dtypes.uint16, dtypes.int16) else val.cast(dtypes.uint32) & UOp.const(dtypes.uint32, slice_mask)
             raw_stores.append(('vgpr_slice', (lo_bit, width, val_bits)))
             continue
-        val = _apply_clamp(val, clmp)
+        # For integer ops with clamp, use pre-computed saturated value; for floats, clamp to [0,1]
+        if int_saturate is not None: val = int_saturate
+        elif clmp and val.dtype in (dtypes.float32, dtypes.half, dtypes.float64):
+          val = val.maximum(UOp.const(val.dtype, 0.0)).minimum(UOp.const(val.dtype, 1.0))
         if val.dtype in (dtypes.uint64, dtypes.int64, dtypes.float64):
           lo, hi = _split64(val)
           raw_stores.extend([('vgpr', self.wvgpr_dyn(vdst_reg, lane, lo, exec_mask)), ('vgpr', self.wvgpr_dyn(vdst_reg + _c(1), lane, hi, exec_mask))])
@@ -732,6 +745,7 @@ def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD, ctx: _Ctx) -> UOp:
   _, assigns = parse_pcode(pcode, srcs)
 
   has_per_lane_vcc = any('[laneId]' in dest for dest, _ in assigns if dest.startswith('VCC') or dest.startswith('D0.u64'))
+  clmp = getattr(inst, 'clmp', 0)
   if has_per_lane_vcc:
     # VCC computation: RANGE+REDUCE gets axis ID first (lower ID = runs first)
     # This ensures VCC reads source values BEFORE VGPR stores modify them
@@ -743,11 +757,17 @@ def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD, ctx: _Ctx) -> UOp:
     final_vcc = ctx.unroll_lanes(get_vcc_bit, exec_mask)
     # VGPR stores: RANGE gets axis ID second (higher ID = runs after VCC loop)
     lane3 = ctx.range()
-    d0_val = None
+    d0_val, vcc_per_lane = None, None
     for dest, val in parse_pcode(pcode, load_srcs(lane3))[1]:
       if dest.startswith('D0') and '[laneId]' not in dest: d0_val = val
+      if dest.startswith('VCC') or (dest.startswith('D0.u64') and '[laneId]' in dest): vcc_per_lane = val
     vgpr_stores = []
     if d0_val is not None:
+      # Apply clamp using carry/borrow bit: ADD overflow->0xFFFFFFFF, SUB underflow->0
+      if clmp and vcc_per_lane is not None:
+        is_sub = 'SUB' in inst.op.name
+        sat_val = _c(0) if is_sub else _c(0xFFFFFFFF)
+        d0_val = vcc_per_lane.cast(dtypes.bool).where(sat_val, d0_val.cast(dtypes.uint32))
       if d0_val.dtype in (dtypes.uint64, dtypes.int64, dtypes.float64):
         lo, hi = _split64(d0_val)
         vgpr_stores.extend([ctx.wvgpr_dyn(vdst_reg, lane3, lo, exec_mask), ctx.wvgpr_dyn(vdst_reg + _c(1), lane3, hi, exec_mask)])
@@ -800,9 +820,10 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P, ctx: _Ctx) -> UOp:
   lane = ctx.range()
   exec_mask = ctx.rsgpr_dyn(_c(EXEC_LO.offset))
   vdst_reg = ctx.inst_field(type(inst).vdst)
-  src0 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src0), lane, 16)
-  src1 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src1), lane, 16)
-  src2 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src2), lane, 16)
+  do_cast = any(x in op_name for x in ('F16', 'F32', 'BF16')) and 'IU' not in op_name
+  src0 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src0), lane, 16, do_cast=do_cast)
+  src1 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src1), lane, 16, do_cast=do_cast)
+  src2 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src2), lane, 16, do_cast=do_cast)
   opsel, opsel_hi = getattr(inst, 'opsel', 0) or 0, getattr(inst, 'opsel_hi', 3) if getattr(inst, 'opsel_hi', 3) is not None else 3
   opsel_hi2 = getattr(inst, 'opsel_hi2', 1) if getattr(inst, 'opsel_hi2', 1) is not None else 1
   neg, neg_hi = getattr(inst, 'neg', 0) or 0, getattr(inst, 'neg_hi', 0) or 0
@@ -1120,8 +1141,8 @@ class WaveState:
     self.n_lanes = n_lanes
     self.vgpr_buf = Buffer('CPU', VGPR_SIZE, dtypes.uint32).ensure_allocated()
     self.sgpr_buf = Buffer('CPU', SGPR_COUNT, dtypes.uint32).ensure_allocated()
-    self._vgpr_mv = self.vgpr_buf.as_buffer(force_zero_copy=True).cast('I')
-    self._sgpr_mv = self.sgpr_buf.as_buffer(force_zero_copy=True).cast('I')
+    self._vgpr_mv = self.vgpr_buf.as_memoryview(force_zero_copy=True).cast('I')
+    self._sgpr_mv = self.sgpr_buf.as_memoryview(force_zero_copy=True).cast('I')
     # Zero memory using ctypes memset (much faster than Python loops)
     ctypes.memset(self.vgpr_buf._buf.va_addr, 0, VGPR_SIZE * 4)
     ctypes.memset(self.sgpr_buf._buf.va_addr, 0, SGPR_COUNT * 4)

extra/assembly/amd/pcode.py

@@ -440,7 +440,7 @@ class Parser:
       self.eat('COMMA')
       lo = self.parse()
       self.eat('RBRACE')
-      return (hi.cast(dtypes.uint64) << _u64(32)) | lo.cast(dtypes.uint64)
+      return (hi.cast(dt:=_BITS_DT.get((s:=lo.dtype.bitsize) * 2, dtypes.uint64)) << _const(dt, s)) | lo.cast(dt)
     if self.at('NUM'):
       num = self.eat('NUM').val
       if self.try_eat('QUOTE'):

extra/assembly/amd/test/hw/test_sop.py

@@ -87,6 +87,50 @@ class TestBasicScalar(unittest.TestCase):
     self.assertEqual(st.sgpr[2], f2i(2.0))
 
 
+class TestPack(unittest.TestCase):
+  """Tests for S_PACK instructions."""
+
+  def test_s_pack_ll_b32_b16(self):
+    """S_PACK_LL_B32_B16 packs low 16 bits of two sources into one 32-bit result."""
+    instructions = [
+      s_mov_b32(s[0], 0xDEADAAAA),
+      s_mov_b32(s[1], 0xDEADBBBB),
+      s_pack_ll_b32_b16(s[2], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[2], 0xBBBBAAAA)
+
+  def test_s_pack_lh_b32_b16(self):
+    """S_PACK_LH_B32_B16: D0 = { S1[31:16], S0[15:0] }."""
+    instructions = [
+      s_mov_b32(s[0], 0xDEADAAAA),
+      s_mov_b32(s[1], 0xDEADBBBB),
+      s_pack_lh_b32_b16(s[2], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[2], 0xDEADAAAA)
+
+  def test_s_pack_hh_b32_b16(self):
+    """S_PACK_HH_B32_B16: D0 = { S1[31:16], S0[31:16] }."""
+    instructions = [
+      s_mov_b32(s[0], 0xDEADAAAA),
+      s_mov_b32(s[1], 0xDEADBBBB),
+      s_pack_hh_b32_b16(s[2], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[2], 0xDEADDEAD)
+
+  def test_s_pack_hl_b32_b16(self):
+    """S_PACK_HL_B32_B16: D0 = { S1[15:0], S0[31:16] }."""
+    instructions = [
+      s_mov_b32(s[0], 0xDEADAAAA),
+      s_mov_b32(s[1], 0xDEADBBBB),
+      s_pack_hl_b32_b16(s[2], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[2], 0xBBBBDEAD)
+
+
 class TestQuadmaskWqm(unittest.TestCase):
   """Tests for S_QUADMASK_B32 and S_WQM_B32."""
 

extra/assembly/amd/test/hw/test_vop3.py

@@ -2949,6 +2949,306 @@ class TestVOP3Clamp(unittest.TestCase):
     self.assertAlmostEqual(i2f(st.vgpr[3][1]), 1.0, places=5, msg="lane 3: 2.5 should clamp to 1.0")
 
 
+class TestVOP3ClampUint32(unittest.TestCase):
+  """Tests for VOP3 clamp modifier on unsigned 32-bit integer operations."""
+
+  def test_v_sub_nc_u32_e64_clamp_underflow(self):
+    """V_SUB_NC_U32_E64 with clamp: 0 - 1 should saturate to 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 1),
+      v_sub_nc_u32_e64(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_sub_nc_u32_e64_clamp_no_underflow(self):
+    """V_SUB_NC_U32_E64 with clamp: 100 - 50 = 50 (no saturation needed)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 50),
+      v_sub_nc_u32_e64(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 50, f"expected 50, got {st.vgpr[0][2]}")
+
+  def test_v_add_nc_u32_e64_clamp_overflow(self):
+    """V_ADD_NC_U32_E64 with clamp: 0xFFFFFFFF + 1 should saturate to 0xFFFFFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFFFFFF),
+      v_mov_b32_e32(v[1], 1),
+      v_add_nc_u32_e64(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_add_nc_u32_e64_clamp_no_overflow(self):
+    """V_ADD_NC_U32_E64 with clamp: 100 + 50 = 150 (no saturation needed)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 50),
+      v_add_nc_u32_e64(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 150, f"expected 150, got {st.vgpr[0][2]}")
+
+
+class TestVOP3ClampUint16(unittest.TestCase):
+  """Tests for VOP3 clamp modifier on unsigned 16-bit integer operations."""
+
+  def test_v_sub_nc_u16_clamp_underflow(self):
+    """V_SUB_NC_U16 with clamp: 0 - 1 should saturate to 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 1),
+      v_sub_nc_u16(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2] & 0xFFFF, 0, f"expected 0, got 0x{st.vgpr[0][2] & 0xFFFF:04x}")
+
+  def test_v_sub_nc_u16_clamp_no_underflow(self):
+    """V_SUB_NC_U16 with clamp: 100 - 50 = 50 (no saturation needed)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 50),
+      v_sub_nc_u16(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2] & 0xFFFF, 50, f"expected 50, got {st.vgpr[0][2] & 0xFFFF}")
+
+  def test_v_add_nc_u16_clamp_overflow(self):
+    """V_ADD_NC_U16 with clamp: 0xFFFF + 1 should saturate to 0xFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFF),
+      v_mov_b32_e32(v[1], 1),
+      v_add_nc_u16(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2] & 0xFFFF, 0xFFFF, f"expected 0xFFFF, got 0x{st.vgpr[0][2] & 0xFFFF:04x}")
+
+  def test_v_add_nc_u16_clamp_no_overflow(self):
+    """V_ADD_NC_U16 with clamp: 100 + 50 = 150 (no saturation needed)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 50),
+      v_add_nc_u16(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2] & 0xFFFF, 150, f"expected 150, got {st.vgpr[0][2] & 0xFFFF}")
+
+
+class TestVOP3ClampInt32(unittest.TestCase):
+  """Tests for VOP3 clamp modifier on signed 32-bit integer operations."""
+
+  def test_v_add_nc_i32_clamp_overflow(self):
+    """V_ADD_NC_I32 with clamp: INT_MAX + 1 should saturate to INT_MAX."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x7FFFFFFF),  # S0 = INT_MAX
+      v_mov_b32_e32(v[1], 1),            # S1 = 1
+      v_add_nc_i32(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF, f"expected 0x7FFFFFFF, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_add_nc_i32_clamp_underflow(self):
+    """V_ADD_NC_I32 with clamp: INT_MIN + (-1) should saturate to INT_MIN."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x80000000),  # S0 = INT_MIN
+      v_mov_b32_e32(v[1], 0xFFFFFFFF),  # S1 = -1
+      v_add_nc_i32(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x80000000, f"expected 0x80000000, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_sub_nc_i32_clamp_underflow(self):
+    """V_SUB_NC_I32 with clamp: INT_MIN - 1 should saturate to INT_MIN."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x80000000),  # S0 = INT_MIN
+      v_mov_b32_e32(v[1], 1),            # S1 = 1
+      v_sub_nc_i32(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x80000000, f"expected 0x80000000, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_sub_nc_i32_clamp_overflow(self):
+    """V_SUB_NC_I32 with clamp: INT_MAX - (-1) should saturate to INT_MAX."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x7FFFFFFF),  # S0 = INT_MAX
+      v_mov_b32_e32(v[1], 0xFFFFFFFF),  # S1 = -1
+      v_sub_nc_i32(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF, f"expected 0x7FFFFFFF, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_add_nc_i32_no_saturation_positive(self):
+    """V_ADD_NC_I32 with clamp: 100 + 200 = 300 (no saturation needed)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 200),
+      v_add_nc_i32(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 300, f"expected 300, got {st.vgpr[0][2]}")
+
+  def test_v_add_nc_i32_no_saturation_negative(self):
+    """V_ADD_NC_I32 with clamp: -100 + -200 = -300 (no saturation needed)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFFFF9C),  # -100
+      v_mov_b32_e32(v[1], 0xFFFFFF38),  # -200
+      v_add_nc_i32(v[2], v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    expected = 0xFFFFFED4  # -300
+    self.assertEqual(st.vgpr[0][2], expected, f"expected 0x{expected:08x}, got 0x{st.vgpr[0][2]:08x}")
+
+
+class TestVOP3ClampCarry(unittest.TestCase):
+  """Tests for VOP3 clamp modifier on carry operations (VOP3SD)."""
+
+  def test_v_add_co_u32_clamp_overflow(self):
+    """V_ADD_CO_U32 with clamp: 0xFFFFFFFF + 1 should saturate to 0xFFFFFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFFFFFF),
+      v_mov_b32_e32(v[1], 1),
+      v_add_co_u32(v[2], VCC, v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_add_co_u32_clamp_no_overflow(self):
+    """V_ADD_CO_U32 with clamp: 100 + 200 = 300 (no saturation)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 200),
+      v_add_co_u32(v[2], VCC, v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 300, f"expected 300, got {st.vgpr[0][2]}")
+
+  def test_v_sub_co_u32_clamp_underflow(self):
+    """V_SUB_CO_U32 with clamp: 0 - 1 should saturate to 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 1),
+      v_sub_co_u32(v[2], VCC, v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_sub_co_u32_clamp_no_underflow(self):
+    """V_SUB_CO_U32 with clamp: 300 - 100 = 200 (no saturation)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 300),
+      v_mov_b32_e32(v[1], 100),
+      v_sub_co_u32(v[2], VCC, v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 200, f"expected 200, got {st.vgpr[0][2]}")
+
+  def test_v_subrev_co_u32_clamp_underflow(self):
+    """V_SUBREV_CO_U32 with clamp: 1 - 0 reversed = 0 - 1 should saturate to 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 1),  # This becomes the subtrahend
+      v_mov_b32_e32(v[1], 0),  # This becomes the minuend (0 - 1)
+      v_subrev_co_u32(v[2], VCC, v[0], v[1], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_add_co_ci_u32_clamp_overflow(self):
+    """V_ADD_CO_CI_U32 with clamp: 0xFFFFFFFF + 1 + 0 should saturate to 0xFFFFFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFFFFFF),
+      v_mov_b32_e32(v[1], 1),
+      s_mov_b64(VCC, 0),  # No carry in
+      v_add_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_add_co_ci_u32_clamp_overflow_with_carry(self):
+    """V_ADD_CO_CI_U32 with clamp: 0xFFFFFFFE + 1 + 1 should saturate to 0xFFFFFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFFFFFE),
+      v_mov_b32_e32(v[1], 1),
+      s_mov_b64(VCC, 1),  # Carry in = 1
+      v_add_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_sub_co_ci_u32_clamp_underflow(self):
+    """V_SUB_CO_CI_U32 with clamp: 0 - 1 - 0 should saturate to 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 1),
+      s_mov_b64(VCC, 0),  # No borrow in
+      v_sub_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
+
+  def test_v_subrev_co_ci_u32_clamp_underflow(self):
+    """V_SUBREV_CO_CI_U32 with clamp: reversed 1 - 0 - 0 = 0 - 1 should saturate to 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 1),
+      v_mov_b32_e32(v[1], 0),
+      s_mov_b64(VCC, 0),
+      v_subrev_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
+
+
+class TestVOP3ClampMAD(unittest.TestCase):
+  """Tests for VOP3 clamp modifier on MAD (multiply-add) operations."""
+
+  def test_v_mad_u16_clamp_overflow(self):
+    """V_MAD_U16 with clamp: 0xFFFF * 2 + 0 should saturate to 0xFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFF),
+      v_mov_b32_e32(v[1], 2),
+      v_mov_b32_e32(v[2], 0),
+      v_mad_u16(v[3], v[0], v[1], v[2], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFF, f"expected 0xFFFF, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
+
+  def test_v_mad_u16_clamp_overflow_with_add(self):
+    """V_MAD_U16 with clamp: 0x8000 * 2 + 0x1000 should saturate to 0xFFFF."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x8000),  # 32768
+      v_mov_b32_e32(v[1], 2),        # * 2 = 65536
+      v_mov_b32_e32(v[2], 0x1000),   # + 4096 = 69632 > 0xFFFF
+      v_mad_u16(v[3], v[0], v[1], v[2], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFF, f"expected 0xFFFF, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
+
+  def test_v_mad_u16_no_overflow(self):
+    """V_MAD_U16 with clamp: 100 * 100 + 50 = 10050 (no saturation)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100),
+      v_mov_b32_e32(v[1], 100),
+      v_mov_b32_e32(v[2], 50),
+      v_mad_u16(v[3], v[0], v[1], v[2], clmp=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 10050, f"expected 10050, got {st.vgpr[0][3] & 0xFFFF}")
+
+  def test_v_mad_u16_no_clamp(self):
+    """V_MAD_U16 without clamp: 0xFFFF * 2 + 0 should wrap to 0xFFFE."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0xFFFF),
+      v_mov_b32_e32(v[1], 2),
+      v_mov_b32_e32(v[2], 0),
+      v_mad_u16(v[3], v[0], v[1], v[2], clmp=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # 0xFFFF * 2 = 0x1FFFE, low 16 bits = 0xFFFE
+    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFE, f"expected 0xFFFE, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
+
+
 class TestCvtPkF16(unittest.TestCase):
   """Tests for V_CVT_PK_RTZ_F16_F32 - pack two f32 to f16 with round toward zero."""
 

extra/assembly/amd/test/hw/test_vop3p.py

@@ -390,6 +390,24 @@ class TestVOP3P(unittest.TestCase):
     self.assertAlmostEqual(lo, 6.0, places=1)
     self.assertAlmostEqual(hi, 0.0, places=1)
 
+  def test_v_pk_add_u16_float_inline_const_opsel(self):
+    """V_PK_ADD_U16 with float inline constant 2.0
+    Regression test: for integer packed ops, do not perform the f32->f16 conversion.
+    """
+    # src1 = inline float constant 2.0
+    instructions = [
+      s_mov_b32(s[0], 0x00030005),  # packed u16: hi=3, lo=5
+      v_mov_b32_e32(v[0], s[0]),
+      v_pk_add_u16(v[1], v[0], SrcEnum.POS_TWO, opsel_hi=3, opsel_hi2=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    # lo = 5 + 0x0000 = 0x0005, hi = 3 + 0x4000 = 0x4003
+    self.assertEqual(lo, 0x0005, f"lo: expected 0x0005, got 0x{lo:04x}")
+    self.assertEqual(hi, 0x4003, f"hi: expected 0x4003, got 0x{hi:04x}")
+
 
 class TestWMMAF16(unittest.TestCase):
   """Tests for WMMA F16 output variant (V_WMMA_F16_16X16X16_F16).

extra/gemm/asm/cdna/gemm.py

@@ -35,9 +35,13 @@ def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
     return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
   batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
   N = b.shape[1]
-  if isinstance(a.device, tuple): batch //= len(a.device)
+  if isinstance(a.device, tuple):
+    batch //= len(a.device)
+    dname = a.device[0]
+  else: dname = a.device
+  arch = getattr(Device[dname].renderer, "arch", "")
   if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
-  if (key:=(M, N, K)) not in GEMM_ARGS: return todo(f"GEMM shape not supported {key}")
+  if (key:=(M, N, K)) not in GEMM_ARGS and arch == "gfx950": return todo(f"GEMM shape not supported {key} on {arch}")
   return True
 
 # ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
@@ -62,10 +66,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
   assert all_same([gradient.device, a.device, b.device, out.device])
   a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
   grad_a = (g_t @ b_t.T).uop
-  a_T = a_t.transpose(-2, -1)
-  a_T = a_T.reshape(*a_T.shape[:-1], 1, a_T.shape[-1])
-  g_r = g_t.reshape(*g_t.shape[:-2], 1, *g_t.shape[-2:]).transpose(-1, -2)
-  grad_b = (a_T * g_r).sum((-1, 0)).uop
+  grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
   return (None, grad_a, grad_b)
 
 # ** main gemm function
@@ -86,7 +87,7 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
     out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
 
   dname = a.device[0] if is_multi else a.device
-  arch = getattr(Device[dname].renderer, "arch", None)
+  arch = getattr(Device[dname].renderer, "arch", "")
   if arch.startswith("gfx950") and getenv("USE_ASM", 1):
     numWG = GEMM_ARGS[(M, N, K)][0]
     out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=numWG, arch=arch), grad_fxn=custom_gemm_bw)[0]

extra/gemm/asm/cdna/test_asm_gemm.py

@@ -1,46 +0,0 @@
-import unittest
-from tinygrad import Tensor, Device, dtypes, Context
-from tinygrad.helpers import getenv
-from extra.gemm.asm.cdna.gemm import asm_gemm
-
-def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, multi=False) -> None:
-  Tensor.manual_seed(0)
-  a_rand = Tensor.randn((batch, M, K), dtype=dtypes.float).sub(0.5).cast(dtype)
-  b_rand = Tensor.randn((K, N), dtype=dtypes.float).sub(0.5).cast(dtype)
-  with Context(DEBUG=0):
-    Tensor.realize(a_rand, b_rand)
-
-  devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8)) if multi else None
-
-  a, b = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
-  if multi: a, b = a.shard(devs, axis=0), b.shard(devs, axis=None)
-  tst = asm_gemm(a, b)
-  tst.sum().backward()
-  Tensor.realize(tst, a.grad, b.grad)
-
-  a_ref, b_ref = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
-  if multi: a_ref, b_ref = a_ref.shard(devs, axis=0), b_ref.shard(devs, axis=None)
-  with Context(ASM_GEMM=0): ref = a_ref @ b_ref
-  ref.sum().backward()
-  Tensor.realize(ref, a_ref.grad, b_ref.grad)
-
-  with Context(DEBUG=0):
-    assert (tst - ref).square().max().float().item() < 1e-6, "forward mismatch"
-    assert (a.grad - a_ref.grad).square().max().float().item() < 1e-3, "grad_a mismatch"
-    assert (b.grad - b_ref.grad).square().max().float().item() < 1e-3, "grad_b mismatch"
-
-class TestGemm(unittest.TestCase):
-  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=dtypes.half)
-
-  def test_gemm1(self): verify_asm_gemm(8, 8192, 4096, 14336, multi=True)
-  def test_gemm2(self): verify_asm_gemm(8, 8192, 128256, 4096, multi=True)
-  def test_gemm3(self): verify_asm_gemm(8, 8192, 14336, 4096, multi=True)
-  def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, multi=True)
-  def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, multi=True)
-  def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, multi=True)
-  def test_gemm_unsupported(self):
-    with self.assertRaisesRegex(AssertionError, "shape not supported"):
-      verify_asm_gemm(8, 8192, 1024, 4096, multi=True)
-
-if __name__ == "__main__":
-  unittest.main()

extra/gemm/intel_xmx.py

@@ -37,7 +37,7 @@ b.copyin(row.data)
 c.copyin(mat.data)
 ret = prog(a._buf, b._buf, c._buf, global_size=[1,1,1], local_size=[8,1,1], wait=True)
 print(ret)
-out = np.frombuffer(a.as_buffer(), np.float32)
+out = np.frombuffer(a.as_memoryview(), np.float32)
 real = row.astype(np.float32)@mat.T.astype(np.float32)
 print("out:", out)
 print("real", real)

extra/gemm/triton_nv_matmul.py

@@ -98,10 +98,10 @@ if __name__ == "__main__":
   # check correctness
   if getenv("VERIFY"):
     from tinygrad.engine.realize import run_schedule
-    triton_buf = np.frombuffer(si.bufs[0].as_buffer(), np.float16).reshape(M,N)
+    triton_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
     print(triton_buf)
     run_schedule(sched)
-    tinygrad_buf = np.frombuffer(si.bufs[0].as_buffer(), np.float16).reshape(M,N)
+    tinygrad_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
     print(tinygrad_buf)
     np.testing.assert_allclose(triton_buf, tinygrad_buf)
     print("correct!")

extra/hip_large_kernel.py

@@ -18,7 +18,7 @@ prg = dev.runtime("write_ones", mbin)
 prg(buf0._buf, global_size=(1,65537,1), local_size=(1,1,1), wait=True)
 
 import numpy as np
-def to_np(buf): return np.frombuffer(buf.as_buffer().cast(buf.dtype.base.fmt), dtype=_to_np_dtype(buf.dtype.base))
+def to_np(buf): return np.frombuffer(buf.as_memoryview().cast(buf.dtype.base.fmt), dtype=_to_np_dtype(buf.dtype.base))
 
 big = to_np(buf0)
 print(big)

extra/huggingface_onnx/huggingface_manager.py

@@ -8,14 +8,14 @@ from tinygrad.helpers import _ensure_downloads_dir
 DOWNLOADS_DIR = _ensure_downloads_dir() / "models"
 from tinygrad.helpers import tqdm
 
-def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, cache_dir: str|Path|None=None,
+def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, local_dir: str|Path|None=None,
                                  tries: int=2, **kwargs) -> Path:
   for attempt in range(tries):
     try:
       return Path(snapshot_download(
         repo_id=repo_id,
         allow_patterns=allow_patterns,
-        cache_dir=str(cache_dir) if cache_dir is not None else None,
+        local_dir=str(local_dir) if local_dir is not None else None,
         **kwargs
       ))
     except Exception as e:
@@ -144,14 +144,14 @@ class HuggingFaceONNXManager:
         root_path = snapshot_download_with_retry(
           repo_id=model_id,
           allow_patterns=allow_patterns,
-          cache_dir=str(self.models_dir)
+          local_dir=str(self.models_dir / model_id)
         )
 
         # Download config files (usually small)
         snapshot_download_with_retry(
           repo_id=model_id,
           allow_patterns=["*config.json"],
-          cache_dir=str(self.models_dir)
+          local_dir=str(self.models_dir / model_id)
         )
 
         model_data["download_path"] = str(root_path)

extra/huggingface_onnx/run_models.py

@@ -88,8 +88,8 @@ if __name__ == "__main__":
       # repo id
       # validates all onnx models inside repo
       repo_id = "/".join(path)
-      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=DOWNLOADS_DIR)
-      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], local_dir=DOWNLOADS_DIR / repo_id)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], local_dir=DOWNLOADS_DIR / repo_id)
       config = get_config(root_path)
       for onnx_model in root_path.rglob("*.onnx"):
         rtol, atol = get_tolerances(onnx_model.name)
@@ -101,8 +101,8 @@ if __name__ == "__main__":
       onnx_model = path[-1]
       assert path[-1].endswith(".onnx")
       repo_id, relative_path = "/".join(path[:2]), "/".join(path[2:])
-      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=DOWNLOADS_DIR)
-      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], local_dir=DOWNLOADS_DIR / repo_id)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], local_dir=DOWNLOADS_DIR / repo_id)
       config = get_config(root_path)
       rtol, atol = get_tolerances(onnx_model)
       print(f"validating {relative_path} with truncate={args.truncate}, {rtol=}, {atol=}")

extra/models/clip.py

@@ -5,7 +5,6 @@ from tinygrad.nn import Linear, LayerNorm, Embedding, Conv2d
 from typing import List, Optional, Union, Tuple, Dict
 from abc import ABC, abstractmethod
 from functools import lru_cache
-from PIL import Image
 import numpy as np
 import re, gzip
 
@@ -444,7 +443,8 @@ class OpenClipEncoder:
   # TODO:
   # Should be doable in pure tinygrad, would just require some work and verification.
   # This is very desirable since it would allow for full generation->evaluation in a single JIT call.
-  def prepare_image(self, image:Image.Image) -> Tensor:
+  def prepare_image(self, image) -> Tensor:
+    from PIL import Image
     SIZE = 224
     w, h = image.size
     scale = min(SIZE / h, SIZE / w)

extra/models/llama.py

@@ -61,6 +61,9 @@ class Attention:
       xq = self.q_norm(xq)
       xk = self.k_norm(xk)
 
+    # cast_float_to_bf16 is expensive in reduction loops, break it out
+    if x.dtype == dtypes.bfloat16: xq, xk = xq.contiguous_backward(), xk.contiguous_backward()
+
     xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
     xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
     xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
@@ -86,11 +89,14 @@ class Attention:
       assert start_pos == 0
       keys, values = xk, xv
 
-    keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
-    xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
-    attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
+    if Tensor.training:
+      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
+      attn = xq.scaled_dot_product_attention(keys, values, is_causal=True, enable_gqa=True).transpose(1, 2)
+    else:
+      keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
+      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
+      attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
     if getenv("STUB_ATTENTION"):
-      # TODO: do we need mask?
       from tinygrad.uop.ops import UOp, KernelInfo
       def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
         return UOp.sink(arg=KernelInfo(name="fa_custom_forward"))
@@ -197,7 +203,9 @@ class Transformer:
     h = self.tok_embeddings(tokens)
     freqs_cis = self.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :]
 
-    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None
+    if not Tensor.training and seqlen > 1:
+      mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1)
+    else: mask = None
     for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
     logits = self.output(self.norm(h))
     if math.isnan(temperature): return logits

extra/optimization/generate_dataset.sh

@@ -7,7 +7,7 @@ export CAPTURE_PROCESS_REPLAY=1
 rm "$LOGOPS" 2>/dev/null || true
 test/external/process_replay/reset.py
 
-CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/unit/test_winograd.py test/models/test_real_world.py --durations=20
+CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/unit/test_winograd.py test/null/test_real_world.py --durations=20
 CL=1 python3 -m pytest test/test_tiny.py
 
 # extract, sort and uniq

extra/thunder/tiny/fa.py

@@ -55,9 +55,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
       batch = ker.blockIdx_z
       q_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
 
-      k_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
-      v_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
-
       q_reg_fl = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
       q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
       q_reg_transposed = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
@@ -89,11 +86,8 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
       num_kv_blocks = (q_seq + 1) if is_causal else (N // KV_BLOCK_SIZE)
       for kv_idx in ker.range(num_kv_blocks):
-        k_smem = warp.load(k_smem, k, (), (batch, kv_idx, head_kv, 0), axis=1)
-        v_smem = warp.load(v_smem, v, (), (batch, kv_idx, head_kv, 0), axis=1)
-
-        k_reg = warp.load(k_reg, k_smem)
-        v_reg = warp.load(v_reg, v_smem)
+        k_reg = warp.load(k_reg, k, (), (batch, kv_idx, head_kv, 0), axis=1)
+        v_reg = warp.load(v_reg, v, (), (batch, kv_idx, head_kv, 0), axis=1)
 
         # mma qk^t
         att_block = warp.zero(att_block.after(kv_idx))
@@ -169,9 +163,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
       batch = ker.blockIdx_z
       q_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
 
-      k_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
-      v_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
-
       q_reg_fl = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
       q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
       q_reg_t = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
@@ -212,14 +203,12 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
       num_kv_blocks = (q_seq + 1) if is_causal else (N // KV_BLOCK_SIZE)
       for kv_idx in ker.range(num_kv_blocks):
-        k_smem = warp.load(k_smem, k, (), (batch, kv_idx, head_kv, 0), axis=1)
-        v_smem = warp.load(v_smem, v, (), (batch, kv_idx, head_kv, 0), axis=1)
+        k_reg = warp.load(k_reg, k, (), (batch, kv_idx, head_kv, 0), axis=1)
+        k_reg_col = warp.load(k_reg_col, k, (), (batch, kv_idx, head_kv, 0), axis=1)
+        v_reg = warp.load(v_reg, v, (), (batch, kv_idx, head_kv, 0), axis=1)
 
-        k_reg = warp.load(k_reg, k_smem)
         k_reg_t = warp.transpose(k_reg_t, k_reg)
-        k_reg_col = warp.load(k_reg_col, k_smem)
         k_reg_col_t = warp.transpose(k_reg_col_t, k_reg_col)
-        v_reg = warp.load(v_reg, v_smem)
 
         # mma qk^t
         att_block = warp.zero(att_block.after(kv_idx))
@@ -273,8 +262,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
       batch = ker.blockIdx_z
       kv_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
 
-      q_smem = ker.st((Q_BLOCK_SIZE, D), dtypes.bfloat16)
-      do_smem = ker.st((Q_BLOCK_SIZE, D), dtypes.bfloat16)
       att_smem = ker.st((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16)
 
       q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
@@ -308,19 +295,17 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
       k_reg_t = warp.transpose(k_reg_t, k_reg)
       v_reg = warp.load(v_reg, v, (), (batch, kv_seq, head_kv, 0), axis=1)
 
-      for q_idx in ker.range(N // Q_BLOCK_SIZE):
+      q_start = kv_seq if is_causal else 0
+      for q_idx in ker.range(q_start, N // Q_BLOCK_SIZE):
         for g in ker.range(GROUP_SIZE):
           head_q = head_kv * GROUP_SIZE + g
 
-          # load q and do
-          q_smem = warp.load(q_smem, q, (), (batch, q_idx, head_q, 0), axis=1)
-          do_smem = warp.load(do_smem, do, (), (batch, q_idx, head_q, 0), axis=1)
+          q_reg = warp.load(q_reg, q, (), (batch, q_idx, head_q, 0), axis=1)
+          q_reg_col = warp.load(q_reg_col, q, (), (batch, q_idx, head_q, 0), axis=1)
+          do_reg = warp.load(do_reg, do, (), (batch, q_idx, head_q, 0), axis=1)
+          do_reg_col = warp.load(do_reg_col, do, (), (batch, q_idx, head_q, 0), axis=1)
 
-          q_reg = warp.load(q_reg, q_smem)
           q_reg_t = warp.transpose(q_reg_t, q_reg)
-          q_reg_col = warp.load(q_reg_col, q_smem)
-          do_reg = warp.load(do_reg, do_smem)
-          do_reg_col = warp.load(do_reg_col, do_smem)
 
           # load l_vec and delta_vec
           l_vec_reg = warp.load(l_vec_reg, l_vec, (), (batch, head_q, 0, q_idx), axis=2)

extra/thunder/tiny/tk/kernel.py

@@ -12,7 +12,7 @@ class _tk_range:
   def __next__(self):
     if not self.done:
       self.done = True
-      self._rng = UOp.range(self.end // self.step, self.rid, axis_type=self.axis_type) * self.step + self.start
+      self._rng = UOp.range((self.end - self.start) // self.step, self.rid, axis_type=self.axis_type) * self.step + self.start
       return self._rng
     raise StopIteration
 

extra/torch_backend/backend.py

@@ -29,6 +29,10 @@ def wrap(x: Tensor) -> torch.Tensor:
   x._strides = strides_for_shape(x.shape) # always recalculate
   if (not hasattr(x, '_storage_offset')) or (not x.uop.is_realized): x._storage_offset = calculate_storage_offset(x)
   return mod.wrap(x, _to_torch_dtype(x.dtype), _to_torch_device(x.device).index)
+def _update_torch_metadata(tensor: torch.Tensor, tiny: Tensor) -> None:
+  tiny._strides = strides_for_shape(tiny.shape)
+  tiny._storage_offset = calculate_storage_offset(tiny)
+  mod.update_metadata(tensor, tiny.shape, tiny._strides, tiny._storage_offset)
 def unwrap(x:torch.Tensor) -> Tensor:
   assert isinstance(x, torch.Tensor), f"x isn't {type(x)}"
   return mod.unwrap(x)
@@ -344,7 +348,7 @@ def scatter_add(self, dim, index, src, out):
 def _copy_between_devices(src, dest, cast_dtype, to_device, non_blocking=False):
   if src.is_tiny and dest.is_tiny:
     src_t, dest_t = unwrap(src), unwrap(dest)
-    if dest_t.uop.is_contiguous() or dest_t.uop.is_realized: src_t = src_t.contiguous()
+    if dest_t.uop.has_buffer_identity() or dest_t.uop.is_realized: src_t = src_t.contiguous()
     _apply_inplace(dest_t, src_t.cast(cast_dtype).to(to_device))
   elif src.is_tiny and dest.is_cpu:
     dest.resize_(src.numel()).resize_(src.shape)
@@ -611,7 +615,10 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.fill_.Tensor": lambda self, value: Tensor.full(self.shape, value.reshape(()).item(), device=self.device, dtype=self.dtype),
   "aten.flip": Tensor.flip,
   "aten.scatter_reduce.two": Tensor.scatter_reduce,
-  "aten.squeeze_.dim": lambda self, dim: self.replace(self.squeeze(dim), allow_shape_mismatch=True), # TODO: inplace view op, here?
+  "aten.squeeze_.dim": Tensor.squeeze,
+  "aten.unsqueeze_": Tensor.unsqueeze,
+  "aten.transpose_": Tensor.transpose,
+  "aten.t_": Tensor.transpose,
   "aten.add.Tensor": lambda input,other,alpha=1: input+alpha*other,
   "aten.linspace": lambda start, stop, steps, dtype=None, **kwargs:
     Tensor.linspace(start, stop, steps, **({"dtype": _from_torch_dtype(dtype)} if dtype is not None else {})),
@@ -655,6 +662,13 @@ inplace_ops = {
   "aten.masked_fill_.Tensor",
 }
 
+inplace_view_ops = {
+  "aten.squeeze_.dim",
+  "aten.unsqueeze_",
+  "aten.transpose_",
+  "aten.t_",
+}
+
 def wrap_fxn(k,f):
   def nf(*args, **kwargs):
     if TORCH_DEBUG:
@@ -675,8 +689,42 @@ def wrap_inplace(k,f):
     return orig
   return nf
 
+def wrap_inplace_view_op(k,f):
+  def nf(*args, **kwargs):
+    orig = args[0]
+    args, kwargs = unwrap_args(args, kwargs)
+    target = args[0]
+    new_view = f(*args, **kwargs)
+    if new_view is target or new_view.uop is target.uop:
+      _update_torch_metadata(orig, target)
+      return orig
+    base = canonical_base(target)
+    op = (f, args[1:], kwargs)
+    if target is base:
+      views = derived_views(base)
+      if views:
+        old_base = Tensor(base.uop, device=base.device)
+        old_base.requires_grad = base.requires_grad
+        old_base._views = getattr(base, "_views", set())
+        for v in views: v._view_base = old_base
+        base._views = set()
+        base._view_base = old_base
+        base._view_ops = [op]
+        old_base._views.add(weakref.ref(base))
+    else:
+      target._view_base = base
+      base._views = getattr(base, "_views", set())
+      base._views.add(weakref.ref(target))
+      target._view_ops = _get_view_ops(target) + [op]
+    target.uop = new_view.uop
+    _update_torch_metadata(orig, target)
+    return orig
+  return nf
+
 for k,v in tiny_backend.items():
-  wrapper = wrap_inplace if k in inplace_ops else wrap_fxn
+  if k in inplace_view_ops: wrapper = wrap_inplace_view_op
+  elif k in inplace_ops: wrapper = wrap_inplace
+  else: wrapper = wrap_fxn
   torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrapper(k,v))
 
 @torch.library.impl("aten::equal", "privateuseone")

extra/torch_backend/test_inplace.py

@@ -67,5 +67,37 @@ class TestTorchBackendInplace(unittest.TestCase):
     d += torch.arange(4)
     np.testing.assert_array_equal(a.cpu(), torch.arange(4).cpu())
 
+  def test_inplace_view_metadata(self):
+    a = torch.arange(6, dtype=torch.float32).reshape(1, 2, 3)
+    ret = a.squeeze_(0)
+    self.assertIs(ret, a)
+    self.assertEqual(a.shape, torch.Size([2, 3]))
+    ret = a.unsqueeze_(1)
+    self.assertIs(ret, a)
+    self.assertEqual(a.shape, torch.Size([2, 1, 3]))
+    ret = a.transpose_(0, 2)
+    self.assertIs(ret, a)
+    self.assertEqual(a.shape, torch.Size([3, 1, 2]))
+
+  def test_t_inplace_metadata(self):
+    a = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    ret = a.t_()
+    self.assertIs(ret, a)
+    self.assertEqual(a.shape, torch.Size([3, 2]))
+    expected = torch.arange(6, dtype=torch.float32).reshape(2, 3).t()
+    np.testing.assert_array_equal(a.cpu().numpy(), expected.cpu().numpy())
+
+  def test_squeeze_matmul(self):
+    # squeeze_ is used internally by PyTorch for vector-matrix matmul (unsqueeze -> mm -> squeeze_)
+    a = torch.arange(65, dtype=torch.float32)
+    b = torch.arange(65*45, dtype=torch.float32).reshape(65, 45)
+    result = a.matmul(b)
+    self.assertEqual(result.shape, torch.Size([45]))
+    # verify correctness
+    a_cpu = torch.arange(65, dtype=torch.float32, device='cpu')
+    b_cpu = torch.arange(65*45, dtype=torch.float32, device='cpu').reshape(65, 45)
+    expected = a_cpu.matmul(b_cpu)
+    np.testing.assert_allclose(result.cpu().numpy(), expected.numpy(), rtol=1e-4, atol=1e-4)
+
 if __name__ == "__main__":
   unittest.main()

extra/torch_backend/wrapped_tensor.cpp

@@ -131,7 +131,15 @@ py::object unwrap_tensor(const at::Tensor &tensor) {
   return py::reinterpret_borrow<py::object>(tiny->ptr(getPyInterpreter()));
 }
 
+void update_metadata(const at::Tensor &tensor, const std::vector<int64_t> &sizes,
+                     const std::vector<int64_t> &strides, int64_t storage_offset) {
+  auto* impl = tensor.unsafeGetTensorImpl();
+  impl->set_allow_tensor_metadata_change(true);
+  impl->set_sizes_and_strides(sizes, strides, storage_offset);
+}
+
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("wrap", &wrap_tensor);
   m.def("unwrap", &unwrap_tensor);
+  m.def("update_metadata", &update_metadata);
 }

extra/viz/cli.py

@@ -3,7 +3,8 @@ import argparse, pathlib
 from typing import Iterator
 from tinygrad.viz import serve as viz
 from tinygrad.uop.ops import RewriteTrace
-from tinygrad.helpers import temp, ansistrip, colored
+from tinygrad.helpers import temp, ansistrip, colored, time_to_str, ansilen
+from test.null.test_viz import load_profile
 
 def optional_eq(val:dict, arg:str|None) -> bool: return arg is None or ansistrip(val["name"]) == arg
 
@@ -27,11 +28,46 @@ if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument('--kernel', type=str, default=None, metavar="NAME", help='Select a kernel by name (optional name, default: only list names)')
   parser.add_argument('--select', type=str, default=None, metavar="NAME",
-                      help='Select an item within the chosen kernel (optional name, default: only list names)')
+                      help='Rewrites: Select an item within the chosen kernel (optional name, default: only list names)')
+  parser.add_argument('--profile', action="store_true", help="View profiling trace (default: views rewrites)")
+  parser.add_argument('--device', type=str, default=None, metavar="NAME", help="Profile only: Select a device (default: prints all devices)")
+  parser.add_argument('--profile-path', type=pathlib.Path, metavar="PATH", help='Path to profile (optional file, default: latest profile)',
+                      default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  parser.add_argument('--rewrites-path', type=pathlib.Path, metavar="PATH", help='Path to rewrites (optional file, default: latest rewrites)',
+                      default=pathlib.Path(temp("rewrites.pkl", append_user=True)))
   args = parser.parse_args()
 
-  viz.trace = viz.load_pickle(pathlib.Path(temp("rewrites.pkl", append_user=True)), default=RewriteTrace([], [], {}))
+  viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
   viz.ctxs = viz.get_rewrites(viz.trace)
+
+  if args.profile:
+    from tabulate import tabulate
+    profile = load_profile(viz.load_pickle(args.profile_path, default=[]))
+    agg, total, n = {}, 0, 0
+    for k,v in profile["layout"].items():
+      if not optional_eq({"name":k}, args.device): continue
+      print(k)
+      if args.device is None: continue
+      for e in v.get("events", []):
+        et = e["dur"]*1e-6
+        if args.kernel is not None:
+          if ansistrip(e["name"]) == args.kernel and n < 10:
+            ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
+            name = e["name"]+(" " * (46 - ansilen(e["name"])))
+            print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms  "+e['fmt'].replace('\n', ' | ')+"  ")
+            n += 1
+        else:
+          a = agg.setdefault(e["name"], [0.0, 0])
+          a[0] += et
+          a[1] += 1
+          total += et
+
+    if agg:
+      rows = [[n, t, time_to_str(t, w=9), t / c if c else 0.0, c, (t / total * 100.0) if total else 0.0] for n, (t, c) in agg.items()]
+      rows.sort(key=lambda r: r[1], reverse=True)
+      print(tabulate([[r[0], r[2], r[4], f"{r[5]:.2f}%"] for r in rows[:30]], headers=["name", "total", "count", "pct"], tablefmt="github"))
+    exit(0)
+
   for k in viz.ctxs:
     if not optional_eq(k, args.kernel): continue
     print(k["name"])

test/device/test_hcq.py

@@ -119,7 +119,7 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
 
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
 
   def test_exec_2_kernels_100_times(self):
@@ -135,7 +135,7 @@ class TestHCQ(unittest.TestCase):
       q.submit(TestHCQ.d0, {virt_val.expr: TestHCQ.d0.timeline_value})
       TestHCQ.d0.timeline_value += 1
 
-    val = TestHCQ.a.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 200.0, f"got val {val}"
 
   @unittest.skipIf(Device.DEFAULT in {"CPU"}, "No globals/locals on LLVM/CPU")
@@ -151,9 +151,9 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
 
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 0.0, f"got val {val}, should not be updated"
 
   @unittest.skipIf(Device.DEFAULT in {"CPU"}, "No globals/locals on LLVM/CPU")
@@ -186,7 +186,7 @@ class TestHCQ(unittest.TestCase):
           TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
           TestHCQ.d0.timeline_value += 1
 
-          res_sum = sum(x for x in zt.as_buffer().cast("I"))
+          res_sum = sum(x for x in zt.as_memoryview().cast("I"))
           assert x * y * z == res_sum, f"want {x * y * z}, got {res_sum}"
 
   # Test copy
@@ -200,7 +200,7 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
 
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 1.0, f"got val {val}"
 
   def test_copy_long(self):
@@ -218,7 +218,7 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
 
-    mv_buf1 = buf1.as_buffer().cast('Q')
+    mv_buf1 = buf1.as_memoryview().cast('Q')
     assert libc.memcmp(mv_address(mv_buf1), buf2._buf.va_addr, sz) == 0
 
   @slow
@@ -242,7 +242,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
 
-      mv_buf1 = buf1.as_buffer()
+      mv_buf1 = buf1.as_memoryview()
       assert libc.memcmp(mv_address(mv_buf1), buf2._buf.va_addr, sz) == 0
 
   def test_update_copy(self):
@@ -260,7 +260,7 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
 
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 1.0, f"got val {val}"
 
   def test_update_copy_long(self):
@@ -283,7 +283,7 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
 
-    mv_buf1 = buf1.as_buffer().cast('Q')
+    mv_buf1 = buf1.as_memoryview().cast('Q')
     for i in range(sz//8): assert mv_buf1[i] == 0x0101010101010101, f"offset {i*8} differs, not all copied, got {hex(mv_buf1[i])}"
 
   # Test bind api
@@ -421,7 +421,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
 
-      assert buf1.as_buffer()[0] == i
+      assert buf1.as_memoryview()[0] == i
 
   def test_small_copies_from_host_buf_intercopy(self):
     if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -440,7 +440,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
 
-      assert buf2.as_buffer()[0] == i
+      assert buf2.as_memoryview()[0] == i
 
   def test_small_copies_from_host_buf_transfer(self):
     if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -463,7 +463,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
 
-      assert buf2.as_buffer()[0] == i
+      assert buf2.as_memoryview()[0] == i
 
   def test_memory_barrier(self):
     a = Tensor([0, 1], device=Device.DEFAULT, dtype=dtypes.int8).realize()
@@ -486,7 +486,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
 
-      assert buf1.as_buffer()[0] == (i + 1), f"has {buf1.as_buffer()[0]}, need {i + 1}"
+      assert buf1.as_memoryview()[0] == (i + 1), f"has {buf1.as_memoryview()[0]}, need {i + 1}"
 
   def test_memory_barrier_before_copy(self):
     if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -511,7 +511,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
 
-      assert buf2.as_buffer()[0] == i
+      assert buf2.as_memoryview()[0] == i
 
   def test_map_cpu_buffer_to_device(self):
     if Device[Device.DEFAULT].hw_copy_queue_t is None: self.skipTest("skip device without copy queue")

test/external/external_fuzz_qcom_cpu_cache.py

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

test/external/external_test_amd.py

@@ -20,7 +20,7 @@ class TestAMD(unittest.TestCase):
                               global_size=TestAMD.d0_runner.global_size, local_size=TestAMD.d0_runner.local_size)
       TestAMD.d0_runner.clprg(TestAMD.a.uop.buffer._buf, TestAMD.b.uop.buffer._buf,
                               global_size=TestAMD.d0_runner.global_size, local_size=TestAMD.d0_runner.local_size)
-    val = TestAMD.a.uop.buffer.as_buffer().cast("f")[0]
+    val = TestAMD.a.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 4000.0, f"got val {val}"
 
 if __name__ == "__main__":

test/external/external_test_hcq.py

@@ -65,7 +65,7 @@ class TestHCQ(unittest.TestCase):
     q.submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.a.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 2000.0, f"got val {val}"
 
   def test_run_1000_times(self):
@@ -81,7 +81,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.compute_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
       TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.a.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 2000.0, f"got val {val}"
 
   def test_run_to_3(self):
@@ -95,7 +95,7 @@ class TestHCQ(unittest.TestCase):
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 3.0, f"got val {val}"
 
   def test_update_exec(self):
@@ -106,9 +106,9 @@ class TestHCQ(unittest.TestCase):
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 0.0, f"got val {val}, should not be updated"
 
   @unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind")
@@ -126,7 +126,7 @@ class TestHCQ(unittest.TestCase):
       TestHCQ.compute_queue().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
       TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.a.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 2000.0, f"got val {val}"
 
   @unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind")
@@ -141,9 +141,9 @@ class TestHCQ(unittest.TestCase):
     q.submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 0.0, f"got val {val}, should not be updated"
 
   @unittest.skipIf(CI, "Can't handle async update on CPU")
@@ -174,7 +174,7 @@ class TestHCQ(unittest.TestCase):
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
 
   def test_submit_empty_queues(self):
@@ -206,7 +206,7 @@ class TestHCQ(unittest.TestCase):
     q.submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
 
   def test_copy_1000_times(self):
@@ -221,7 +221,7 @@ class TestHCQ(unittest.TestCase):
     # confirm the signal didn't exceed the put value
     with self.assertRaises(RuntimeError):
       TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value + 1, timeout=50)
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 0.0, f"got val {val}"
 
   def test_copy(self):
@@ -231,7 +231,7 @@ class TestHCQ(unittest.TestCase):
     q.submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 1.0, f"got val {val}"
 
   @unittest.skipUnless(Device.DEFAULT == "NV", "Only NV supports bind")
@@ -248,7 +248,7 @@ class TestHCQ(unittest.TestCase):
     # confirm the signal didn't exceed the put value
     with self.assertRaises(RuntimeError):
       TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value + 1, timeout=50)
-    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
+    val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
     assert val == 0.0, f"got val {val}"
 
   def test_copy_bandwidth(self):
@@ -288,7 +288,7 @@ class TestHCQ(unittest.TestCase):
     q.submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
-    val = TestHCQ.a.uop.buffer.as_buffer().cast("f")[0]
+    val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
     assert val == 1.0, f"got val {val}"
 
   def test_cross_device_signal(self):
@@ -319,7 +319,7 @@ class TestHCQ(unittest.TestCase):
       q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
       TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1
-      val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[0]
+      val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
       assert val == 1.0, f"got val {val}"
 
 if __name__ == "__main__":

test/external/fuzz_graph.py

@@ -30,7 +30,7 @@ def alloc_rawbuffer(device, fill=False):
   if fill:
     with Context(DEBUG=0):
       data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
-      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_buffer())
+      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_memoryview())
   return rawbuf
 
 def gen_kernel_ji(device, deps):
@@ -93,7 +93,7 @@ def run_jit(jis, all_buffers, input_buffers, var_vals):
 
   with Context(DEBUG=0):
     res_buffers = []
-    for rawbuf in all_buffers: res_buffers.append(rawbuf.as_buffer())
+    for rawbuf in all_buffers: res_buffers.append(rawbuf.as_memoryview())
     return res_buffers
 
 def fuzz_graph(jis, all_buffers, input_buffers):

test/models/test_onnx.py

@@ -73,7 +73,7 @@ class TestHuggingFaceOnnxModels(unittest.TestCase):
     onnx_model_path = snapshot_download_with_retry(
       repo_id=repo_id,
       allow_patterns=["*.onnx", "*.onnx_data"],
-      cache_dir=str(DOWNLOADS_DIR)
+      local_dir=DOWNLOADS_DIR / repo_id
     )
     onnx_model_path = onnx_model_path / model_file
     file_size = onnx_model_path.stat().st_size

test/null/test_attention.py

@@ -0,0 +1,39 @@
+import unittest
+from tinygrad import Tensor, dtypes, TinyJit, UOp
+from tinygrad.apps.llm import apply_rope as apply_rope_new, precompute_freqs_cis
+
+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()
+    # attention has 4 kernels now
+    self.assertEqual(len(sched), 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))
+    noprune_size = len(rope_noprune.captured.jit_cache)
+    prune_size = len(rope_prune.captured.jit_cache)
+
+    self.assertGreater(noprune_size, prune_size)
+    self.assertGreaterEqual(noprune_size, 2)
+    self.assertEqual(prune_size, 1)
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_const_folding.py

@@ -0,0 +1,149 @@
+import unittest, itertools, math
+from tinygrad import Tensor, dtypes, Context
+from tinygrad.dtype import DType, ConstType
+from tinygrad.uop.ops import Ops, UOp
+from tinygrad.codegen import full_rewrite_to_sink
+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
+  schedule = t.schedule()
+  asts = [s for s in schedule if s.ast.op is Ops.SINK]
+  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_idiv_literal_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
+  def test_idiv_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_to_sink(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):
+      r = full_rewrite_to_sink(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src[0]
+    self.assertEqual(r.op, Ops.VECTORIZE)
+    self.assertEqual(r.dtype, dtypes.uint32.vec(3))
+    self.assertEqual(tuple(x.arg for x in r.src), (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).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).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()

test/null/test_dtype_spec.py

@@ -0,0 +1,413 @@
+import unittest, math, struct, operator
+from tinygrad.tensor import Tensor, dtypes
+from tinygrad.dtype import DTYPES_DICT, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
+from tinygrad.device import is_dtype_supported
+
+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 is_dtype_supported(dt)]
+dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and is_dtype_supported(dt)]
+
+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(dtypes.min(dt), -math.inf)
+        np.testing.assert_equal(dtypes.max(dt), 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(dtypes.min(dt), info.min)
+        np.testing.assert_equal(dtypes.max(dt), 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(dtypes.min(dt), False)
+        np.testing.assert_equal(dtypes.max(dt), 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())
+
+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
+
+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
+
+  @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()

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, requires_grad=True).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, requires_grad=True).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, requires_grad=True).realize()
+    y = x * 2
+    y.sum().gradient(x)[0].realize()
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_helpers.py

@@ -0,0 +1,469 @@
+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
+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}"
+
+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_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()

test/null/test_indexing.py

@@ -0,0 +1,143 @@
+# test cases are modified from pytorch test_indexing.py
+
+import unittest
+import numpy as np
+
+from tinygrad import Tensor, dtypes
+
+def numpy_testing_assert_equal_helper(a, b):
+  if isinstance(a, Tensor): a = a.numpy()
+  if isinstance(b, Tensor): b = b.numpy()
+  np.testing.assert_equal(a, b)
+
+class TestIndexing(unittest.TestCase):
+  def test_single_int(self):
+    v = Tensor.randn(5, 7, 3)
+    numpy_testing_assert_equal_helper(v[4].shape, (7, 3))
+
+  def test_multiple_int(self):
+    v = Tensor.randn(5, 7, 3)
+    numpy_testing_assert_equal_helper(v[4].shape, (7, 3))
+    numpy_testing_assert_equal_helper(v[4, :, 1].shape, (7,))
+
+  def test_none(self):
+    v = Tensor.randn(5, 7, 3)
+    numpy_testing_assert_equal_helper(v[None].shape, (1, 5, 7, 3))
+    numpy_testing_assert_equal_helper(v[:, None].shape, (5, 1, 7, 3))
+    numpy_testing_assert_equal_helper(v[:, None, None].shape, (5, 1, 1, 7, 3))
+    numpy_testing_assert_equal_helper(v[..., None].shape, (5, 7, 3, 1))
+
+  def test_int_indices(self):
+    v = Tensor.randn(5, 7, 3)
+    numpy_testing_assert_equal_helper(v[[0, 4, 2]].shape, (3, 7, 3))
+    numpy_testing_assert_equal_helper(v[:, [0, 4, 2]].shape, (5, 3, 3))
+    numpy_testing_assert_equal_helper(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], :, :]
+    numpy_testing_assert_equal_helper(res.shape, src.shape)
+
+  def test_empty_slice(self):
+    x = Tensor.randn(2, 3, 4, 5)
+    y = x[:, :, :, 1]
+    z = y[:, 1:1, :]
+    numpy_testing_assert_equal_helper((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])
+
+  def test_take_along_dim(self):
+    # NOTE: the actual test logic is inside _test_against_numpy which is never called
+    # This test effectively does nothing but defines a function
+    def _test_against_numpy(t: Tensor, indices: Tensor, dim):
+      actual = t.gather(dim, indices)
+      t_np = t.numpy()
+      indices_np = indices.numpy()
+      expected = np.take_along_axis(t_np, indices_np, axis=dim)
+      numpy_testing_assert_equal_helper(actual, expected)
+
+      # TODO argsort
+      '''
+      for shape in [(3, 2), (2, 3, 5), (2, 4, 0), (2, 3, 1, 4)]:
+        for noncontiguous in [True, False]:
+          for dtype in (dtypes.float32, dtypes.int64):
+            t = make_tensor(shape, dtype=dtype, noncontiguous=noncontiguous)
+            for dim in list(range(t.ndim)) + [None]:
+              if dim is None:
+                indices = argsort(t.reshape(-1))
+              else:
+                indices = argsort(t, dim=dim)
+
+          _test_against_numpy(t, indices, dim)
+      '''
+
+      # test broadcasting
+      t = Tensor.ones((3, 4, 1))
+      indices = Tensor.ones((1, 2, 5), dtype=dtypes.int64)
+
+      _test_against_numpy(t, indices, 1)
+
+      # test empty indices
+      t = Tensor.ones((3, 4, 5))
+      indices = Tensor.ones((3, 0, 5), dtype=dtypes.int64)
+
+      _test_against_numpy(t, indices, 1)
+
+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])
+    numpy_testing_assert_equal_helper(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()

test/null/test_llm_server.py

@@ -0,0 +1,137 @@
+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.end_turn = Mock(return_value=[998])
+
+    cls.mock_model = Mock()
+    cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
+
+    cls.bos_id = 1
+    cls.eos_id = 999
+
+    import tinygrad.apps.llm as llm_module
+    llm_module.model = cls.mock_model
+    llm_module.tok = cls.mock_tok
+    llm_module.bos_id = cls.bos_id
+    llm_module.eos_id = cls.eos_id
+
+    from tinygrad.apps.llm import Handler
+    from tinygrad.viz.serve import TCPServerWithReuse
+
+    cls.server = TCPServerWithReuse(('127.0.0.1', 0), Handler)
+    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)
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_multitensor.py

@@ -0,0 +1,58 @@
+import unittest
+from tinygrad import Tensor, GlobalCounters, dtypes
+
+class TestMultiRamUsage(unittest.TestCase):
+  def setUp(self):
+    self.baseline = GlobalCounters.mem_used
+    self.N = 100
+  def assertUsed(self, amt, strict=True):
+    used = GlobalCounters.mem_used - self.baseline
+    print(f"used {used} bytes")
+    if strict: self.assertEqual(used, amt)
+    else: self.assertLessEqual(used, amt)
+
+  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)
+
+  @unittest.skip("flaky")
+  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)
+
+  @unittest.skip("flaky")
+  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"))
+
+  @unittest.skip("flaky")
+  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_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)
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_real_world.py

@@ -12,7 +12,7 @@ from test.helpers import derandomize_model
 
 from examples.gpt2 import Transformer as GPT2Transformer
 from examples.hlb_cifar10 import SpeedyResNet, hyp
-from examples.llama import Transformer as LLaMaTransformer
+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

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()

test/null/test_schedule.py

@@ -0,0 +1,143 @@
+# schedule tests that pass on NULL backend (no copyout needed)
+import unittest
+from tinygrad import Tensor
+from tinygrad.uop.ops import UOp
+from tinygrad.helpers import DEBUG, Context
+from tinygrad.engine.realize import CompiledRunner, run_schedule
+
+class KernelCountException(Exception): pass
+def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
+  if to_prerealize:
+    with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize)
+  if isinstance(t, Tensor): sched = t.schedule()
+  elif isinstance(t, list) and isinstance(t[0], Tensor): sched = Tensor.schedule(*t)
+  else:
+    assert isinstance(t, UOp), f"can't schedule {t}"
+    sched = Tensor(t).schedule()
+  # test lowering all the ExecItems
+  for si in sched: si.lower()
+  kernel_cnt = len([si for si in sched if isinstance(si.prg, CompiledRunner) or not filter_sink])
+  if kernel_cnt != allowed:
+    print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}")
+    if DEBUG >= 3:
+      for i,s in enumerate(sched):
+        print("kernel", i+1)
+        print(s.ast)
+    raise KernelCountException(f"{kernel_cnt} != {allowed}")
+  return sched
+
+class TestBufferUOp(unittest.TestCase):
+  # BUFFER has a ShapeTracker of shape=(n,) and stride=(1,)
+  def test_buffer_has_buffer(self):
+    buf = Tensor.empty(10)
+    self.assertIsNotNone(buf.uop.buffer)
+    self.assertEqual(buf.uop.shape, (10,))
+    # the device Buffer remains unallocated until it's we run the schedule
+    self.assertFalse(buf.uop.buffer.is_allocated())
+    add = buf+1
+    sched = add.schedule()
+    self.assertFalse(buf.uop.buffer.is_allocated())
+    run_schedule(sched)
+    self.assertTrue(buf.uop.buffer.is_allocated())
+
+  def test_buffer_has_unique_buffer(self):
+    buf = Tensor.empty(10)
+    buf1 = buf.uop.buffer
+    buf2 = buf.uop.buffer
+    self.assertIs(buf1, buf2)
+
+  # we also allow VIEW(BUFFER) to access the underlying device Buffer, as long as it's contiguous
+  def test_buffer_view_allowed(self):
+    add = Tensor.empty(1, 1)+Tensor.empty(1, 1)
+    add.realize()
+    self.assertIsNotNone(add.uop.buffer)
+    self.assertEqual(add.uop.shape, (1, 1))
+
+  def test_buffer_view_not_allowed(self):
+    permuted_view = Tensor.empty(1, 2, 3).permute(0, 2, 1)
+    with self.assertRaisesRegex(AssertionError, "can only be RESHAPE"):
+      permuted_view.uop.buffer # cannot access Buffer of a non contiguous VIEW
+
+  def test_buffer_only_after_realize(self):
+    a = Tensor([1])+Tensor([2])
+    # accessing realized will return None
+    self.assertIsNone(a.uop.realized)
+    # accessing Buffer will assert
+    with self.assertRaisesRegex(AssertionError, "must be BUFFER"):
+      a.uop.buffer # there is no BUFFER on an unrealized ADD
+    # Buffer only exists once we realize it
+    a.realize()
+    self.assertIsNotNone(a.uop.buffer)
+
+  def test_const_does_not_realize(self):
+    a = Tensor(1)+Tensor(2)
+    run_schedule(check_schedule(a, 0))
+    self.assertIsNone(a.uop.base.realized)
+
+  def test_var_does_not_realize(self):
+    a = Tensor(UOp.variable("a", 0, 10).bind(1))
+    run_schedule(check_schedule(a, 0))
+    self.assertIsNone(a.uop.base.realized)
+
+  def test_unused_var_not_in_var_vals(self):
+    # unused variable should not appear in var_vals even when there's other work
+    a = Tensor(UOp.variable("unused", 0, 10).bind(1))
+    b = Tensor.empty(3) + 1
+    _, var_vals = Tensor.schedule_with_vars(a, b)
+    self.assertEqual(var_vals, {})
+    self.assertIsNone(a.uop.base.realized)
+
+  def test_view_does_not_realize(self):
+    a = Tensor.randn(1, 4).expand(4, 4)
+    a.realize()
+    self.assertEqual(a.uop.base.realized.size, 4)
+    a2 = a.contiguous().realize()
+    self.assertEqual(a2.uop.base.realized.size, 16)
+
+class TestContiguous(unittest.TestCase):
+  def test_contiguous_buffer(self):
+    a = Tensor.empty(4)
+    b = a.contiguous()
+    check_schedule(b, 0)
+
+  def test_contiguous_buffer_view(self):
+    a = Tensor.empty(4)
+    b = a.reshape((2, 2)).contiguous()
+    check_schedule(b, 0)
+
+  def test_non_contiguous_buffer_view(self):
+    a = Tensor.empty(4, 1)
+    b = a.expand((4, 4)).contiguous()
+    check_schedule(b, 1)
+
+  def test_size_change_buffer_view(self):
+    a = Tensor.empty(4)
+    b = a.reshape((1, 1, 4)).shrink(((0, 1), (0, 1), (0, 3))).contiguous()
+    check_schedule(b, 1)
+
+  def test_double_contiguous_realizes_once(self):
+    a = Tensor.empty(4, 1)
+    b = a.expand((4, 4)).contiguous().contiguous()
+    check_schedule(b, 1)
+
+  def test_view_does_not_realize(self):
+    a = Tensor.empty(4)
+    b = a.expand((4, 4))
+    check_schedule(b, 0)
+    self.assertEqual(b.uop.base.buffer.size, 4)
+
+  def test_contiguous_view_realizes(self):
+    a = Tensor.empty(4)
+    b = a.expand((4, 4)).contiguous()
+    check_schedule(b, 1)
+    self.assertEqual(b.uop.base.buffer.size, 16)
+
+class TestSimpleSchedule(unittest.TestCase):
+  def test_reduce_doesnt_split(self):
+    a = Tensor.empty(16,16).sum(axis=1)
+    a1 = a.reshape(4,4)
+    a2 = a.reshape(16,1,1)
+    self.assertEqual(len(Tensor.schedule(a1, a2)), 1)
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)

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.engine.schedule import schedule_cache
+
+def schedule_one():
+  Tensor([1]).schedule()
+
+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.schedule_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()

test/null/test_symbolic_tensor.py

@@ -0,0 +1,101 @@
+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)
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_tensor.py

@@ -0,0 +1,166 @@
+# tensor tests that pass on NULL backend (no copyout needed)
+import numpy as np
+import unittest
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.device import is_dtype_supported
+from tinygrad.uop.ops import Ops, UOp
+from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
+from tinygrad.engine.realize import get_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, requires_grad=True)
+    m = Tensor(m_init, requires_grad=True)
+    W = Tensor(W_init, requires_grad=True)
+    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
+    assert W.requires_grad
+
+  def test_no_grad_mode_context_manager(self):
+    x = Tensor(x_init, requires_grad=True)
+    m = Tensor(m_init, requires_grad=True)
+    W = Tensor(W_init, requires_grad=True)
+    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):
+    schedule, _ = a.schedule_with_vars()
+    for s in schedule:
+      if s.ast.op is Ops.SINK:
+        renderer = Device[s.bufs[0].device].renderer
+        prg = get_program(s.ast, renderer)
+        return prg.uops
+
+  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]
+      assert idx_val.dtype is 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(is_dtype_supported(dtypes.long), "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(is_dtype_supported(dtypes.long), "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, 2048, 2048, 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)
+
+  def test_arange_raise_overflow(self):
+    with self.assertRaises(ValueError):
+      self._schedule_render(Tensor.arange(2**33, dtype=dtypes.int))
+
+  @unittest.skipIf(is_dtype_supported(dtypes.long), "int64 is supported")
+  def test_int64_unsupported_overflow_sym(self):
+    with self.assertRaises(KeyError):
+      self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 1, 2048).bind(32))
+
+  @unittest.skipIf(is_dtype_supported(dtypes.long), "int64 is supported")
+  @unittest.expectedFailure  # bug in gpu dims limiting
+  def test_int64_unsupported_overflow(self):
+    with self.assertRaises(KeyError):
+      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_eye_bufs_unique(self):
+    a = Tensor.eye(10).contiguous()
+    b = Tensor.eye(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)
+
+if __name__ == '__main__':
+  unittest.main()

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()

test/null/test_uops.py

@@ -0,0 +1,334 @@
+# 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
+from tinygrad.dtype import dtypes, ConstFloat  # noqa: F401
+from tinygrad.device import Device
+from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu
+from tinygrad.uop.spec import shared_spec
+from tinygrad.uop.symbolic import sym
+from test.helpers import get_uops
+
+def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
+  sink = UOp.group(*u)
+  for r in sink.ranges: sink = sink.end(r)
+  ret = get_uops(sink.sink(arg=KernelInfo(opts_to_apply=())), ren)
+  assert ret[-1].op is Ops.SINK
+  return ret
+
+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.IDIV, dtypes.int8, (8, 2)), 4)
+    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (7, 3)), 2)
+    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (7, -3)), -2)
+    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (-50, 6)), -8)
+
+    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 TestConstantFolding(unittest.TestCase):
+  def test_cast_const(self):
+    t = Tensor(1, dtype=dtypes.float).cast(dtypes.int)
+    si = t.schedule()
+    assert len(si) == 0
+
+class TestGatedStoreRewrite(unittest.TestCase):
+  def test_tiny_gate_store(self):
+    gmem = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
+    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)
+
+  def test_gate_some_stores(self):
+    gmem0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
+    gmem1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
+    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)
+
+  # scaled down version of TestLinearizerDumb.test_unmerged_ifs
+  @unittest.skip("we don't merge ifs anymore")
+  def test_merge_ifs_alt(self):
+    gmem0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
+    gmem1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
+    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, gate))
+    idx1 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem1, idx, 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)
+
+@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(Ops.PARAM, dt.ptr(), (), 0)
+      c = UOp.const(dt, 2)
+      l = g.index(c)
+      a = UOp(Ops.IDIV, 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.IDIV, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
+
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
+  def test_fast_idiv_and_mod(self):
+    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
+    c = UOp.const(dtypes.uint, 3)
+    l = g.index(c)
+    a = UOp(Ops.IDIV, 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.IDIV, ops)
+
+    b = UOp(Ops.MOD, 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.MOD, ops)
+
+  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(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
+    c = UOp.const(dtypes.uint, 7)
+    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
+    a = UOp(Ops.IDIV, 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.IDIV, ops)
+
+  def test_disable_fast_idiv(self):
+    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
+    c = UOp.const(dtypes.uint, 3)
+    l = g.index(c)
+    a = UOp(Ops.IDIV, 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.IDIV, 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).schedule_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(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    self.assertIs(x.replace(arg=None).arg, None)
+    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.VECTORIZE, 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(shared_spec.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.VECTORIZE, dtype=dtypes.int.vec(0), src=())
+    self.assertEqual(u.render(simplify=False), "{}")
+  def test_render_vectorize_empty_simplified(self):
+    u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(0), src=())
+    self.assertEqual(u.render(), "{}")
+  def test_render_vectorize_same(self):
+    u = UOp(Ops.VECTORIZE, 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.VECTORIZE, 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.VECTORIZE, 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.VECTORIZE, 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()

test/null/test_viz.py

@@ -353,8 +353,8 @@ def load_profile(lst:list[ProfileEvent]) -> dict:
     event_type, event_count = u("<BI")
     if event_type == 0:
       for _ in range(event_count):
-        name, ref, key, st, dur, _ = u("<IIIIfI")
-        v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur})
+        name, ref, key, st, dur, fmt = u("<IIIIfI")
+        v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur, "fmt":strings[fmt]})
     else:
       v["peak"] = u("<Q")[0]
       for _ in range(event_count):

test/null/test_winograd.py

@@ -0,0 +1,59 @@
+import unittest, sys
+from tinygrad import Tensor, GlobalCounters, dtypes, Context
+from tinygrad.helpers import CI, Profiling, 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_profile(self):
+    x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
+    with Profiling(enabled=not CI, sort='time'):
+      Tensor.conv2d(x,w).realize()
+
+  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()), 2)
+
+  def test_backward_kernels(self):
+    x,w = Tensor.empty(1,4,9,9,requires_grad=True).realize(), Tensor.empty(4,4,3,3,requires_grad=True).realize()
+    out = Tensor.conv2d(x,w, padding=1)
+    out.mean().backward()
+    backward_schedule = Tensor.schedule(x.grad, w.grad)
+    self.assertEqual(len(backward_schedule), 4)
+
+  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, 3)
+
+  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)

test/opt/test_tensor_cores.py

@@ -149,10 +149,10 @@ class TestTensorCores(unittest.TestCase):
         if _to_np_dtype(real_bufs[0].dtype) is None: continue
         real_bufs[0].copyin(np.zeros((real_bufs[0].size, ), dtype=_to_np_dtype(real_bufs[0].dtype)).data) # Zero to check that all values are filled
         prg.exec(real_bufs)
-        result = np.frombuffer(real_bufs[0].as_buffer(), _to_np_dtype(real_bufs[0].dtype))
+        result = np.frombuffer(real_bufs[0].as_memoryview(), _to_np_dtype(real_bufs[0].dtype))
 
         # ensure the results for each choice of axis matches
-        if golden_result is None: golden_result = np.frombuffer(real_bufs[0].as_buffer(), _to_np_dtype(real_bufs[0].dtype))
+        if golden_result is None: golden_result = np.frombuffer(real_bufs[0].as_memoryview(), _to_np_dtype(real_bufs[0].dtype))
         np.testing.assert_allclose(result, golden_result, atol=0.1, rtol=0.2)
 
   @Context(ALLOW_TF32=1)

test/test_arange.py

@@ -1,7 +1,7 @@
 import unittest
 import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
-from tinygrad.helpers import Context, getenv
+from tinygrad.helpers import Context, getenv, EMULATE
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.engine.schedule import ExecItem
@@ -188,5 +188,33 @@ class TestIndexing(unittest.TestCase):
     for i in idx.flatten().numpy(): expected_grad[i] += 2
     np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
 
+  # ~10x overhead in fused matmul bw with rope in bf16 vs float16
+  @unittest.skipUnless(Device.DEFAULT == "AMD" or (Device.DEFAULT == "NULL" and EMULATE.value.startswith("AMD")), "tests AMD bf16 cast overhead")
+  def base_test_llama_8b_rope_backward(self, dtype, ops_scale):
+    from extra.models.llama import precompute_freqs_cis, apply_rotary_emb
+    Tensor.training = True
+    bs, seqlen, dim, n_heads = 1, 512, 256, 4
+    head_dim = dim // n_heads
+    x = Tensor.randn(bs, seqlen, dim, dtype=dtype)
+    wq = Tensor.randn(dim, dim, dtype=dtype, requires_grad=True)
+    freqs_cis = precompute_freqs_cis(head_dim, seqlen).cast(dtype)
+    Tensor.realize(x, wq, freqs_cis)
+    xq = (x @ wq.T)
+    # main llama does not fuse it
+    #xq = xq.contiguous_backward()
+    xq = xq.reshape(bs, seqlen, n_heads, head_dim)
+    xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis)
+    xq_rope.sum().backward()
+    sched = wq.grad.schedule()
+    assert len(sched) == 1, f"expected one kernel for backward, got: {len(sched)}"
+    prg = sched[0].lower().prg.p
+    bwd_ops = prg.estimates.ops
+    expected_ops = bs*seqlen*dim*dim*ops_scale
+    print(f"rope matmul bwd ({dtype}): {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops")
+    self.assertLess(bwd_ops, expected_ops, f"rope bwd ops {bwd_ops:,} should be < {ops_scale} per (got {bwd_ops/(bs*seqlen*dim*dim):.1f})")
+
+  def test_llama_8b_rope_backward_f16(self): self.base_test_llama_8b_rope_backward(dtypes.float16, 1)
+  def test_llama_8b_rope_backward_bf16(self): self.base_test_llama_8b_rope_backward(dtypes.bfloat16, 11)
+
 if __name__ == "__main__":
   unittest.main()

test/test_const_folding.py

@@ -1,8 +1,7 @@
-import unittest, itertools, math
-from tinygrad import Tensor, Device, dtypes, Context
-from tinygrad.dtype import DType, ConstType, DTYPES_DICT
-from tinygrad.uop.ops import Ops, UOp
-from tinygrad.codegen import full_rewrite_to_sink
+import unittest, math
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.dtype import DTYPES_DICT
+from tinygrad.uop.ops import Ops
 from tinygrad.device import is_dtype_supported
 import numpy as np
 from test.helpers import not_support_multi_device
@@ -13,140 +12,6 @@ def _check_ast_count(desired_count:int, t:Tensor):
   asts = [s for s in schedule if s.ast.op is Ops.SINK]
   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_idiv_literal_one(self):
-    _check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
-  def test_idiv_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_to_sink(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):
-      r = full_rewrite_to_sink(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src[0]
-    self.assertEqual(r.op, Ops.VECTORIZE)
-    self.assertEqual(r.dtype, dtypes.uint32.vec(3))
-    self.assertEqual(tuple(x.arg for x in r.src), (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).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).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)])
-
 class TestMovedConstFolding(unittest.TestCase):
   def test_add_shrunk_zero(self):
     _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(6).shrink(((1, 5),)))

test/test_dtype.py

@@ -8,7 +8,6 @@ from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float,
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Context, Device, Tensor, dtypes
-from tinygrad.uop import Ops
 from hypothesis import given, settings, strategies as strat
 from test.helpers import rand_for_dtype
 from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
@@ -247,6 +246,17 @@ class TestBFloat16DTypeCast(unittest.TestCase):
 
 class TestHalfDType(TestDType): DTYPE = dtypes.half
 
+class TestEmulatedHalf(TestHalfDType):
+  @classmethod
+  def setUpClass(cls):
+    cls.stack = contextlib.ExitStack()
+    cls.stack.enter_context(Context(EMULATED_DTYPES="half"))
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+
+  @classmethod
+  def tearDownClass(cls): cls.stack.close()
+
+
 class TestFloatDType(TestDType):
   DTYPE = dtypes.float
 
@@ -339,7 +349,6 @@ class TestUint32DType(TestDType): DTYPE = dtypes.uint32
 
 class TestInt64DType(TestDType): DTYPE = dtypes.int64
 
-@unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
 class TestEmulatedInt64DType(TestInt64DType):
   @classmethod
@@ -356,7 +365,6 @@ class TestUint64DType(TestDType):
   def test_uint64_load(self):
     assert Tensor(2**64 - 1, dtype=dtypes.uint64).numpy() == 2**64 - 1
 
-@unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
 class TestEmulatedUInt64DType(TestUint64DType):
   @classmethod

test/test_dtype_alu.py

@@ -1,13 +1,12 @@
 import unittest, operator, math
 from tinygrad import Context, Tensor, dtypes, Device
 from tinygrad.dtype import DType, truncate
-from tinygrad.helpers import CI, getenv
+from tinygrad.helpers import CI, EMULATED_DTYPES, getenv
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
 from tinygrad.runtime.ops_python import from_storage_scalar
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.uop import Ops
 import numpy as np
 import pytest
 from hypothesis import assume, given, strategies as strat, settings
@@ -64,7 +63,10 @@ def universal_test(a, b, dtype, op):
   numpy_value = op[1](ta.numpy(), tb.numpy())
   if dtype in dtypes.fp8s: numpy_value = truncate[dtype](numpy_value.item())
   if dtype in dtypes.floats:
-    atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype, (1e-10, 1e-7))
+    if not is_dtype_supported(dtype) or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
+      fe, fm = dtypes.finfo(dtype)
+      atol, rtol = 2 ** (2 - (1 << (fe - 1))), 2 ** (-fm)
+    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype, (1e-10, 1e-7))
     np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
   else: np.testing.assert_equal(tensor_value, numpy_value)
 
@@ -117,6 +119,10 @@ class TestDTypeALU(unittest.TestCase):
   @given(ht.float16, ht.float16, strat.sampled_from(binary_operations))
   def test_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)
 
+  @given(ht.float16, ht.float16, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="half")
+  def test_emulated_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)
+
   @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
   @given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
   def test_bfloat16(self, a, b, op):
@@ -139,6 +145,10 @@ class TestDTypeALU(unittest.TestCase):
   @given(ht.float16, strat.sampled_from(unary_operations))
   def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
 
+  @given(ht.float16, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="half")
+  def test_emulated_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
+
   @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
   @given(ht.bfloat16, strat.sampled_from(unary_operations))
   def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
@@ -170,7 +180,6 @@ class TestDTypeALU(unittest.TestCase):
   @given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
   def test_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
 
-  @unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
   @given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
   @Context(EMULATED_DTYPES="long")
@@ -189,7 +198,6 @@ class TestDTypeALU(unittest.TestCase):
   @given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
   def test_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
 
-  @unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
   @given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
   @Context(EMULATED_DTYPES="long")
@@ -210,7 +218,6 @@ class TestDTypeALU(unittest.TestCase):
   @given(ht.uint64, strat.sampled_from(integer_unary_operations))
   def test_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)
 
-  @unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
   @given(ht.uint64, strat.sampled_from(integer_unary_operations))
   @Context(EMULATED_DTYPES="long")
@@ -229,7 +236,6 @@ class TestDTypeALU(unittest.TestCase):
   @given(ht.int64, strat.sampled_from(integer_unary_operations))
   def test_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64, op)
 
-  @unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "long decomp requires bitshift")
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
   @given(ht.int64, strat.sampled_from(integer_unary_operations))
   @Context(EMULATED_DTYPES="long")

test/test_graph.py

@@ -40,7 +40,7 @@ def helper_alloc_rawbuffer(device, fill=False):
   if fill:
     with Context(DEBUG=0):
       data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
-      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_buffer())
+      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_memoryview())
   return rawbuf
 
 def helper_create_offset_rawbuffer(base, offset=0):
@@ -54,7 +54,7 @@ def helper_run_jit(jis, bufs, out_buffers):
     rawbuf.copyin(mv)
 
   for ei in jis: ei.run({}, jit=True)
-  return [rawbuf.as_buffer() for rawbuf in bufs]
+  return [rawbuf.as_memoryview() for rawbuf in bufs]
 
 def helper_test_graphs(graph_impl, graphs, runs=RUN_CNT):
   reg_ji = []

test/test_image_dtype.py

@@ -12,7 +12,7 @@ class TestImageCopy(unittest.TestCase):
   def test_image_copyout_1x8(self, img_type=dtypes.imagef):
     it = Tensor.arange(32).cast(img_type((1,8,4))).realize()
     buf = it.uop.buffer
-    out = buf.as_buffer()
+    out = buf.as_memoryview()
     np.testing.assert_equal(out.cast(it.dtype.fmt).tolist(), np.arange(32))
 
   @unittest.skipUnless(is_dtype_supported(dtypes.half, device="PYTHON"), "need half")
@@ -26,14 +26,14 @@ class TestImageCopy(unittest.TestCase):
   def test_image_copyout_2x4(self):
     it = Tensor.arange(2*4*4).cast(dtypes.imagef((2,4,4))).realize()
     buf = it.uop.buffer
-    out = buf.as_buffer()
+    out = buf.as_memoryview()
     np.testing.assert_equal(out.cast('f').tolist(), np.arange(2*4*4))
 
   def test_image_roundtrip(self):
     sz = (4,2,4)
     it = Tensor.rand(prod(sz)).cast(dtypes.imagef(sz)).realize()
     buf = it.uop.buffer
-    out = buf.as_buffer()
+    out = buf.as_memoryview()
 
     it2 = Tensor.rand(prod(sz)).cast(dtypes.imagef(sz)).realize()
     buf2 = it2.uop.buffer
@@ -190,7 +190,7 @@ class TestImageDType(unittest.TestCase):
       for s in sched:
         s.run()
         if s.bufs[0].dtype == dtypes.float:
-          lst = s.bufs[0].as_buffer().cast("f").tolist()
+          lst = s.bufs[0].as_memoryview().cast("f").tolist()
           print(lst)
           assert not np.any(np.isnan(lst))
       # NOTE: the w1 grad must realize to a separate kernel

test/test_linearizer.py

@@ -504,7 +504,7 @@ def helper_linearizer_opt(r:Tensor|list[Tensor], *args, **kwargs):
   return realized_ast
 
 def copyout_outputs(outbufs:list[Buffer]) -> list[np.ndarray]:
-  return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]
+  return [np.frombuffer(x.as_memoryview(), _to_np_dtype(x.dtype)) for x in outbufs]
 
 def reset_bufs(bufs:list[Buffer]):
   for buf in bufs: buf.copyin(np.zeros((buf.size*buf.dtype.itemsize,), dtype=np.uint8).data)

test/test_multitensor.py

@@ -1248,58 +1248,6 @@ class TestTensorOps(unittest.TestCase):
   def test_bitcast(self):
     helper_test_shard_op([(256,), (256,)], lambda x: x.bitcast(dtypes.int))
 
-@unittest.skipIf(not_support_multi_device(), "no multi")
-class TestMultiRamUsage(unittest.TestCase):
-  def setUp(self):
-    self.baseline = GlobalCounters.mem_used
-    self.N = 100
-  def assertUsed(self, amt, strict=True):
-    used = GlobalCounters.mem_used - self.baseline
-    print(f"used {used} bytes")
-    if strict: self.assertEqual(used, amt)
-    else: self.assertLessEqual(used, amt)
-
-  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)
-
-  @unittest.skip("flaky")
-  def test_zeros_copy(self):
-    _ = 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)
-
-  @unittest.skip("flaky")
-  def test_zeros_shard(self, devices=(d1, d2)):
-    _ = 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((d0, d1))
-
-  @unittest.skip("flaky")
-  def test_zeros_contiguous_shard(self):
-    _ = 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_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)
-
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiFromUnrenderable(unittest.TestCase):
   @needs_second_gpu

test/test_outerworld.py

@@ -38,7 +38,7 @@ class TestOuterRange(unittest.TestCase):
     vi = UOp.variable("i", i.vmin, i.vmax).bind(i)
     out = Tensor(acc.uop.after(acc_i.store(acc_i + a[:, vi].uop).end(i)))
     out.realize()
-    assert all(x == 10.0 for x in out.tolist())
+    self.assertEqual(out.tolist(), [10.0]*10)
 
   def test_range_matmul(self):
     vec = Tensor.randn(1, 10).realize()

test/test_pickle.py

@@ -76,7 +76,7 @@ class TestPickle(unittest.TestCase):
     del a
     del buffer
     a2:UOp = pickle.loads(s)
-    self.assertListEqual(a2.base.realized.as_buffer().cast("I").tolist(), [0, 1, 2, 3])
+    self.assertListEqual(a2.base.realized.as_memoryview().cast("I").tolist(), [0, 1, 2, 3])
 
   def test_pickle_unrealized_tensor(self):
     t = Tensor.ones(10, 10)

test/test_renderer_failures.py

@@ -22,7 +22,7 @@ def _test_uop_result(inputs:list[Tensor], prg, local_size=None):
   if local_size is not None: prg = replace(prg, local_size=local_size)
   ei = CompiledRunner(prg)
   ei.exec(outbufs+inbufs)
-  return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]
+  return [np.frombuffer(x.as_memoryview(), _to_np_dtype(x.dtype)) for x in outbufs]
 
 def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
   dtype = alu_src_uops[0].dtype

test/test_schedule.py

@@ -175,6 +175,14 @@ class TestSchedule(unittest.TestCase):
     child.realize()
     assert a.uop.is_realized
 
+  def test_realize_view_of_realized_has_empty_schedule(self):
+    # views of realized buffers produce an empty schedule
+    t = Tensor.zeros((3, 3)).contiguous().realize()
+    v = t[1]  # view - is_realized but not has_buffer_identity
+    assert v.uop.is_realized
+    sched, _ = Tensor.schedule_with_vars(v)
+    self.assertEqual(len(sched), 0)
+
   # NOTE: because empty does not have a lowered ExecItem if realize is called on a childless empty, it never gets allocated.
   def test_childless_empty_never_allocates(self):
     a = Tensor.empty(10)
@@ -2079,7 +2087,7 @@ class TestCopyFolding(unittest.TestCase):
     check_schedule(b, 1, filter_sink=False) # TODO: 0?
 
   def test_copy_to_same_device_sched(self):
-    a = Tensor.ones(4).contiguous().realize().uop.as_buf()
+    a = Tensor.ones(4).contiguous().realize().uop.buf_uop
     t = Tensor(a.copy_to_device(a.device))
     sched = t.schedule()
     assert len([s for s in sched if s.ast.op is Ops.COPY]) == 0
@@ -2116,14 +2124,14 @@ class TestCopyFolding(unittest.TestCase):
     self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
 
   def test_permute_on_disk(self):
-    with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_buffer())
+    with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_memoryview())
     a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}")
     b = a.reshape(2, 2).permute(1, 0).to("CPU")
     b.realize()
     self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
 
   def test_permute_on_disk_contiguous(self):
-    with open(temp('dt_arange_4_permute_contig'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_buffer())
+    with open(temp('dt_arange_4_permute_contig'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_memoryview())
     a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute_contig')}")
     b = a.reshape(2, 2).permute(1, 0).contiguous().to("CPU")
     b.realize()
@@ -2137,118 +2145,12 @@ class TestCopyFolding(unittest.TestCase):
 
   # NOTE: disk permute must come after COPY
   def test_permute_after_shrink_on_disk(self):
-    with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).realize().uop.base.buffer.as_buffer())
+    with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).realize().uop.base.buffer.as_memoryview())
     a = Tensor.empty(5, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_5_permute')}")
     b = a.shrink(((0, 4),)).reshape(2, 2).permute(1, 0).to("CPU")
     b.realize()
     self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
 
-class TestBufferUOp(unittest.TestCase):
-  # BUFFER has a ShapeTracker of shape=(n,) and stride=(1,)
-  def test_buffer_has_buffer(self):
-    buf = Tensor.empty(10)
-    self.assertIsNotNone(buf.uop.buffer)
-    self.assertEqual(buf.uop.shape, (10,))
-    # the device Buffer remains unallocated until it's we run the schedule
-    self.assertFalse(buf.uop.buffer.is_allocated())
-    add = buf+1
-    sched = add.schedule()
-    self.assertFalse(buf.uop.buffer.is_allocated())
-    run_schedule(sched)
-    self.assertTrue(buf.uop.buffer.is_allocated())
-
-  def test_buffer_has_unique_buffer(self):
-    buf = Tensor.empty(10)
-    buf1 = buf.uop.buffer
-    buf2 = buf.uop.buffer
-    self.assertIs(buf1, buf2)
-
-  # we also allow VIEW(BUFFER) to access the underlying device Buffer, as long as it's contiguous
-  def test_buffer_view_allowed(self):
-    add = Tensor.empty(1, 1)+Tensor.empty(1, 1)
-    add.realize()
-    self.assertIsNotNone(add.uop.buffer)
-    self.assertEqual(add.uop.shape, (1, 1))
-
-  def test_buffer_view_not_allowed(self):
-    permuted_view = Tensor.empty(1, 2, 3).permute(0, 2, 1)
-    with self.assertRaisesRegex(AssertionError, "can only be RESHAPE"):
-      permuted_view.uop.buffer # cannot access Buffer of a non contiguous VIEW
-
-  def test_buffer_only_after_realize(self):
-    a = Tensor([1])+Tensor([2])
-    # accessing realized will return None
-    self.assertIsNone(a.uop.realized)
-    # accessing Buffer will assert
-    with self.assertRaisesRegex(AssertionError, "must be BUFFER"):
-      a.uop.buffer # there is no BUFFER on an unrealized ADD
-    # Buffer only exists once we realize it
-    a.realize()
-    self.assertIsNotNone(a.uop.buffer)
-
-  def test_const_does_not_realize(self):
-    a = Tensor(1)+Tensor(2)
-    run_schedule(check_schedule(a, 0))
-    self.assertIsNone(a.uop.base.realized)
-
-  def test_var_does_not_realize(self):
-    a = Tensor(UOp.variable("a", 0, 10).bind(1))
-    run_schedule(check_schedule(a, 0))
-    self.assertIsNone(a.uop.base.realized)
-
-  def test_unused_var_not_in_var_vals(self):
-    # unused variable should not appear in var_vals even when there's other work
-    a = Tensor(UOp.variable("unused", 0, 10).bind(1))
-    b = Tensor.empty(3) + 1
-    _, var_vals = Tensor.schedule_with_vars(a, b)
-    self.assertEqual(var_vals, {})
-    self.assertIsNone(a.uop.base.realized)
-
-  def test_view_does_not_realize(self):
-    a = Tensor.randn(1, 4).expand(4, 4)
-    a.realize()
-    self.assertEqual(a.uop.base.realized.size, 4)
-    a2 = a.contiguous().realize()
-    self.assertEqual(a2.uop.base.realized.size, 16)
-
-class TestContiguous(unittest.TestCase):
-  def test_contiguous_buffer(self):
-    a = Tensor.empty(4)
-    b = a.contiguous()
-    check_schedule(b, 0)
-
-  def test_contiguous_buffer_view(self):
-    a = Tensor.empty(4)
-    b = a.reshape((2, 2)).contiguous()
-    check_schedule(b, 0)
-
-  def test_non_contiguous_buffer_view(self):
-    a = Tensor.empty(4, 1)
-    b = a.expand((4, 4)).contiguous()
-    check_schedule(b, 1)
-
-  def test_size_change_buffer_view(self):
-    a = Tensor.empty(4)
-    b = a.reshape((1, 1, 4)).shrink(((0, 1), (0, 1), (0, 3))).contiguous()
-    check_schedule(b, 1)
-
-  def test_double_contiguous_realizes_once(self):
-    a = Tensor.empty(4, 1)
-    b = a.expand((4, 4)).contiguous().contiguous()
-    check_schedule(b, 1)
-
-  def test_view_does_not_realize(self):
-    a = Tensor.empty(4)
-    b = a.expand((4, 4))
-    check_schedule(b, 0)
-    self.assertEqual(b.uop.base.buffer.size, 4)
-
-  def test_contiguous_view_realizes(self):
-    a = Tensor.empty(4)
-    b = a.expand((4, 4)).contiguous()
-    check_schedule(b, 1)
-    self.assertEqual(b.uop.base.buffer.size, 16)
-
 class TestUOpBecome(unittest.TestCase):
   # the simplest case, if we create a new BUFFER for this tensor UOp
   def test_new_buffer(self):
@@ -2388,12 +2290,5 @@ class TestUOpBecome(unittest.TestCase):
     b.shrink(((0,4),)).assign(a_view).realize()
     self.assertListEqual(b.tolist(), [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
 
-class TestSimpleSchedule(unittest.TestCase):
-  def test_reduce_doesnt_split(self):
-    a = Tensor.empty(16,16).sum(axis=1)
-    a1 = a.reshape(4,4)
-    a2 = a.reshape(16,1,1)
-    self.assertEqual(len(Tensor.schedule(a1, a2)), 1)
-
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/test_setitem.py

@@ -54,7 +54,32 @@ class TestSetitem(unittest.TestCase):
     t = Tensor.ones(4)
     with self.assertRaises(RuntimeError): t[1] = 5
 
-  @unittest.skip("TODO: flaky")
+  def test_setitem_chained_indexing(self):
+    # N[i][j] must work the same as N[i, j]
+    N1 = Tensor.zeros((3, 3)).contiguous().realize()
+    N1[1, 2] = 5
+    N2 = Tensor.zeros((3, 3)).contiguous().realize()
+    N2[1][2] = 5
+    np.testing.assert_equal(N1.numpy(), N2.numpy())
+
+  def test_setitem_detach(self):
+    # setitem on detached tensor should work
+    t = Tensor.zeros((3, 3)).contiguous().realize()
+    t.detach()[1, 2] = 5
+    self.assertEqual(t[1, 2].item(), 5.0)
+
+  def test_setitem_permute(self):
+    # setitem on permuted tensor should modify original
+    t = Tensor.zeros((2, 3)).contiguous().realize()
+    t.T[1, 0] = 5  # t.T is (3, 2), so [1, 0] maps to t[0, 1]
+    self.assertEqual(t[0, 1].item(), 5.0)
+
+  def test_setitem_flip(self):
+    # setitem on flipped tensor should modify original
+    t = Tensor.zeros((3,)).contiguous().realize()
+    t[::-1][0] = 5  # flip, then set first element (which is last in original)
+    self.assertEqual(t[2].item(), 5.0)
+
   def test_setitem_inplace_operator(self):
     t = Tensor.arange(4).reshape(2, 2).contiguous()
     t[1] += 2

test/test_subbuffer.py

@@ -13,26 +13,26 @@ class TestSubBuffer(unittest.TestCase):
 
   def test_subbuffer(self):
     vbuf = self.buf.view(2, dtypes.uint8, offset=3).ensure_allocated()
-    tst = vbuf.as_buffer().tolist()
+    tst = vbuf.as_memoryview().tolist()
     assert tst == [3, 4]
 
   def test_subbuffer_cast(self):
     # NOTE: bitcast depends on endianness
     vbuf = self.buf.view(2, dtypes.uint16, offset=3).ensure_allocated()
-    tst = vbuf.as_buffer().cast("H").tolist()
+    tst = vbuf.as_memoryview().cast("H").tolist()
     assert tst == [3|(4<<8), 5|(6<<8)]
 
   def test_subbuffer_double(self):
     vbuf = self.buf.view(4, dtypes.uint8, offset=3).ensure_allocated()
     vvbuf = vbuf.view(2, dtypes.uint8, offset=1).ensure_allocated()
-    tst = vvbuf.as_buffer().tolist()
+    tst = vvbuf.as_memoryview().tolist()
     assert tst == [4, 5]
 
   def test_subbuffer_len(self):
     vbuf = self.buf.view(5, dtypes.uint8, 2).ensure_allocated()
-    mv = vbuf.as_buffer()
+    mv = vbuf.as_memoryview()
     assert len(mv) == 5
-    mv = vbuf.as_buffer(allow_zero_copy=True)
+    mv = vbuf.as_memoryview(allow_zero_copy=True)
     assert len(mv) == 5
 
   def test_subbuffer_used(self):
@@ -63,7 +63,7 @@ class TestSubBuffer(unittest.TestCase):
 
       vbuf.ensure_allocated()
 
-      tst = vbuf.as_buffer().tolist()
+      tst = vbuf.as_memoryview().tolist()
       assert tst == [13, 14]
 
   def test_subbuffer_is_allocated(self):
@@ -112,8 +112,8 @@ class TestSubBuffer(unittest.TestCase):
     sub_buf.copyout(memoryview(data_out_sub))
     assert data_out_sub == bytearray(range(3, 6))
     sub_buf.copyin(memoryview(bytearray(range(3))))
-    assert sub_buf.as_buffer().tolist() == list(range(3))
-    assert self.buf.as_buffer().tolist()[3:6] == list(range(3))
+    assert sub_buf.as_memoryview().tolist() == list(range(3))
+    assert self.buf.as_memoryview().tolist()[3:6] == list(range(3))
     sub_buf.copyout(memoryview(data_out_sub))
     assert data_out_sub == bytearray(range(3))
     data_out_base = bytearray([0]*10)
@@ -145,17 +145,17 @@ class TestSubBuffer(unittest.TestCase):
     sub_buf = self.buf.view(4, dtypes.int8, offset=3)
     sub_buf.allocate()
     sub_buf.copyin(memoryview(bytearray(range(10, 14))))
-    assert self.buf.as_buffer().tolist()[3:7] == sub_buf.as_buffer().tolist()
+    assert self.buf.as_memoryview().tolist()[3:7] == sub_buf.as_memoryview().tolist()
 
     sub_buf = self.buf_unalloc.view(4, dtypes.int8, offset=3)
     sub_buf.allocate()
     sub_buf.copyin(memoryview(bytearray(range(10, 14))))
-    assert self.buf_unalloc.as_buffer().tolist()[3:7] == sub_buf.as_buffer().tolist()
+    assert self.buf_unalloc.as_memoryview().tolist()[3:7] == sub_buf.as_memoryview().tolist()
 
   def test_subbuffer_dealloc(self):
     sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
     sub_buf.deallocate()
-    assert self.buf.as_buffer().tolist() == list(range(10))
+    assert self.buf.as_memoryview().tolist() == list(range(10))
 
   def test_subbuffer_double_dealloc(self):
     sub_buf = self.buf.view(3, dtypes.uint8, offset=4).ensure_allocated()
@@ -168,17 +168,17 @@ class TestSubBuffer(unittest.TestCase):
 
   def test_subbuffer_uaf(self):
     sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
-    assert self.buf.as_buffer().tolist(), list(range(10))
+    assert self.buf.as_memoryview().tolist(), list(range(10))
     sub_buf.deallocate()
     with self.assertRaises(AssertionError):
-      sub_buf.as_buffer().tolist()
-    assert self.buf.as_buffer().tolist(), list(range(10))
+      sub_buf.as_memoryview().tolist()
+    assert self.buf.as_memoryview().tolist(), list(range(10))
 
     sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
-    assert sub_buf.as_buffer().tolist(), list(range(3, 7))
+    assert sub_buf.as_memoryview().tolist(), list(range(3, 7))
     self.buf.deallocate()
     with self.assertRaises(AssertionError):
-      sub_buf.as_buffer().tolist()
+      sub_buf.as_memoryview().tolist()
 
 if __name__ == '__main__':
   unittest.main()

test/test_tensor.py

@@ -6,11 +6,7 @@ from tinygrad.helpers import getenv, temp, mv_address
 from extra.gradcheck import numerical_jacobian, jacobian, gradcheck
 from hypothesis import given, settings, strategies as strat
 from tinygrad.device import is_dtype_supported
-from tinygrad.uop.ops import Ops, UOp
-from tinygrad.renderer.ptx import PTXRenderer
-from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.engine.realize import get_program
-from tinygrad.dtype import DType, DTYPES_DICT
+from tinygrad.dtype import DTYPES_DICT
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -752,154 +748,5 @@ class TestTensorCreationDevice(unittest.TestCase):
     x = y.one_hot(10)
     x.realize()
 
-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, requires_grad=True)
-    m = Tensor(m_init, requires_grad=True)
-    W = Tensor(W_init, requires_grad=True)
-    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
-    assert W.requires_grad
-
-  def test_no_grad_mode_context_manager(self):
-    x = Tensor(x_init, requires_grad=True)
-    m = Tensor(m_init, requires_grad=True)
-    W = Tensor(W_init, requires_grad=True)
-    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):
-    schedule, _ = a.schedule_with_vars()
-    for s in schedule:
-      if s.ast.op is Ops.SINK:
-        renderer = Device[s.bufs[0].device].renderer
-        prg = get_program(s.ast, renderer)
-        return prg.uops
-
-  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]
-      assert idx_val.dtype is 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(is_dtype_supported(dtypes.long), "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(is_dtype_supported(dtypes.long), "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, 2048, 2048, 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)
-
-  def test_arange_raise_overflow(self):
-    with self.assertRaises(ValueError):
-      self._schedule_render(Tensor.arange(2**33, dtype=dtypes.int))
-
-  @unittest.skipIf(is_dtype_supported(dtypes.long), "int64 is supported")
-  def test_int64_unsupported_overflow_sym(self):
-    with self.assertRaises(KeyError):
-      self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 1, 2048).bind(32))
-
-  @unittest.skipIf(is_dtype_supported(dtypes.long), "int64 is supported")
-  @unittest.expectedFailure  # bug in gpu dims limiting
-  def test_int64_unsupported_overflow(self):
-    with self.assertRaises(KeyError):
-      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_eye_bufs_unique(self):
-    a = Tensor.eye(10).contiguous()
-    b = Tensor.eye(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)
-
 if __name__ == '__main__':
   unittest.main()

test/test_uops.py

@@ -2,15 +2,13 @@ from typing import Optional, Any
 import unittest, math
 import numpy as np
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.helpers import CI, DEBUG, getenv, Timing, Context
+from tinygrad.helpers import CI, getenv, Context
 from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat  # noqa: F401
 from tinygrad.device import Buffer, Device
-from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu, AxisType
-from tinygrad.uop.spec import shared_spec
+from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.engine.realize import CompiledRunner, get_program, get_runner
 from tinygrad.engine.schedule import ExecItem
-from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer
@@ -181,153 +179,6 @@ class TestBoolUOps(TestUOps):
   def test_cmplt_bool(self): self._test_bop_bool_fxn(Ops.CMPLT, lambda a,b: a < b)
   def test_where_bool(self): self._test_top_bool_fxn(Ops.WHERE, lambda a,b,c: b if a else c)
 
-class TestSafeCast(TestUOps):
-  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(TestUOps):
-  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.IDIV, dtypes.int8, (8, 2)), 4)
-    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (7, 3)), 2)
-    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (7, -3)), -2)
-    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (-50, 6)), -8)
-
-    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 TestConstantFolding(unittest.TestCase):
-  def test_cast_const(self):
-    t = Tensor(1, dtype=dtypes.float).cast(dtypes.int)
-    si = t.schedule()
-    assert len(si) == 0
-
-class TestGatedStoreRewrite(unittest.TestCase):
-  def test_tiny_gate_store(self):
-    gmem = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
-    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 DEBUG >= 4: print(Device[Device.DEFAULT].renderer.render(uops))
-    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)
-
-  def test_gate_some_stores(self):
-    gmem0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
-    gmem1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
-    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 DEBUG >= 4: print(Device[Device.DEFAULT].renderer.render(uops))
-    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)
-
-  # scaled down version of TestLinearizerDumb.test_unmerged_ifs
-  @unittest.skip("we don't merge ifs anymore")
-  def test_merge_ifs_alt(self):
-    gmem0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
-    gmem1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
-    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, gate))
-    idx1 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem1, idx, gate))
-    val = UOp.const(dtypes.float, 42.0)
-    stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
-    uops = to_uops_list(stores)
-    if DEBUG >= 4: print(Device[Device.DEFAULT].renderer.render(uops))
-    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)
-
 class TestLocalAccess(unittest.TestCase):
   # NOTE: this is failing on METAL CI, no idea why. Works locally.
   @unittest.skipIf(Device.DEFAULT == "METAL" and CI, "failing only in CI")
@@ -375,73 +226,6 @@ class TestLocalAccess(unittest.TestCase):
     sres = uop(uops, Ops.LOAD, dtypes.int32, (smem.index(ofs),))
     self.assertEqual(_test_uops_result(dtypes.int32, uops, sres), 42)
 
-@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(Ops.PARAM, dt.ptr(), (), 0)
-      c = UOp.const(dt, 2)
-      l = g.index(c)
-      a = UOp(Ops.IDIV, 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.IDIV, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
-
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
-  def test_fast_idiv_and_mod(self):
-    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
-    c = UOp.const(dtypes.uint, 3)
-    l = g.index(c)
-    a = UOp(Ops.IDIV, 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.IDIV, ops)
-
-    b = UOp(Ops.MOD, 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.MOD, ops)
-
-  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(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
-    c = UOp.const(dtypes.uint, 7)
-    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
-    a = UOp(Ops.IDIV, 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.IDIV, ops)
-
-  def test_disable_fast_idiv(self):
-    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
-    c = UOp.const(dtypes.uint, 3)
-    l = g.index(c)
-    a = UOp(Ops.IDIV, 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.IDIV, ops)
-
-
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends")
 class TestAssembly(unittest.TestCase):
   def test_bitshift_left(self):
@@ -480,118 +264,6 @@ class TestAssembly(unittest.TestCase):
     self.assertIn(Ops.CMPEQ, ops)
     self.assertNotIn(Ops.CMPNE, 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).schedule_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(Ops.PARAM, dtypes.int.ptr(), (), 0)
-    self.assertIs(x.replace(arg=None).arg, None)
-    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.VECTORIZE, 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(shared_spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
-    test_upat = UPat(Ops.CONST, dtypes.bool)
-    self.assertEqual(test_upat.location[0].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):
-  # LOL, running this test breaks all instances of "4"
-  """
-  @unittest.expectedFailure
-  def test_immutable(self):
-    const_4 = UOp.const(dtypes.int, 4)
-    with self.assertRaises(Exception):
-      const_4.arg = 5
-  """
-
-  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.VECTORIZE, dtype=dtypes.int.vec(0), src=())
-    self.assertEqual(u.render(simplify=False), "{}")
-  def test_render_vectorize_empty_simplified(self):
-    u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(0), src=())
-    self.assertEqual(u.render(), "{}")
-  def test_render_vectorize_same(self):
-    u = UOp(Ops.VECTORIZE, 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.VECTORIZE, 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.VECTORIZE, 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.VECTORIZE, 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)")
-
 class TestZeroRange(unittest.TestCase):
   def test_reduce_variable(self):
     for i in range(3,-1,-1):

test/test_zero_copy.py

@@ -6,7 +6,7 @@ def time_tensor_numpy(out:Tensor):
   times = []
   for _ in range(5):
     st = time.perf_counter()
-    out.uop.base.realized.as_buffer(allow_zero_copy=True)
+    out.uop.base.realized.as_memoryview(allow_zero_copy=True)
     et = time.perf_counter() - st
     times.append(et)
   return min(times)

test/testextra/test_asm_gemm.py

@@ -0,0 +1,59 @@
+import unittest
+from tinygrad import Tensor, Device, dtypes, Context
+from tinygrad.device import is_dtype_supported
+from tinygrad.helpers import getenv, CI
+from extra.gemm.asm.cdna.gemm import asm_gemm
+
+def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
+  Tensor.manual_seed(0)
+  a_rand = Tensor.randn((batch, M, K), dtype=dtypes.float).sub(0.5).cast(dtype)
+  b_rand = Tensor.randn((K, N), dtype=dtypes.float).sub(0.5).cast(dtype)
+  with Context(DEBUG=0):
+    Tensor.realize(a_rand, b_rand)
+
+  devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(gpus)) if (multi:=gpus>1) else None
+
+  a, b = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
+  if multi: a, b = a.shard(devs, axis=0), b.shard(devs, axis=None)
+  tst = asm_gemm(a, b)
+  tst.sum().backward()
+  Tensor.realize(tst, a.grad, b.grad)
+
+  a_ref, b_ref = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
+  if multi: a_ref, b_ref = a_ref.shard(devs, axis=0), b_ref.shard(devs, axis=None)
+  with Context(ASM_GEMM=0): ref = a_ref @ b_ref
+  ref.sum().backward()
+  Tensor.realize(ref, a_ref.grad, b_ref.grad)
+
+  with Context(DEBUG=0):
+    assert (tst - ref).square().max().float().item() < 1e-6, "forward mismatch"
+    assert (a.grad - a_ref.grad).square().max().float().item() < 1e-3, "grad_a mismatch"
+    assert (b.grad - b_ref.grad).square().max().float().item() < 1e-3, "grad_b mismatch"
+
+SCALE = 128 if CI else 1
+
+@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
+class TestGemm(unittest.TestCase):
+  def test_simple(self): verify_asm_gemm(1, N:=(getenv("N", 4096)//SCALE), N, N, dtype=dtypes.half)
+  def test_gemm(self): verify_asm_gemm(1, 8192//SCALE, 4096//SCALE, 14336//SCALE)
+  def test_gemm_batched(self): verify_asm_gemm(2, 8192//SCALE, 4096//SCALE, 4096//SCALE)
+  def test_gemm_multi(self): verify_asm_gemm(2, 8192//SCALE, 4096//SCALE, 4096//SCALE, gpus=2)
+
+class TestGemmLarge(unittest.TestCase):
+  def setUp(self):
+    if getattr(Device[Device.DEFAULT].renderer, "arch", "") != "gfx950":
+      self.skipTest("very slow on non mi350x")
+
+  def test_gemm1(self): verify_asm_gemm(8, 8192, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
+  def test_gemm2(self): verify_asm_gemm(8, 8192, 128256, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_gemm3(self): verify_asm_gemm(8, 8192, 14336, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
+  def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
+  def test_gemm7(self): verify_asm_gemm(1, 8192, 128256, 4096)
+  def test_gemm_unsupported(self):
+    with self.assertRaisesRegex(AssertionError, "shape not supported"):
+      verify_asm_gemm(8, 1024, 1024, 4096, gpus=8)
+
+if __name__ == "__main__":
+  unittest.main()

test/unit/test_assign.py

@@ -131,7 +131,7 @@ class TestAssign(unittest.TestCase):
 
   @unittest.skip("assign to contiguous shouldn't change the base buffer")
   def test_assign_changes_buffer_alt(self):
-    a, b = [Tensor(Tensor(0).contiguous().realize().uop.as_buf()) for _ in range(2)]
+    a, b = [Tensor(Tensor(0).contiguous().realize().uop.buf_uop) for _ in range(2)]
     Tensor.realize(a.contiguous().assign(1), b.contiguous().assign(2))
     self.assertEqual((a + b).item(), 3)
 

test/unit/test_attention.py

@@ -1,7 +1,6 @@
 import unittest
-from tinygrad import Tensor, dtypes, TinyJit, UOp
+from tinygrad import Tensor, dtypes
 from tinygrad.apps.llm import apply_rope as apply_rope_new, precompute_freqs_cis
-#from tinygrad.engine.realize import run_schedule
 
 def apply_rope(x:Tensor, start_pos:int):
   B, H, T, Hd = x.shape
@@ -9,21 +8,7 @@ def apply_rope(x:Tensor, start_pos:int):
   freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
   return apply_rope_new(x, freqs_cis)
 
-# TODO: test_scheduler, but just in uint
 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()
-    # attention has 4 kernels now
-    self.assertEqual(len(sched), 4)
-    # softmax_inputs = sched[1:4]
-    # for i,si in enumerate(softmax_inputs):
-    #   assert all(b.dtype == dtypes.half for b in si.bufs), f"non half {si.bufs=} in kernel {i}"
-
   def test_apply_rope(self):
     x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32)
     result = apply_rope(x, 0)
@@ -32,21 +17,5 @@ class TestAttention(unittest.TestCase):
     self.assertGreater((result - apply_rope(x, 5)).abs().max().item(), 1e-6)
     with self.assertRaises(AssertionError): apply_rope(Tensor.randn(1, 1, 4, 7, dtype=dtypes.float32), 0)
 
-  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))
-    noprune_size = len(rope_noprune.captured.jit_cache)
-    prune_size = len(rope_prune.captured.jit_cache)
-
-    self.assertGreater(noprune_size, prune_size)
-    self.assertGreaterEqual(noprune_size, 2)
-    self.assertEqual(prune_size, 1)
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_dtype_spec.py

@@ -1,6 +1,6 @@
-import unittest, math, operator, subprocess, struct
+import unittest, math, subprocess
 from tinygrad.tensor import Tensor, dtypes, Device
-from tinygrad.dtype import DType, DTYPES_DICT, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
+from tinygrad.dtype import DType, DTYPES_DICT
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, DEBUG
 from test.helpers import slow
@@ -28,195 +28,12 @@ def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float
   except AssertionError as e:
     raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e
 
-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(dtypes.min(dt), -math.inf)
-        np.testing.assert_equal(dtypes.max(dt), 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(dtypes.min(dt), info.min)
-        np.testing.assert_equal(dtypes.max(dt), 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(dtypes.min(dt), False)
-        np.testing.assert_equal(dtypes.max(dt), 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):
-    # TODO: fuzz this better
-    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):
-    # In f32, NaN = exp 0xFF and mantissa ≠ 0. Quiet-vs-signaling is bit 22 of the mantissa: 1 = qNaN, 0 = sNaN.
-    # qNaN(+/-), sNaN(+/-) overflow(+/-)
-    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):
-    # round_to_nearest_even
-    uppers = [0x3f800000, 0x41230000, 0xC1460000] # 1.0, 10.1875, -12.375
-    for upper in uppers:
-      base = upper & 0xFFFF0000
-      base_f32 = u32_to_f32(base)
-      base_f32_round_up = u32_to_f32(base + 0x00010000)
-
-      # low < 0x8000(0.5ULP) -> round down
-      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)
-
-      # low > 0x8000(0.5ULP) -> round up
-      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)
-
-      # low == 0x8000(0.5ULP) and LSB even -> round down
-      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)
-      # low == 0x8000(0.5ULP) and LSB odd -> round up
-      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):
-    # bf16 max finite: exp=0xFE, faction=0x7F => 0x7F7F0000(f32)
-    # bf16 inf(+/-):   exp=0xFF
-    base = 0x7F7F0000
-    inf_u32 = 0x7F800000
-
-    # low < 0.5ULP
-    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)
-
-    # low > 0.5ULP -> overflows to +inf
-    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)
-
-    # low == 0.5ULP and LSB odd -> overflows to +inf
-    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())
-
 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
-
   @unittest.skip("this test is slow and spawning whole pythons")
   def test_env_set_default_float(self):
     # check default
@@ -323,97 +140,12 @@ class TestTypeSpec(unittest.TestCase):
     _assert_eq(Tensor.arange(3, 5, -2), dtypes.default_int, np.arange(3, 5, -2))
     _assert_eq(Tensor.arange(5.0, 3.0), dtypes.default_float, np.arange(5.0, 3.0))
 
-  @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 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
-    # similar to jax but we don't use weak type
-    assert least_upper_dtype(dtypes.int64, dtypes.uint64) == dtypes.uint64  # is this correct?
-    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
-
 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([d for d in core_dtypes if dtypes.is_int(d) and is_dtype_supported(d)]))
   def test_int_to_float_unary_func(self, dtype):
     for func in [
@@ -432,34 +164,6 @@ class TestAutoCastType(unittest.TestCase):
       # float16 can have larger precision errors
       np.testing.assert_allclose(func(Tensor(a, dtype=dtype)).numpy(), func(torch.tensor(a)), rtol=1e-3, atol=1e-3)
 
-  @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
-
   @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16")
   def test_sum_dtype_arg(self):
     t = Tensor([40000, 40000], dtype=dtypes.float16)
@@ -477,105 +181,6 @@ class TestAutoCastType(unittest.TestCase):
     assert t.prod(dtype=dtypes.float32).dtype == dtypes.float32
     np.testing.assert_allclose(t.prod(dtype=dtypes.float32).numpy(), 20000)
 
-  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
-
-  @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))
-    # if dtype is specified, return in 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))
-    # if dtype is specified, return in 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
-
   def test_gradient_dtype(self):
     old_default_float = dtypes.default_float
 

test/unit/test_gradient.py

@@ -1,67 +1,7 @@
-from typing import Callable
-import unittest, math
-import torch
+import unittest
 import numpy as np
 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 TestTensorGradient(unittest.TestCase):
   def test_example(self):
@@ -128,36 +68,13 @@ class TestTensorGradient(unittest.TestCase):
     np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0+2*3.0])
     self.assertIs(x.grad, old_grad)
 
-class TestRealizeMeansRealize(unittest.TestCase):
-  def test_randn_realizes(self):
-    x = Tensor.randn(2, 3, 64, 64, requires_grad=True).realize()
-    assert x.uop is not x.uop.base
-    assert x.uop.is_realized
-
-  #@unittest.expectedFailure
-  # update: passing after delete_forced_realize
-  def test_uniform_realizes(self):
-    x = Tensor.uniform(16, 3, 3, 3, requires_grad=True).realize()
-    print(x.uop)
-    assert x.uop is not x.uop.base
-    assert x.uop.is_realized
-
-  # NOTE: even though it doesn't realize, this seems fine
-  def test_uniform_gradient(self):
-    x = Tensor.uniform(16, 3, 3, 3, requires_grad=True).realize()
-    y = x * 2
-    y.sum().gradient(x)[0].realize()
-
 class TestViewGradient(unittest.TestCase):
   def test_expand(self):
-    # this test shows that if Tensors collapse to the views and create a disconnected graph
-    # there's no way to recover the proper gradient
     x = Tensor.randn(5,2)
     a = Tensor([3.], requires_grad=True)
     aex = a.expand(10)
     (aex.reshape(5,2) * x).sum().backward()
     np.testing.assert_allclose(aex.grad.numpy(), x.reshape(10).numpy())
-    # NOTE: aex.grad is *not* a.grad.expand(10)!
     with self.assertRaises(AssertionError):
       np.testing.assert_allclose(aex.grad.numpy(), a.grad.expand(10).numpy())
 

test/unit/test_helpers.py

@@ -1,473 +1,13 @@
-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
-from tinygrad.tensor import Tensor, get_shape
+import unittest
 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}"
-
-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]])
+from tinygrad.helpers import polyN, is_numpy_ndarray
+from tinygrad.tensor import Tensor
 
 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_tensor(self):
-    from tinygrad.tensor import Tensor
     np.testing.assert_allclose(polyN(Tensor([1.0, 2.0, 3.0, 4.0]), [1.0, -2.0, 1.0]).numpy(), [0.0, 1.0, 4.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_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]))))
   def test_tensor_numpy(self):
     self.assertTrue(is_numpy_ndarray(Tensor([1, 2, 3]).numpy()))
 

test/unit/test_indexing.py

@@ -193,22 +193,6 @@ class TestIndexing(unittest.TestCase):
     numpy_testing_assert_equal_helper(w.grad, m * a)
   '''
 
-  def test_single_int(self):
-    v = Tensor.randn(5, 7, 3)
-    numpy_testing_assert_equal_helper(v[4].shape, (7, 3))
-
-  def test_multiple_int(self):
-    v = Tensor.randn(5, 7, 3)
-    numpy_testing_assert_equal_helper(v[4].shape, (7, 3))
-    numpy_testing_assert_equal_helper(v[4, :, 1].shape, (7,))
-
-  def test_none(self):
-    v = Tensor.randn(5, 7, 3)
-    numpy_testing_assert_equal_helper(v[None].shape, (1, 5, 7, 3))
-    numpy_testing_assert_equal_helper(v[:, None].shape, (5, 1, 7, 3))
-    numpy_testing_assert_equal_helper(v[:, None, None].shape, (5, 1, 1, 7, 3))
-    numpy_testing_assert_equal_helper(v[..., None].shape, (5, 7, 3, 1))
-
   def test_step(self):
     v = Tensor.arange(10)
     numpy_testing_assert_equal_helper(v[::1], v)
@@ -418,34 +402,6 @@ class TestIndexing(unittest.TestCase):
     out = scripted_fn2(clone(data))
     numpy_testing_assert_equal_helper(out, ref)
 
-  def test_int_indices(self):
-    v = Tensor.randn(5, 7, 3)
-    numpy_testing_assert_equal_helper(v[[0, 4, 2]].shape, (3, 7, 3))
-    numpy_testing_assert_equal_helper(v[:, [0, 4, 2]].shape, (5, 3, 3))
-    numpy_testing_assert_equal_helper(v[:, [[0, 1], [4, 3]]].shape, (5, 2, 2, 3))
-
-  # TODO fancy setitem
-  '''
-  def test_index_put_src_datatype(self, dtype):
-    src = Tensor.ones(3, 2, 4, dtype=dtype)
-    vals = Tensor.ones(3, 2, 4, dtype=dtype)
-    indices = (np.array([0, 2, 1]),)
-    res = index_put_(src, indices, vals, accumulate=True)
-    numpy_testing_assert_equal_helper(res.shape, src.shape)
-  '''
-
-  def test_index_src_datatype(self):
-    src = Tensor.ones(3, 2, 4)
-    # test index
-    res = src[[0, 2, 1], :, :]
-    numpy_testing_assert_equal_helper(res.shape, src.shape)
-    # test index_put, no accum
-    # TODO fancy setitem
-    '''
-    src[[0, 2, 1], :, :] = res
-    numpy_testing_assert_equal_helper(res.shape, src.shape)
-    '''
-
   def test_int_indices2d(self):
     # From the NumPy indexing example
     x = Tensor.arange(0, 12).reshape(4, 3)
@@ -496,12 +452,6 @@ class TestIndexing(unittest.TestCase):
     with self.assertRaises(IndexError):
       x[:, [0, 1]]
 
-  def test_empty_slice(self):
-    x = Tensor.randn(2, 3, 4, 5)
-    y = x[:, :, :, 1]
-    z = y[:, 1:1, :]
-    numpy_testing_assert_equal_helper((2, 0, 4), z.shape)
-
   @unittest.skip("bool indexing not supported")
   def test_index_getitem_copy_bools_slices(self):
     true = Tensor(1, dtype=dtypes.uint8)
@@ -738,17 +688,6 @@ class TestIndexing(unittest.TestCase):
         (2, -3))
   '''
 
-  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])
-
   def test_zero_dim_index(self):
     x = Tensor(10)
     numpy_testing_assert_equal_helper(x, x.item())
@@ -773,41 +712,6 @@ class TestIndexing(unittest.TestCase):
     numpy_testing_assert_equal_helper(out, Tensor.zeros(2))
   '''
 
-  def test_take_along_dim(self):
-    def _test_against_numpy(t: Tensor, indices: Tensor, dim):
-      actual = t.gather(dim, indices)
-      t_np = t.numpy()
-      indices_np = indices.numpy()
-      expected = np.take_along_axis(t_np, indices_np, axis=dim)
-      numpy_testing_assert_equal_helper(actual, expected)
-
-      # TODO argsort
-      '''
-      for shape in [(3, 2), (2, 3, 5), (2, 4, 0), (2, 3, 1, 4)]:
-        for noncontiguous in [True, False]:
-          for dtype in (dtypes.float32, dtypes.int64):
-            t = make_tensor(shape, dtype=dtype, noncontiguous=noncontiguous)
-            for dim in list(range(t.ndim)) + [None]:
-              if dim is None:
-                indices = argsort(t.reshape(-1))
-              else:
-                indices = argsort(t, dim=dim)
-
-          _test_against_numpy(t, indices, dim)
-      '''
-
-      # test broadcasting
-      t = Tensor.ones((3, 4, 1))
-      indices = Tensor.ones((1, 2, 5), dtype=dtypes.int64)
-
-      _test_against_numpy(t, indices, 1)
-
-      # test empty indices
-      t = Tensor.ones((3, 4, 5))
-      indices = Tensor.ones((3, 0, 5), dtype=dtypes.int64)
-
-      _test_against_numpy(t, indices, 1)
-
   # TODO argsort
   '''
   def test_take_along_dim_invalid(self):
@@ -840,38 +744,6 @@ class TestIndexing(unittest.TestCase):
   '''
 
 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])
-    numpy_testing_assert_equal_helper(a[None].ndim, a.ndim+1)
-
   def test_empty_tuple_index(self):
     # Empty tuple index creates a view
     a = Tensor([1, 2, 3])
@@ -1034,19 +906,6 @@ class TestNumpy(unittest.TestCase):
     numpy_testing_assert_equal_helper(a[c], Tensor([[1, 2, 3], [4, 5, 6]]))
     numpy_testing_assert_equal_helper(a[c, c], Tensor([1, 5]))
 
-  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)
-
   # TODO out of bound getitem does not raise error
   '''
   def test_trivial_fancy_out_of_bounds(self):

test/unit/test_llm_server.py

@@ -1,5 +1,5 @@
-import unittest, threading, time
-from unittest.mock import Mock, patch
+import unittest
+from unittest.mock import patch
 from tinygrad import Tensor
 
 class TestTransformerGenerate(unittest.TestCase):
@@ -25,137 +25,5 @@ class TestTransformerGenerate(unittest.TestCase):
     self.assertEqual(captured_inputs[0][0][-1], 2)  # shape should be (1, 2)
     self.assertEqual(captured_inputs[0][1], 3)  # start_pos should be 3, not 0
 
-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.end_turn = Mock(return_value=[998])
-
-    cls.mock_model = Mock()
-    cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
-
-    cls.bos_id = 1
-    cls.eos_id = 999
-
-    import tinygrad.apps.llm as llm_module
-    llm_module.model = cls.mock_model
-    llm_module.tok = cls.mock_tok
-    llm_module.bos_id = cls.bos_id
-    llm_module.eos_id = cls.eos_id
-
-    from tinygrad.apps.llm import Handler
-    from tinygrad.viz.serve import TCPServerWithReuse
-
-    cls.server = TCPServerWithReuse(('127.0.0.1', 0), Handler)
-    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)
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_rearrange_einops.py

@@ -9,66 +9,6 @@ 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)
-
   def test_tensor_train_example_numpy(self):
     # kept here just for a collection, only tested for numpy
     # https://arxiv.org/pdf/1509.06569.pdf, (5)
@@ -109,40 +49,6 @@ class test_rearrange_examples(unittest.TestCase):
 
 
 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)
-
   def test_rearrange_ellipsis_ops(self):
     identity_patterns = [
       "...->...",
@@ -237,69 +143,8 @@ class test_rearrange_ops(unittest.TestCase):
 
       assert np.array_equal(result, expected_result)
 
-def check_expression_helper(expression: str):
-  Tensor.ones((1, 2, 3, 4, 5, 6, 7))
 
 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...")
-
   def test_unicode_ellipsis(self):
     equivalent_rearrange_patterns = [
       ("a b … -> (a b) … ", "a b ... -> (a b) ... "),

test/unit/test_schedule_cache.py

@@ -1,16 +1,12 @@
 import unittest
 import functools
-from tinygrad import Tensor, Variable, UOp, Context
-from tinygrad.helpers import cpu_events
+from tinygrad import Tensor, Variable, UOp
 from tinygrad.uop.ops import KernelInfo
 from tinygrad.engine.schedule import schedule_cache
 
 def custom_set0_kernel(A:UOp, num:int) -> UOp:
   return A[0].set(num).sink(arg=KernelInfo(f"custom_set0_{num}"))
 
-def schedule_one():
-  Tensor([1]).schedule()
-
 class TestScheduleCache(unittest.TestCase):
   def test_bound_variable_reuses_cache(self):
     schedule_cache.clear()
@@ -27,14 +23,6 @@ class TestScheduleCache(unittest.TestCase):
     self.assertEqual(t2.item(), 110.0)
     self.assertEqual(len(schedule_cache), cache_size_after_first)
 
-  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.schedule_with_vars()
-    self.assertEqual(var_vals, {'pos': 42})
-
   def test_custom_kernel(self):
     for i in range(4):
       a = Tensor.empty(1)
@@ -77,27 +65,5 @@ class TestScheduleCache(unittest.TestCase):
       print(num)
     self.assertEqual(len(schedule_cache), start_len_schedule_cache)
 
-  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()

test/unit/test_symbolic_tensor.py

@@ -2,101 +2,6 @@ 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 TestSymbolicPad(unittest.TestCase):
   def test_pad(self):
     v = Variable("v", 1, 100).bind(5)

test/unit/test_tensor_io.py

@@ -1,25 +1,8 @@
 import unittest
-from tinygrad import Tensor, dtypes
+from tinygrad import Tensor
 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)
-
   def test_read(self):
     data = b"Hello World!"
     fobj = TensorIO(Tensor(data))

test/unit/test_winograd.py

@@ -1,7 +1,7 @@
 import unittest, sys
 import numpy as np
-from tinygrad import Tensor, GlobalCounters, dtypes, Context, nn
-from tinygrad.helpers import CI, Profiling, WINO
+from tinygrad import Tensor, GlobalCounters, Context, nn
+from tinygrad.helpers import WINO
 
 @unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
 class TestWinogradClose(unittest.TestCase):
@@ -27,52 +27,6 @@ class TestWinograd(unittest.TestCase):
   def tearDown(self):
     WINO.value = self.old
 
-  def test_profile(self):
-    x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
-    with Profiling(enabled=not CI, sort='time'):
-      Tensor.conv2d(x,w).realize()
-
-  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()), 2)
-
-  def test_backward_kernels(self):
-    x,w = Tensor.empty(1,4,9,9,requires_grad=True).realize(), Tensor.empty(4,4,3,3,requires_grad=True).realize()
-    out = Tensor.conv2d(x,w, padding=1)
-    out.mean().backward()
-    backward_schedule = Tensor.schedule(x.grad, w.grad)
-    self.assertEqual(len(backward_schedule), 4)
-
-  def test_counters(self):
-    IC, OC, X, Y = 4,4,9,9
-    #OC, IC, X, Y = 512, 256, 8, 8
-    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, 3)
-
-  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)
-
   def test_padded_conv2d(self):
     # tests padding order in winograd
     x,w = Tensor.rand(1,3,11,28).realize(), Tensor.rand(4,3,3,3).realize()

tinygrad/codegen/__init__.py

@@ -11,7 +11,7 @@ from tinygrad.codegen.opt import Opt
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load
-from tinygrad.uop.decompositions import get_late_rewrite_patterns
+from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_unsupported_dtypes_patterns, get_transcendental_patterns
 from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render, pm_add_loads
@@ -92,9 +92,12 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
 
   # decompositions
   supported_ops = tuple(ren.code_for_op.keys())
-  pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, ren.device, TRANSCENDENTAL>=2, bool(DISABLE_FAST_IDIV),
-                                                        tuple(EMULATED_DTYPES.tolist(dtypes)))
+  pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, ren.device, bool(DISABLE_FAST_IDIV))
+  pm_unsupported = get_unsupported_dtypes_patterns(ren.device, tuple(EMULATED_DTYPES.tolist(dtypes)))
+  pm_transcendental = symbolic_simple+get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
   sink = graph_rewrite(sink, pm_decomp, ctx=ren.device, name="decompositions")
+  sink = graph_rewrite(sink, pm_unsupported, ctx=ren.device, name="unsupported dtypes", bottom_up=True)
+  sink = graph_rewrite(sink, pm_transcendental, ctx=ren.device, name="transcendental")
 
   # final rules for the renderer (without sym)
   extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])

tinygrad/codegen/gpudims.py

@@ -1,4 +1,4 @@
-import math, functools, operator
+import math
 from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, sint_to_uop
 from tinygrad.helpers import all_int, dedup, get_contraction
 from tinygrad.dtype import dtypes, AddrSpace, Invalid
@@ -88,12 +88,11 @@ def add_gpudims(ctx:Renderer, s:UOp):
   subs = {}
   for r in s_topo:
     # look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID
-    if r.op is Ops.STORE and r.buf_target().ptrdtype.addrspace == AddrSpace.GLOBAL:
-      idx = r.src[0]
+    if r.op is Ops.STORE and (idx := r.src[0]).src[0].ptrdtype.addrspace == AddrSpace.GLOBAL:
       missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges]
       if len(missing_locals):
         assert len(idx.src) == 2, "index has 2 sources"
-        mask: UOp = functools.reduce(operator.and_, [x.eq(0) for x in missing_locals])
+        mask: UOp = UOp.prod(*[x.eq(0) for x in missing_locals])
         subs[idx] = idx.replace(src=(idx.src[0], mask.broadcast(idx.src[1].dtype.count).where(idx.src[1], Invalid)))
     if r.op is not Ops.RANGE: continue
     try:

tinygrad/codegen/late/devectorizer.py

@@ -1,5 +1,5 @@
 from typing import Any, cast
-import functools, operator, itertools
+import functools, itertools
 from collections import defaultdict
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
@@ -41,7 +41,7 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
           break
 
   if not drop_stmt and idx is start_idx: return None
-  new_valid = functools.reduce(operator.and_, ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
+  new_valid = UOp.prod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
   return buf.index(idx.valid(new_valid) if new_valid is not None else idx, ptr=True)
 
 

tinygrad/device.py

@@ -177,20 +177,16 @@ class Buffer:
   def as_dmaref(self) -> DMARef:
     assert hasattr(self.allocator, "_as_dmaref"), f"Device {self.device} doesn't support DMA"
     return self.allocator._as_dmaref(self._buf)
-  def as_buffer(self, allow_zero_copy=False, force_zero_copy=False) -> memoryview:
-    # zero copy with as_buffer (disabled by default due to use after free)
+  def as_memoryview(self, allow_zero_copy=False, force_zero_copy=False) -> memoryview:
+    # zero copy with as_memoryview (disabled by default due to use after free)
     if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer') and (self.options is None or self.options.image is None):
       return self.allocator._as_buffer(self._buf)
     assert not force_zero_copy, "force zero copy was passed, but copy is required"
     return self.copyout(memoryview(bytearray(self.nbytes)))
-  def as_typed_buffer(self, shape=None, allow_zero_copy=False, force_zero_copy=False) -> memoryview:
-    assert self.dtype.base.fmt is not None, f"no fmt dtype for {self.dtype.base}"
-    assert self.dtype.base.fmt != "e" or sys.version_info >= (3, 12)
-    return self.as_buffer(allow_zero_copy, force_zero_copy).cast(self.dtype.base.fmt, shape if shape is not None else (self.size,))
   def numpy(self) -> 'np.ndarray': # type: ignore [name-defined] # noqa: F821
     import numpy as np
     assert _to_np_dtype(self.dtype.base) is not None, f"no np dtype for {self.dtype.base}"
-    return np.frombuffer(self.as_buffer(), dtype=_to_np_dtype(self.dtype.base))
+    return np.frombuffer(self.as_memoryview(), dtype=_to_np_dtype(self.dtype.base))
   def copyin(self, mv:memoryview):
     mv = flat_mv(mv)
     assert len(mv) == self.nbytes, f"size mismatch, {len(mv)=} != {self.dtype=} {self.size=}"

tinygrad/engine/jit.py

@@ -207,7 +207,7 @@ class CapturedJit(Generic[ReturnType]):
     asgn = _internal_memory_planner([[b for item in self.jit_cache for b in item.bufs if b is not None and b not in blacklist]], ignore_checks=True)
     self.jit_cache = [replace(item, bufs=[asgn.get(b,b) if b is not None else None for b in item.bufs]) for item in self.jit_cache]
     for old, new in asgn.items():
-      if old.is_allocated(): new.ensure_allocated().copyin(old.as_buffer())
+      if old.is_allocated(): new.ensure_allocated().copyin(old.as_memoryview())
     self.__post_init__()
 
   # jit exec
@@ -219,7 +219,7 @@ class CapturedJit(Generic[ReturnType]):
     # copy aliased inputs to prevent read-after-write hazard
     for i, ib in enumerate(input_buffers):
       if (writer := self._output_to_writer.get(ib)) is not None and self._input_to_max_reader.get(i, -1) > writer:
-        input_buffers[i] = Buffer(ib.device, ib.size, ib.dtype).ensure_allocated().copyin(ib.as_buffer())
+        input_buffers[i] = Buffer(ib.device, ib.size, ib.dtype).ensure_allocated().copyin(ib.as_memoryview())
 
     for (j,i),input_idx in self._input_replace.items(): self._jit_cache[j].bufs[i] = input_buffers[input_idx]
 
@@ -257,8 +257,7 @@ def _prepare_jit_inputs(args, kwargs):
   input_uops: list[UOp] = flatten([t.uop.src if t.uop.op is Ops.MULTI else [t.uop] for t in tensors])
   if any(u.base.op is Ops.CONST for u in input_uops):
     raise JitError("JIT inputs cannot be const, create a buffer with .contiguous()")
-  input_buffers: list[Buffer] = flatten([b.bufs if isinstance(b:=u.base.realized, MultiBuffer) else [b]
-                                         for u in input_uops if u.base.realized is not None])
+  input_buffers: list[Buffer] = flatten([b.bufs if isinstance(b, MultiBuffer) else [b] for u in input_uops if (b:=u.base.realized) is not None])
   if len(set(input_buffers)) != len(input_buffers): raise JitError("duplicate inputs to JIT")
   inputs = [(*(u.substitute({u.base:UOp(Ops.NOOP)}, extra_pm=mop_cleanup).unbind_all()), u.dtype, u.device) for u in input_uops]
   _var_vals = merge_dicts([x[1] for x in inputs] + [dict(v.unbind() for v in (args + tuple(kwargs.values())) if isinstance(v, UOp))])
@@ -311,7 +310,7 @@ class TinyJit(Generic[ReturnType]):
       assert self.fxn is not None
       with Context(BEAM=0 if getenv("IGNORE_JIT_FIRST_BEAM") else BEAM.value):
         ret = self.fxn(*args, **kwargs)
-        if len(params:=get_parameters(ret)): Tensor.realize(params[0], *params[1:])
+        if len(params:=get_parameters(ret)): Tensor.realize(*params)
     elif self.cnt == 1:
       # jit capture
       assert self.fxn is not None
@@ -323,7 +322,7 @@ class TinyJit(Generic[ReturnType]):
         capturing.append(self)
         try:
           ret = self.fxn(*args, **kwargs)
-          if len(params:=get_parameters(ret)): Tensor.realize(params[0], *params[1:])
+          if len(params:=get_parameters(ret)): Tensor.realize(*params)
         finally: capturing.clear()
       jit_cache = self._jit_cache
       del self._buffer_replace, self._jit_cache

tinygrad/engine/realize.py

@@ -79,7 +79,7 @@ class BufferCopy(Runner):
       # fast(ish) path, uses readinto in diskbuffers
       src.allocator._copyout(dest.allocator._as_buffer(dest._buf), src._buf)
     else:
-      dest.copyin(src.as_buffer(allow_zero_copy=True))  # may allocate a CPU buffer depending on allow_zero_copy
+      dest.copyin(src.as_memoryview(allow_zero_copy=True))  # may allocate a CPU buffer depending on allow_zero_copy
   def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False):
     dest, src = rawbufs[0:2]
     assert dest.size == src.size and dest.dtype == src.dtype, f"buffer copy mismatch, {dest.size} != {src.size}, {dest.dtype} != {src.dtype}"
@@ -199,7 +199,7 @@ def run_schedule(schedule:list[ExecItem], var_vals:dict[str, int]|None=None, do_
       bufs = [b for b in ei.bufs if b is not None]
       nb: list[Buffer|None] = [Buffer("CPU", b.size, b.dtype) for b in bufs]
       for cpu_b, gpu_b in zip(nb, bufs):
-        if cpu_b is not None and gpu_b.is_allocated(): cpu_b.ensure_allocated().copyin(gpu_b.as_buffer())
+        if cpu_b is not None and gpu_b.is_allocated(): cpu_b.ensure_allocated().copyin(gpu_b.as_memoryview())
 
       # run on GPU
       ei.run(var_vals, do_update_stats=do_update_stats)

tinygrad/engine/schedule.py

@@ -1,7 +1,7 @@
 import time
 from typing import cast
 from collections import deque
-from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, Kernel
+from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, CallInfo
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Buffer, MultiBuffer
 from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize, SCACHE, Metadata
@@ -9,8 +9,8 @@ from tinygrad.engine.realize import ExecItem
 
 # **** schedule linearizer
 
-# ScheduleItem = tuple[AST, buffer UOps, metadata, fixedvars, bound_ranges]
-ScheduleItem = tuple[UOp, tuple[UOp, ...], tuple[Metadata, ...], dict[str, int], tuple[UOp, ...]]
+# ScheduleItem = tuple[AST, buffer UOps, metadata, bound_ranges]
+ScheduleItem = tuple[UOp, tuple[UOp, ...], tuple[Metadata, ...], tuple[UOp, ...]]
 
 # unwrap VIEW/CAST/etc to find the actual data source (kernel output, buffer, or multi-device op)
 def _unwrap_src(s: UOp) -> UOp:
@@ -23,13 +23,12 @@ def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
     children: dict[UOp, list[UOp]] = {}
     in_degree: dict[UOp, int] = {}
     for u in sched_sink.toposort():
-      if u.op is Ops.RANGE:
-        in_degree.setdefault(u, 0)
-        continue
-      if u.op is not Ops.AFTER or u.src[1].op is Ops.RANGE: continue
-      k = u.src[1]
+      if u.op is Ops.RANGE: in_degree.setdefault(u, 0)
+      if u.op is not Ops.AFTER: continue
+      if (k:=u.src[1]).op is Ops.RANGE: continue  # RANGEs are scheduled directly, not through dependency graph
+      assert k.op in {Ops.KERNEL, Ops.END}, f"AFTER src[1] should be KERNEL or END, not {k.op}"
       in_degree.setdefault(k, 0)
-      for s in k.src[0].src if k.op is Ops.END else k.src:
+      for s in k.src[0].src if k.op is Ops.END else k.src[1:]:
         match (s := _unwrap_src(s)).op:
           case Ops.AFTER:
             children.setdefault(s.src[1], []).append(k)
@@ -42,51 +41,51 @@ def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
                 children.setdefault(ss.src[1], []).append(k)
                 in_degree[k] += 1
           case Ops.BUFFER | Ops.BIND:
-            pass  # BUFFER is already realized, BIND is outer range (handled via bound_ranges below)
+            pass  # BUFFER is already realized, BIND is a bound variable (not a buffer dependency)
           case _:
             raise RuntimeError(f"input to kernel must be AFTER, BUFFER, MSELECT, MSTACK, or BIND, not {s.op}")
 
   with cpu_profile(TracingKey("linearize schedule")):
-    queue: deque[UOp] = deque()
-    for k,v in in_degree.items():
-      if v == 0: queue.append(k)
+    queue: deque[UOp] = deque(k for k,v in in_degree.items() if v == 0)
 
-    schedule: list[ScheduleItem|UOp] = []  # ScheduleItem for kernels, UOp for RANGE/END
+    schedule: list[UOp] = []  # RANGE, KERNEL, or END UOps
+    sched_item: dict[UOp, ScheduleItem] = {}
     while len(queue):
       k = rk = queue.popleft()
       if k.op is Ops.END: k = k.src[0]
-      assert k.op in {Ops.RANGE, Ops.KERNEL}, f"unexpected op in queue: {k.op}"
+      assert k.op in {Ops.RANGE, Ops.CALL}, f"unexpected op in queue: {k.op}"
       if k.op is Ops.RANGE: schedule.append(k)
-      elif k.op is Ops.KERNEL:
-        ast = (kernel:=cast(Kernel, k.arg)).ast
-        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src if s.op is not Ops.BIND)
-        bound_ranges = tuple(s for s in k.src if s.op is Ops.BIND and len(s.src) > 1 and s.src[1].op is Ops.RANGE)
-        schedule.append((ast, buf_uops, kernel.metadata, {}, bound_ranges))
+      elif k.op is Ops.CALL:
+        ast = k.src[0]
+        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if s.op is not Ops.BIND)
+        bound_ranges = tuple(s for s in k.src[1:] if s.op is Ops.BIND and len(s.src) > 1 and s.src[1].op is Ops.RANGE)
+        schedule.append((ast, buf_uops, cast(CallInfo, k.arg).metadata, {}, bound_ranges))
         if rk.op is Ops.END: schedule.append(rk)
       for x in children.get(rk, []):
         in_degree[x] -= 1
         if in_degree[x] == 0: queue.append(x)
 
   with cpu_profile(TracingKey("unroll outer ranges")):
-    pre_schedule, buf_uops_list = unroll_outer_ranges(schedule)
+    pre_schedule, buf_uops_list = unroll_outer_ranges(schedule, sched_item)
   return pre_schedule, UOp.sink(*buf_uops_list)
 
-def unroll_outer_ranges(schedule:list[ScheduleItem|UOp]) -> tuple[list[ExecItem], list[UOp]]:
+def unroll_outer_ranges(schedule:list[UOp], sched_item:dict[UOp, ScheduleItem]) -> tuple[list[ExecItem], list[UOp]]:
   pre_schedule: list[ExecItem] = []
   buf_uops_list: list[UOp] = []
   sched_ptr, in_ranges, range_ptrs = 0, dict[UOp, int](), dict[UOp, int]()
   while sched_ptr < len(schedule):
-    if isinstance(si := schedule[sched_ptr], UOp):
-      if si.op is Ops.RANGE:
-        in_ranges[si] = 0
-        range_ptrs[si] = sched_ptr + 1
-      elif si.op is Ops.END:
-        if in_ranges[si.src[1]] < si.src[1].vmax:
-          in_ranges[si.src[1]] += 1
-          sched_ptr = range_ptrs[si.src[1]]
-          continue
+    si = schedule[sched_ptr]
+    if si.op is Ops.RANGE:
+      in_ranges[si] = 0
+      range_ptrs[si] = sched_ptr + 1
+    elif si.op is Ops.END:
+      if in_ranges[si.src[1]] < si.src[1].vmax:
+        in_ranges[si.src[1]] += 1
+        sched_ptr = range_ptrs[si.src[1]]
+        continue
     else:
-      ast, buf_uops, metadata, _, bound_ranges = si
+      assert si.op is Ops.KERNEL, f"unexpected op in schedule: {si.op}"
+      ast, buf_uops, metadata, bound_ranges = sched_item[si]
       fixedvars = {s.src[0].arg[0]:in_ranges[s.src[1]] for s in bound_ranges}
       pre_schedule.append(ExecItem(ast, [], metadata, fixedvars))
       buf_uops_list.append(UOp.sink(*buf_uops))

tinygrad/gradient.py

@@ -14,7 +14,7 @@ def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
   if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
 
 def call_gradient(ctx:UOp, k:UOp):
-  if k.arg is not None: return (None,) + k.arg(ctx, k)
+  if k.arg.grad_fxn is not None: return (None,) + k.arg.grad_fxn(ctx, k)
   # auto-differentiate the function
   fxn, args = k.src[0], k.src[1:]
   params = sorted([x for x in fxn.toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)

tinygrad/runtime/ops_amd.py

@@ -43,10 +43,9 @@ class ProfilePMCEvent(ProfileEvent): device:str; kern:int; sched:list[PMCSample]
 class AMDSignal(HCQSignal):
   def __init__(self, *args, **kwargs): super().__init__(*args, **{**kwargs, 'timestamp_divider': 100})
 
-  def _sleep(self, time_spent_waiting_ms:int) -> bool:
-    # Resonable to sleep for long workloads (which take more than 2s) and only timeline signals.
-    if time_spent_waiting_ms > 2000 and self.is_timeline and self.owner is not None: return self.owner.iface.sleep(200)
-    return False
+  def _sleep(self, time_spent_since_last_sleep_ms:int):
+    # Reasonable to sleep for long workloads (which take more than 200ms) and only timeline signals.
+    if time_spent_since_last_sleep_ms > 200 and self.is_timeline and self.owner is not None: self.owner.iface.sleep(200)
 
 class AMDComputeQueue(HWQueue):
   def __init__(self, dev:AMDDevice):
@@ -778,9 +777,8 @@ class KFDIface:
                         write_ptr=MMIOInterface(queue.write_pointer_address, 8, fmt='Q'),
                         doorbell=MMIOInterface(self.doorbells + queue.doorbell_offset - self.doorbells_base, 8, fmt='Q'))
 
-  def sleep(self, tm:int) -> bool:
+  def sleep(self, tm:int):
     kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=self.queue_event_arr_ptr, num_events=1, wait_for_all=1, timeout=tm)
-    return False
 
   def on_device_hang(self):
     def _collect_str(st): return ' '.join(f'{k[0]}={getattr(st, k[0])}' for k in st._real_fields_)
@@ -857,11 +855,11 @@ class PCIIface(PCIIfaceBase):
     return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q'), put_value=pv,
       read_ptr=gart.cpu_view().view(offset=rptr, size=8, fmt='Q'), write_ptr=gart.cpu_view().view(offset=wptr, size=8, fmt='Q'), params=rcvr_params)
 
-  def sleep(self, timeout) -> bool:
+  def sleep(self, timeout):
     if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
       self.pci_dev.irq_fd.read(8 * events_cnt)
     self.dev_impl.ih.interrupt_handler()
-    return self.dev_impl.is_err_state
+    if self.dev_impl.is_err_state: raise RuntimeError("Device fault detected")
 
   def on_device_hang(self):
     devs:list[AMDDevice] = [d for pg in HCQCompiled.peer_groups.values() for d in pg if isinstance(d, AMDDevice) and d.is_am()]
@@ -905,7 +903,7 @@ class USBIface(PCIIface):
     if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE: self.pci_dev.usb._pci_cacheable += [(ring.cpu_view().addr, ring.size)]
     return super().create_queue(queue_type, ring, gart, rptr, wptr, eop_buffer, cwsr_buffer, ctl_stack_size, ctx_save_restore_size, xcc_id, idx)
 
-  def sleep(self, timeout) -> bool: return False
+  def sleep(self, timeout): pass
 
 class AMDDevice(HCQCompiled):
   def is_am(self) -> bool: return isinstance(self.iface, (PCIIface, USBIface))

tinygrad/runtime/ops_cpu.py

@@ -13,9 +13,8 @@ from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.uop.ops import sint
 
 class CPUSignal(HCQSignal):
-  def _sleep(self, time_spent_waiting_ms:int) -> bool:
+  def _sleep(self, time_spent_since_last_sleep_ms:int):
     if self.is_timeline and self.owner is not None: self.owner.tasks.join()
-    return False
 
 class CPUWorker(threading.Thread):
   def __init__(self, dev, tasks, thread_id):

tinygrad/runtime/ops_nv.py

@@ -24,13 +24,12 @@ nv_gpu = nv_570 # default to 570
 PMA = ContextVar("PMA", abs(VIZ.value)>=2)
 
 @dataclass(frozen=True)
-class ProfilePMAEvent(ProfileEvent): device:str; kern:str; blob:bytes # noqa: E702
+class ProfilePMAEvent(ProfileEvent): device:str; kern:str; blob:bytes; exec_tag:int # noqa: E702
 
 class NVSignal(HCQSignal):
-  def _sleep(self, time_spent_waiting_ms:int) -> bool:
-    # Resonable to sleep for long workloads (which take more than 2s) and only timeline signals.
-    if time_spent_waiting_ms > 2000 and self.is_timeline and self.owner is not None: return self.owner.iface.sleep(200)
-    return False
+  def _sleep(self, time_spent_since_last_sleep_ms:int):
+    # Reasonable to sleep for long workloads (which take more than 200ms) and only timeline signals.
+    if time_spent_since_last_sleep_ms > 200 and self.is_timeline and self.owner is not None: self.owner.iface.sleep(200)
 
 def get_error_str(status): return f"{status}: {nv_gpu.nv_status_codes.get(status, 'Unknown error')}"
 
@@ -315,7 +314,8 @@ class NVProgram(HCQProgram):
     res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait)
     if self.dev.pma_enabled:
       self.dev.synchronize()
-      if pma_blob:=self.dev._prof_readback(): Compiled.profile_events += [ProfilePMAEvent(self.dev.device, self.name, pma_blob)]
+      if pma_blob:=self.dev._prof_readback():
+        Compiled.profile_events += [ProfilePMAEvent(self.dev.device, self.name, pma_blob, self.dev.prof_exec_counter)]
     return res
 
 class NVAllocator(HCQAllocator['NVDevice']):
@@ -525,7 +525,7 @@ class NVKIface:
   def _alloc_gpu_vaddr(self, size, alignment=(4 << 10), force_low=False):
     return NVKIface.low_uvm_vaddr_allocator.alloc(size, alignment) if force_low else NVKIface.uvm_vaddr_allocator.alloc(size, alignment)
 
-  def sleep(self, tm:int) -> bool: return False
+  def sleep(self, tm:int): pass
 
 class PCIIface(PCIIfaceBase):
   gpus:ClassVar[list[str]] = []
@@ -561,9 +561,9 @@ class PCIIface(PCIIfaceBase):
 
   def device_fini(self): self.dev_impl.fini()
 
-  def sleep(self, timeout) -> bool:
+  def sleep(self, timeout):
     for _ in self.dev_impl.gsp.stat_q.read_resp(): pass
-    return self.dev_impl.is_err_state
+    if self.dev_impl.is_err_state: raise RuntimeError("Device fault detected")
 
 class NVDevice(HCQCompiled[NVSignal]):
   def is_nvd(self) -> bool: return isinstance(self.iface, PCIIface)

tinygrad/runtime/ops_qcom.py

@@ -2,21 +2,33 @@ from __future__ import annotations
 import os, ctypes, functools, mmap, struct, array, math, sys, weakref, contextlib
 assert sys.platform != 'win32'
 from typing import Any
-from tinygrad.device import BufferSpec, CompilerSet, CompilerPair
+from tinygrad.device import BufferSpec, CompilerSet, CompilerPair, Device
 from tinygrad.runtime.support.hcq import HCQBuffer, HWQueue, HCQProgram, HCQCompiled, HCQAllocatorBase, HCQSignal, HCQArgsState, BumpAllocator
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
 from tinygrad.runtime.autogen import kgsl, mesa
 from tinygrad.runtime.ops_cl import CLCompiler, CLDevice
 from tinygrad.renderer.cstyle import QCOMRenderer
 from tinygrad.renderer.nir import IR3Renderer
-from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport, cpu_profile, lo32, PROFILE, suppress_finalizing
-from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC
-from tinygrad.dtype import ImageDType
+from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, ceildiv, prod, fromimport, cpu_profile, lo32, suppress_finalizing
+from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC, PROFILE
+from tinygrad.dtype import ImageDType, dtypes
 from tinygrad.runtime.support.system import System
 if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl  # noqa: F401  # pylint: disable=unused-import
 
 BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2
 
+@functools.cache
+def dcache_flush():
+  from tinygrad.uop.ops import UOp, Ops, KernelInfo
+  from tinygrad.codegen import get_program
+  buf, n = UOp(Ops.PARAM, dtypes.uint8.ptr(), arg=0), UOp(Ops.PARAM, dtypes.uint8.ptr(), arg=1)
+  i = UOp.range(n.cast(dtypes.int), 0, dtype=dtypes.int)
+  flush = UOp(Ops.CUSTOM, dtypes.void, (buf.cast(dtypes.ulong) + i.cast(dtypes.ulong) * UOp.const(dtypes.ulong, 64),),
+              arg='__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");')
+  sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, dtypes.void, (), arg='__asm__ volatile("dsb sy" ::: "memory");'), arg=KernelInfo(name="dcache_flush"))
+  ps = get_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="CPU"), UOp(Ops.LINEAR, src=tuple(sink.toposort())))), Device["CPU"].renderer)
+  return Device["CPU"].runtime(ps.function_name, ps.lib)
+
 #Parse C-style defines: <regname>_<field_x>__SHIFT and <regname>_<field_y>__MASK from the adreno module into the following format:
 # qreg.<regname>(<field_x>=..., <field_y>=..., ..., <field_n>=...)
 def _qreg_exec(__reg, __val=0, **kwargs):
@@ -45,11 +57,10 @@ class QCOMCompiler(CLCompiler):
 class QCOMSignal(HCQSignal):
   def __init__(self, *args, **kwargs): super().__init__(*args, **{**kwargs, 'timestamp_divider': 19.2})
 
-  def _sleep(self, time_spent_waiting_ms:int) -> bool:
+  def _sleep(self, time_spent_since_last_sleep_ms:int):
     # Sleep only for timeline signals. Do it immediately to free cpu.
     if self.is_timeline and self.owner is not None:
       kgsl.IOCTL_KGSL_DEVICE_WAITTIMESTAMP_CTXTID(self.owner.fd, context_id=self.owner.ctx, timestamp=self.owner.last_cmd, timeout=0xffffffff)
-    return False
 
 class QCOMComputeQueue(HWQueue):
   def __init__(self, dev:QCOMDevice):
@@ -391,6 +402,7 @@ class QCOMDevice(HCQCompiled):
 
   def _gpu_map(self, ptr:int, size:int, **kwargs) -> HCQBuffer:
     ptr_aligned, size_aligned = (ptr & ~0xfff), round_up(size + (ptr & 0xfff), 0x1000)
+    dcache_flush().fxn(ctypes.c_uint64(ptr_line_aligned:=ptr & ~63), ctypes.c_uint64(ceildiv(ptr + size - ptr_line_aligned, 64)))
     try:
       mi = kgsl.IOCTL_KGSL_MAP_USER_MEM(self.fd, hostptr=ptr_aligned, len=size_aligned, memtype=kgsl.KGSL_USER_MEM_TYPE_ADDR)
       return HCQBuffer(mi.gpuaddr + (ptr - ptr_aligned), size=size, meta=(mi, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self, **kwargs)

tinygrad/runtime/support/am/ip.py

@@ -132,8 +132,9 @@ class AM_GMC(AM_IP):
       self.adev.reg(f"reg{ip}MC_VM_MX_L1_TLB_CNTL").update(enable_l1_tlb=1, system_access_mode=3, enable_advanced_driver_model=1,
         system_aperture_unmapped_access=0, mtype=self.adev.soc.module.MTYPE_UC, inst=inst)
 
-      self.adev.reg(f"reg{ip}VM_L2_CNTL").update(enable_l2_cache=1, enable_l2_fragment_processing=0, enable_default_page_out_to_system_memory=1,
-        l2_pde0_cache_tag_generation_mode=0, pde_fault_classification=0, context1_identity_access_mode=1, identity_mode_fragment_size=0, inst=inst)
+      self.adev.reg(f"reg{ip}VM_L2_CNTL").update(enable_l2_cache=1, enable_default_page_out_to_system_memory=1,
+        l2_pde0_cache_tag_generation_mode=0, pde_fault_classification=0, context1_identity_access_mode=1, identity_mode_fragment_size=0,
+        enable_l2_fragment_processing=int(self.adev.ip_ver[am.GC_HWIP] < (10,0,0)), inst=inst)
       self.adev.reg(f"reg{ip}VM_L2_CNTL2").update(invalidate_all_l1_tlbs=1, invalidate_l2_cache=1, inst=inst)
       self.adev.reg(f"reg{ip}VM_L2_CNTL3").write(l2_cache_4k_associativity=1, l2_cache_bigk_associativity=1,
         bank_select=12 if self.trans_futher else 9, l2_cache_bigk_fragment_size=9 if self.trans_futher else 6, inst=inst)
@@ -263,7 +264,7 @@ class AM_GFX(AM_IP):
       self.adev.regGRBM_CNTL.update(read_timeout=0xff, inst=xcc)
       for i in range(0, 16):
         self._grbm_select(vmid=i, inst=xcc)
-        self.adev.regSH_MEM_CONFIG.write(**({'initial_inst_prefetch':3} if self.adev.ip_ver[am.GC_HWIP][0] >= 10 else {}),
+        self.adev.regSH_MEM_CONFIG.write(**({'initial_inst_prefetch':3} if self.adev.ip_ver[am.GC_HWIP][0]>=10 else {'retry_disable':1, 'f8_mode':1}),
           address_mode=self.adev.soc.module.SH_MEM_ADDRESS_MODE_64, alignment_mode=self.adev.soc.module.SH_MEM_ALIGNMENT_MODE_UNALIGNED, inst=xcc)
 
         # Configure apertures:
@@ -344,14 +345,15 @@ class AM_GFX(AM_IP):
       self.adev.regRLC_CGCG_CGLS_CTRL.update(cgcg_gfx_idle_threshold=0x36, cgcg_en=1, cgls_rep_compansat_delay=0xf, cgls_en=1, inst=xcc)
 
       self.adev.regCP_RB_WPTR_POLL_CNTL.update(poll_frequency=0x100, idle_poll_count=0x90, inst=xcc)
-      self.adev.regCP_INT_CNTL.update(cntx_busy_int_enable=1, cntx_empty_int_enable=1, cmp_busy_int_enable=1, gfx_idle_int_enable=1, inst=xcc)
+      self.adev.regCP_INT_CNTL.update(cntx_busy_int_enable=1, cntx_empty_int_enable=1, cmp_busy_int_enable=1, inst=xcc)
       if self.adev.ip_ver[am.GC_HWIP] >= (10,0,0):
         self.adev.regSDMA0_RLC_CGCG_CTRL.update(cgcg_int_enable=1, inst=xcc)
         self.adev.regSDMA1_RLC_CGCG_CTRL.update(cgcg_int_enable=1, inst=xcc)
 
-      feats_gfx11 = {'perfmon_clock_state':1, 'gfxip_repeater_fgcg_override':0} if self.adev.ip_ver[am.GC_HWIP] >= (11,0,0) else {}
-      self.adev.regRLC_CGTT_MGCG_OVERRIDE.update(**feats_gfx11, gfxip_fgcg_override=0, grbm_cgtt_sclk_override=0, rlc_cgtt_sclk_override=0,
-        gfxip_mgcg_override=0, gfxip_cgls_override=0, gfxip_cgcg_override=0, inst=xcc)
+      feats_gfx9 = {'gfxip_mgls_override':0, 'gfxip_rep_fgcg_override':0} if self.adev.ip_ver[am.GC_HWIP][0] == 9 else {}
+      feats_gfx11 = {'perfmon_clock_state':1, 'gfxip_repeater_fgcg_override':0} if self.adev.ip_ver[am.GC_HWIP][0] >= 11 else {}
+      self.adev.regRLC_CGTT_MGCG_OVERRIDE.update(**feats_gfx9, **feats_gfx11, gfxip_fgcg_override=0, grbm_cgtt_sclk_override=0,
+        rlc_cgtt_sclk_override=0, gfxip_mgcg_override=0, gfxip_cgls_override=0, gfxip_cgcg_override=0, inst=xcc)
 
       self.adev.regRLC_SAFE_MODE.write(message=0, cmd=1, inst=xcc)
 

tinygrad/runtime/support/hcq.py

@@ -243,12 +243,11 @@ class HCQSignal(Generic[HCQDeviceType]):
     """
     return self.timestamp_mv[0] / self.timestamp_divider
 
-  def _sleep(self, time_spent_waiting_ms:int) -> bool:
+  def _sleep(self, time_spent_since_last_sleep_ms:int):
     """
     Optional function which can implement sleep functionality for the signal.
-    Returns True if a fault was detected, False otherwise.
+    Raises RuntimeError if a fault is detected.
     """
-    return False
 
   def wait(self, value:int, timeout:int=getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000)):
     """
@@ -258,12 +257,12 @@ class HCQSignal(Generic[HCQDeviceType]):
       value: The value to wait for.
       timeout: Maximum time to wait in milliseconds. Defaults to 30s.
     """
-    start_time, fault = int(time.perf_counter() * 1000), False
-    while (not_passed:=(prev_value:=self.value) < value) and (time_spent:=int(time.perf_counter() * 1000) - start_time) < timeout:
-      if fault:=self._sleep(time_spent): break
-      if self.value != prev_value: start_time = int(time.perf_counter() * 1000) # progress was made, reset timer
-    if not_passed and self.value < value:
-      raise RuntimeError("Device fault detected" if fault else f"Wait timeout: {timeout} ms! (the signal is not set to {value}, but {self.value})")
+    start_time = last_sleep_time = int(time.perf_counter() * 1000)
+    while (not_passed:=(prev_value:=self.value) < value) and (cur_time:=int(time.perf_counter() * 1000)) - start_time < timeout:
+      self._sleep(cur_time - last_sleep_time)
+      last_sleep_time = int(time.perf_counter() * 1000)
+      if self.value != prev_value: start_time = last_sleep_time # progress was made, reset timer
+    if not_passed and self.value < value: raise RuntimeError(f"Wait timeout: {timeout} ms! (the signal is not set to {value}, but {self.value})")
 
 @contextlib.contextmanager
 def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]|None=None, queue:HWQueue|None=None):

tinygrad/schedule/indexing.py

@@ -1,5 +1,5 @@
 from typing import Iterator
-import functools, operator, itertools
+import functools, itertools
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches
@@ -9,7 +9,7 @@ from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored
 
 ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
                      Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
-                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.KERNEL, Ops.ENCDEC}
+                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.CALL, Ops.ENCDEC}
 
 def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None
 
@@ -20,7 +20,7 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
 def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
   if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
   # if it's a kernel, we don't realize it
-  if a.src[1].op is not Ops.KERNEL: ctx[a] = None
+  if a.src[1].op is not Ops.CALL: ctx[a] = None
 
 pm_generate_realize_map = PatternMatcher([
   # always realize SINK src
@@ -83,7 +83,7 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
 
 def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp):
   if x not in ctx.range_map: return None
-  valid: UOp = functools.reduce(operator.and_, [r.get_valid() for r in ctx.range_map[x][0]], UOp.const(dtypes.bool, True))
+  valid: UOp = UOp.const(dtypes.bool, True).prod(*[r.get_valid() for r in ctx.range_map[x][0]])
   ret = valid.where(x.src[0], UOp.const(x.dtype, 0))
   ctx.range_map[ret] = ctx.range_map[x]
   return ret
@@ -99,7 +99,7 @@ def remove_movement_op_after_rangeify(ctx:IndexingContext, x:UOp):
   if x in ctx.range_map or x.src[0].op is Ops.INDEX: return x.src[0]
 
 def add_third_op_to_assign_to_track_shape(ctx:IndexingContext, assign:UOp):
-  if assign.src[1].op is Ops.KERNEL: return None
+  if assign.src[1].op is Ops.CALL: return None
   to_mop = graph_rewrite(assign.src[0], PatternMatcher([(UPat(GroupOp.Movement, name="x"), lambda x: x.replace(tag=()))]))
   ret = assign.replace(src=assign.src+(to_mop,))
   ctx.range_map[ret] = ctx.range_map[assign]
@@ -127,7 +127,7 @@ def _apply_reshape(in_shape:tuple[sint,...], out_shape:tuple[sint, ...], urngs:U
   for s,src in list(zip(out_shape, urngs.src))[::-1]:
     axes_in.append(acc*src)
     acc *= s
-  combined_axes = sum(axes_in, start=UOp.const(dtypes.index, 0))
+  combined_axes = UOp.const(dtypes.index, 0).sum(*axes_in)
   axes_out:list[UOp] = []
   for s in in_shape[::-1]:
     axes_out.append(combined_axes % s)
@@ -174,7 +174,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue
 
     # no ranges on kernels, they are internal
-    if x.op is Ops.KERNEL: continue
+    if x.op is Ops.CALL: continue
 
     if x.dtype.scalar() == dtypes.index: continue  # TODO: why do I need this?
     ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])
@@ -218,7 +218,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
         # we compare the ranges without their valids
         if all_all_same or (PCONTIG and all_same(local_rngs)):
           # the new valid is the OR of all the children valids
-          minimum_valid = functools.reduce(operator.or_, valids, UOp.const(dtypes.bool, False))
+          minimum_valid = UOp.const(dtypes.bool, False).sum(*valids)
           _out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), symbolic, name="minimum_valid"))
         else:
           _out_rngs.append(rctx.new_range(x.shape[i]))