failing assign test

2025-09-12 12:53:48 +08:00
117 changed files with 1390 additions and 2109 deletions
@@ -28,7 +28,7 @@ jobs:
      # since sudo is required for usbgpu on macos, move the cache to a new location, as some of the files are owned by root
      PYTHONPYCACHEPREFIX: /tmp/tiny_python_pycache
    runs-on: [self-hosted, macOS]
-    timeout-minutes: 60
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -160,7 +160,7 @@ jobs:
  testnvidiabenchmark:
    name: tinybox green Benchmark
    runs-on: [self-hosted, Linux, tinyboxgreen]
-    timeout-minutes: 60
+    timeout-minutes: 30
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -274,7 +274,7 @@ jobs:
  testmorenvidiabenchmark:
    name: tinybox green Training Benchmark
    runs-on: [self-hosted, Linux, tinyboxgreen]
-    timeout-minutes: 60
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -346,7 +346,7 @@ jobs:
  testamdbenchmark:
    name: tinybox red Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 60
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -476,7 +476,7 @@ jobs:
  testmoreamdbenchmark:
    name: tinybox red Training Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 60
+    timeout-minutes: 30
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -511,8 +511,8 @@ jobs:
      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=85 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=188 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
-    # - name: Run 10 CIFAR training steps w BF16
-    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    - name: Run 10 CIFAR training steps w BF16
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    - name: Run 10 CIFAR training steps w winograd
      run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
@@ -539,7 +539,7 @@ jobs:
  testmlperfamdbenchmark:
    name: tinybox red MLPerf Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 60
+    timeout-minutes: 30
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -645,7 +645,7 @@ jobs:
  testreddriverbenchmark:
    name: AM Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
+    timeout-minutes: 15
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -716,7 +716,7 @@ jobs:
  testgreendriverbenchmark:
    name: NV Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
+    timeout-minutes: 15
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -30,6 +30,8 @@ jobs:
        key: llvm-speed
        deps: testing_minimal
        llvm: 'true'
+    - name: External Benchmark Schedule
+      run: python3 test/external/external_benchmark_schedule.py
    - name: Speed Test
      run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
    - name: Speed Test (BEAM=2)
@@ -46,7 +48,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        deps: docs
-        pydeps: "capstone torch"
+        pydeps: "capstone"
    - name: Build wheel and show size
      run: |
        pip install build
@@ -77,8 +79,6 @@ jobs:
      run: |
        python docs/abstractions2.py
        python docs/abstractions3.py
-    - name: Test README
-      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  python README.py
    - name: Test Quickstart
      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' docs/quickstart.md > quickstart.py &&  python quickstart.py
    - name: Test DEBUG
@@ -148,28 +148,18 @@ jobs:
    - name: Test some torch tests (expect failure)
      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true

-  bepython:
-    name: Python Backend
+  tc:
+    name: Tensor Core tests
    runs-on: ubuntu-latest
-    timeout-minutes: 15
+    timeout-minutes: 10
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: be-minimal
+        key: uops-minimal
        deps: testing_minimal
-    - name: Test dtype with Python emulator
-      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
-    - name: Test ops with Python emulator
-      run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20
-    - name: Test uops with Python emulator
-      run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20
-    - name: Test symbolic with Python emulator
-      run: PYTHON=1 python3 test/test_symbolic_ops.py
-    - name: test_renderer_failures with Python emulator
-      run: PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
    - name: Test IMAGE=2 support
      run: |
        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
@@ -216,6 +206,29 @@ jobs:
        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

+  bepython:
+    name: Python Backend
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: be-minimal
+        deps: testing_minimal
+    - name: Test dtype with Python emulator
+      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
+    - name: Test ops with Python emulator
+      run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20
+    - name: Test uops with Python emulator
+      run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20
+    - name: Test symbolic with Python emulator
+      run: PYTHON=1 python3 test/test_symbolic_ops.py
+    - name: test_renderer_failures with Python emulator
+      run: PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
+
  linter:
    name: Linters
    runs-on: ubuntu-latest
@@ -259,23 +272,21 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: unittest-12
-        pydeps: "pillow numpy ftfy regex"
+        pydeps: "pillow"
        deps: testing_unit
+    - name: Test README
+      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  python README.py
    - name: Run unit tests
      run: python -m pytest -n=auto test/unit/ --durations=20
    - name: Run targetted tests on NULL backend
-      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
+      run: NULL=1 python3 test/test_multitensor.py TestMultiTensor.test_data_parallel_resnet_train_step
    - name: Run SDXL on NULL backend
      run: MAX_BUFFER_SIZE=0 NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
-    - name: Run Clip tests for SD MLPerf on NULL backend
-      run: MAX_BUFFER_SIZE=0 NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
    # TODO: support fake weights
    #- name: Run LLaMA 7B on 4 fake devices
    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
    - name: Run GC tests
      run: python test/external/external_uop_gc.py
-    - name: External Benchmark Schedule
-      run: python3 test/external/external_benchmark_schedule.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
    - name: Regen dataset on test_tiny
@@ -312,9 +323,9 @@ jobs:
      run: python test/external/fuzz_shape_ops.py

  testopenclimage:
-    name: CL IMAGE Tests
+    name: 'CL IMAGE Tests'
    runs-on: ubuntu-22.04
-    timeout-minutes: 15
+    timeout-minutes: 10
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -328,15 +339,11 @@ jobs:
        run: |
          CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
          CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
-      - name: Test CL IMAGE=2 ops + training (rangeify)
-        run: |
-          RANGEIFY=1 CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
-          RANGEIFY=1 CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

  testgpumisc:
-    name: CL Misc tests
+    name: 'CL Misc tests'
    runs-on: ubuntu-22.04
    timeout-minutes: 10
    steps:
@@ -361,7 +368,7 @@ jobs:
          path: /tmp/sops.gz

  testopenpilot:
-    name: openpilot Compile Tests
+    name: 'openpilot Compile Tests'
    runs-on: ubuntu-22.04
    timeout-minutes: 15
    steps:
@@ -376,9 +383,7 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2160 ALLOWED_GATED_READ_IMAGE=16 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
-      - name: Test openpilot model with rangeify
-        run: RANGEIFY=1 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+          ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2175 ALLOWED_GATED_READ_IMAGE=16 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot alt model correctness (float32)
        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot fastvits model correctness (float32)
@@ -395,7 +400,7 @@ jobs:
 # ****** ONNX Tests ******

  testonnxcpu:
-    name: ONNX (CPU) Tests
+    name: 'ONNX (CPU) Tests'
    runs-on: ubuntu-22.04
    timeout-minutes: 20

@@ -423,7 +428,7 @@ jobs:
        uses: ./.github/actions/process-replay

  testopencl:
-    name: ONNX (CL)+Optimization Tests
+    name: 'ONNX (GPU)+Optimization Tests'
    runs-on: ubuntu-22.04
    timeout-minutes: 20
    steps:
@@ -511,8 +516,8 @@ jobs:

 # ****** Feature Tests ******

-  testrangeifycpu:
-    name: Linux (rangeify) CPU
+  testrangeify:
+    name: Linux (rangeify)
    runs-on: ubuntu-24.04
    timeout-minutes: 15
    steps:
@@ -527,68 +532,20 @@ jobs:
        llvm: "true"
    - name: Test CPU=1 RANGEIFY=1
      # TODO: add more passing tests here
-      # rangeify diamond cycle gives the wrong answer
+      # test_symbolic_arange_sym_step is passing now
+      # test_threefry_doesnt_use_long is because there's a contig after the long now
      run: |
        CPU=1 CPU_LLVM=0 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \
-          -k "not test_assign_diamond_cycle" \
-          test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_symbolic_ops.py test/test_symbolic_jit.py test/test_tensor_variable.py \
-          test/test_outerworld_range.py test/test_randomness.py test/test_nn.py test/test_arange.py test/test_tensor.py test/test_optim.py \
-          test/test_setitem.py test/test_assign.py test/test_multitensor.py
-    - name: Test CPU=1 CPU_LLVM=1 RANGEIFY=1
-      run: |
-        CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_edgecases.py
-    - name: Test const folding
-      run: CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_const_folding.py -k "not test_cast_padded and not TestReduceOpsConstFolding"
-    # RANGEIFY=2 isn't supported
-    #- name: Test CPU=1 RANGEIFY=2
-    #  run: CPU=1 CPU_LLVM=0 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
+          -k "not test_symbolic_arange_sym_step and not test_threefry_doesnt_use_long" \
+          test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_tensor_variable.py \
+          test/test_outerworld_range.py test/test_sample.py test/test_randomness.py
+    - name: Test GPU=1 RANGEIFY=1
+      run: GPU=1 RANGEIFY=1 pytest -n auto test/test_ops.py
+    - name: Test CPU=1 RANGEIFY=2
+      run: CPU=1 CPU_LLVM=0 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
    # slow (and still wrong on beautiful_mnist)
-    #- name: Test LLVM RANGEIFY=1 (slow tests)
+    #- name: Test LLVM=1 RANGEIFY=1 (slow tests)
    #  run: CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testrangeifycl:
-    name: Linux (rangeify) CL
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: rangeify-cl
-        deps: testing
-        opencl: 'true'
-        llvm: "true"
-    - name: Test CL=1 RANGEIFY=1
-      run: CL=1 RANGEIFY=1 pytest -n auto test/test_ops.py test/test_schedule.py test/test_symbolic_ops.py test/test_jit.py test/unit/test_disk_tensor.py test/models/test_mnist.py test/unit/test_mnist_dataset.py test/test_optim.py --durations 20
-    - name: Test Fuse
-      run: CL=1 RANGEIFY=2 python3 -m pytest --durations 20 test/test_softmax_fusion.py -k "not test_auto_softmax"
-    - name: Test ONNX
-      run: CL=1 RANGEIFY=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testrangeifymacos:
-    name: MacOS (rangeify)
-    runs-on: macos-14
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: metal
-        deps: testing
-    - name: some unit tests
-      run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/unit/test_winograd.py test/unit/test_linalg.py --durations=20
-    - name: Test METAL=1 RANGEIFY=1
-      run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/test_ops.py --durations=20
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay

  testdevectorize:
    name: Linux (devectorize)
@@ -710,7 +667,7 @@ jobs:
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run TestOps.test_add with SQTT
        run: |
-          VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
+          PROFILE=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
      - name: Run process replay tests
        uses: ./.github/actions/process-replay
@@ -41,6 +41,8 @@ BEAM                | [#]        | number of beams in kernel beam search
 DEFAULT_FLOAT       | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
 IMAGE               | [1-2]      | enable 2d specific optimizations
 FLOAT16             | [1]        | use float16 for images instead of float32
+PTX                 | [1]        | enable the specialized [PTX](https://docs.nvidia.com/cuda/parallel-thread-execution/) assembler for Nvidia GPUs. If not set, defaults to generic CUDA codegen backend.
+PROFILE             | [1]        | enable profiling. This feature is supported in NV, AMD, QCOM and METAL backends.
 VISIBLE_DEVICES     | [list[int]]| restricts the NV/AMD devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
 JIT                 | [0-2]      | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
 VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
@@ -2,17 +2,17 @@

 tinygrad supports various runtimes, enabling your code to scale across a wide range of devices. The default runtime can be automatically selected based on the available hardware, or you can force a specific runtime to be default using environment variables (e.g., `CPU=1`).

-| Runtime | Description | Compiler Options | Requirements |
-|---------|-------------|------------------|--------------|
-| [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | nvrtc (default)<br>PTX (`NV_PTX=1`) | Ampere/Ada/Blackwell series GPUs.<br>You can select an interface via `NV_IFACE=(NVK\|PCI)`. See [NV interfaces](#nv-interfaces) for details. |
-| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`AMD_LLVM=1`)<br>HIP/COMGR (`AMD_HIP=1`) | RDNA2 or newer GPUs.<br>You can select an interface via `AMD_IFACE=(KFD\|PCI\|USB)`. See [AMD interfaces](#amd-interfaces) for details. |
-| [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | - | 6xx series GPUs |
-| [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | - | M1+ Macs; Metal 3.0+ for `bfloat` support |
-| [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | nvrtc (default)<br> PTX (`CUDA_PTX=1`) | NVIDIA GPU with CUDA support |
-| [CL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | - | OpenCL 2.0 compatible device |
-| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`CPU_LLVM=1`) | `clang` compiler in system `PATH` |
-| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | - | Dawn library installed and discoverable. Binaries: [pydawn v0.3.0](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0) |
-
+| Runtime | Description | Requirements |
+|---------|-------------|--------------|
+| [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | Ampere/Ada series GPUs |
+| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | RDNA2/RDNA3/RDNA4 series GPUs. You can select one of the interfaces for communication by setting `AMD_IFACE=(KFD|PCI)`. See [AMD interfaces](#amd-interfaces) for more details. |
+| [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | 6xx series GPUs |
+| [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | M1+ Macs; Metal 3.0+ for `bfloat` support |
+| [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | NVIDIA GPU with CUDA support |
+| [OpenCL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | OpenCL 2.0 compatible device |
+| [CPU (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` |
+| [LLVM (LLVM IR)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | llvm libraries installed and findable |
+| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0). |

 ## Interoperability

@@ -70,12 +70,5 @@ AMD backend supports several interfaces for communicating with devices:

 * `KFD`: uses the amdgpu driver
 * `PCI`: uses the [AM driver](developer/am.md)
-* `USB`: USB3 interafce for asm24xx chips.

 You can force an interface by setting `AMD_IFACE` to one of these values. In the case of `AMD_IFACE=PCI`, this may unbind your GPU from the amdgpu driver.
-
-## NV Interfaces
-NV backend supports several interfaces for communicating with devices:
-
-* `NVK`: uses the nvidia driver
-* `PCI`: uses the [NV driver](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/support/nv/nvdev.py)
@@ -26,8 +26,8 @@ class Attention:
      start_pos = start_pos.val

    if HALF: x = x.half()
-    xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
-    xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
+    xqkv = self.c_attn(x)
+    xq, xk, xv = [xqkv.shrink((None, None, (i*self.dim, (i+1)*self.dim))).reshape(None, None, self.n_heads, self.head_dim) for i in range(3)]
    bsz, seqlen, _, _ = xq.shape

    # create kv cache
@@ -35,11 +35,11 @@ class Attention:
      self.cache_kv = Tensor.zeros(2, bsz, MAX_CONTEXT, self.n_heads, self.head_dim, dtype=x.dtype).contiguous().realize()

    # update the cache
-    self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()
+    self.cache_kv.shrink((None, None,(start_pos,start_pos+seqlen),None,None)).assign(Tensor.stack(xk, xv)).realize()

    if start_pos > 0:
-      keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
-      values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
+      keys = self.cache_kv[0].shrink((None, (0, start_pos+seqlen), None, None))
+      values = self.cache_kv[1].shrink((None, (0, start_pos+seqlen), None, None))
    else:
      keys = xk
      values = xv
@@ -64,7 +64,7 @@ class TransformerBlock:

  def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]):
    h = x + self.attn(self.ln_1(x), start_pos, mask).float()
-    return (h + self.mlp(self.ln_2(h))).contiguous()
+    return (h + self.mlp(self.ln_2(h)))

 class Transformer:
  def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):
@@ -229,8 +229,7 @@ def train_cifar():
    if getenv("RANDOM_CROP", 1):
      X = random_crop(X, crop_size=32)
    if getenv("RANDOM_FLIP", 1):
-      # NOTE: RANGEIFY=1 needs this contiguous or the X[perms] is very slow
-      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X).contiguous() # flip LR
+      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X) # flip LR
    X, Y = X[perms], Y[perms]
    return X, Y, *cutmix(X, Y, perms, mask_size=hyp['net']['cutmix_size'])

@@ -279,15 +279,9 @@ def generate(model, tokenizer, prompt: str, n_tokens_to_gen: int = 10, temp: boo
  # Loading in the prompt tokens
  logits = model.forward(Tensor([tks]))[:, -1, :]
  for _ in tqdm(range(n_tokens_to_gen), desc="Speed Gen"):
+    # TODO: topk
    if sample:
-      scaled_logits = logits / temp
-      if top_k is not None:
-        topk_values, topk_indices = scaled_logits.topk(top_k)
-        filtered_logits = Tensor.full_like(scaled_logits, -float("inf"))
-        filtered_logits = filtered_logits.scatter(dim=-1, index=topk_indices, src=topk_values)
-        tok_Tens = filtered_logits.softmax().multinomial()
-      else:
-        tok_Tens = scaled_logits.softmax().multinomial()
+      tok_Tens = (logits/temp).softmax().multinomial()
    else:
      tok_Tens = logits.argmax(axis=-1).unsqueeze(0)
    tok = tok_Tens.item()
@@ -304,7 +298,6 @@ if __name__ == "__main__":
  parser.add_argument("--size", type=str, default="370m",
                      help=f"Size of model to use [{', '.join([k for k in MODELS.keys()])}]")
  parser.add_argument("--n_tokens", type=int, default=10, help="Number of tokens to generate")
-  parser.add_argument("--top_k", type=int, help="Limit sampling to the top k most likely tokens")
  parser.add_argument("--sample", dest="sample", action="store_true", help="Sample flag")
  parser.add_argument("--temp", type=float, default=1.0, help="Sampling temp has to be <=1.0")
  args = parser.parse_args()
@@ -315,9 +308,8 @@ if __name__ == "__main__":
  num_toks = args.n_tokens
  sample = args.sample
  temp = args.temp
-  top_k = args.top_k
  s = time.time()
-  tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp, top_k=top_k)
+  tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp)
  print(tinyoutput)
  print('TIME: ', time.time() - s)
  TORCHOUTPUT = "Why is gravity \nso important?\nBecause it's the only"
@@ -17,10 +17,6 @@ def he_normal(*shape, a: float = 0.00, **kwargs) -> Tensor:
  std = math.sqrt(2.0 / (1 + a ** 2)) / math.sqrt(prod(argfix(*shape)[1:])) / 0.87962566103423978
  return std * rand_truncn(*shape, **kwargs)

-# Stable Diffusion v2 training uses default torch gelu, which doesn't use tanh approximation
-def gelu_erf(x:Tensor) -> Tensor:
-  return 0.5 * x * (1.0 + (x / 1.4142135623730951).erf())
-
 class Conv2dHeNormal(nn.Conv2d):
  def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
    super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
@@ -1,57 +0,0 @@
-#!/usr/bin/env bash
-# adapted from https://github.com/mlcommons/training/blob/4bdf5c8ed218ad76565a2ba1ac27c919ccc6d689/stable_diffusion/README.md
-
-# setup dirs
-
-DATA=/raid/datasets/stable_diffusion
-
-LAION=$DATA/laion-400m/webdataset-moments-filtered 
-COCO=$DATA/coco2014
-mkdir -p $LAION $COCO
-
-CKPT=/raid/weights/stable_diffusion
-mkdir -p $CKPT/clip $CKPT/sd $CKPT/inception
-
-# download data
-
-# if rclone isn't installed system-wide / in your PATH, put the executable path in quotes below
-#RCLONE=""
-RCLONE="rclone"
-
-## VAE-encoded image latents, from 6.1M image subset of laion-400m
-## about 1 TB for whole download
-$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
-$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/ ${LAION} --include="*.tar" -P
-$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/sha512sums.txt ${LAION} -P
-cd $LAION && grep -E '\.tar$' sha512sums.txt | sha512sum -c --quiet - && \
-  echo "All .tar files verified" || { echo "Checksum failure when validating downloaded Laion moments"; exit 1; }
-
-## prompts and FID statistics from 30k image subset of coco2014
-## 33 MB
-$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
-$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k.tsv ${COCO} -P
-
-$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
-$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k_stats.npz ${COCO} -P
-
-# download checkpoints
-
-## clip (needed for text and vision encoders for validation)
-CLIP_WEIGHTS_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin"
-CLIP_WEIGHTS_SHA256="9a78ef8e8c73fd0df621682e7a8e8eb36c6916cb3c16b291a082ecd52ab79cc4"
-CLIP_CONFIG_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/raw/main/open_clip_config.json"
-wget -N -P ${CKPT}/clip ${CLIP_WEIGHTS_URL}
-wget -N -P ${CKPT}/clip ${CLIP_CONFIG_URL}
-echo "${CLIP_WEIGHTS_SHA256}  ${CKPT}/clip/open_clip_pytorch_model.bin"                    | sha256sum -c
-
-## sd (needed for latent->image decoder for validation, also has clip text encoder for training)
-SD_WEIGHTS_URL='https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt'
-SD_WEIGHTS_SHA256="d635794c1fedfdfa261e065370bea59c651fc9bfa65dc6d67ad29e11869a1824"
-wget -N -P ${CKPT}/sd ${SD_WEIGHTS_URL}
-echo "${SD_WEIGHTS_SHA256}  ${CKPT}/sd/512-base-ema.ckpt"                    | sha256sum -c
-
-## inception (needed for validation)
-FID_WEIGHTS_URL='https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
-FID_WEIGHTS_SHA1="bd836944fd6db519dfd8d924aa457f5b3c8357ff"
-wget -N -P ${CKPT}/inception ${FID_WEIGHTS_URL}
-echo "${FID_WEIGHTS_SHA1}  ${CKPT}/inception/pt_inception-2015-12-05-6726825d.pth"                    | sha1sum -c
@@ -437,8 +437,8 @@ if __name__ == "__main__":
    im.show()

  # validation!
-  is_default = args.prompt == default_prompt and args.steps == 10 and args.seed == 0 and args.guidance == 6.0 and args.width == args.height == 1024
-  if is_default and not args.weights and not args.fakeweights:
+  if args.prompt == default_prompt and args.steps == 10 and args.seed == 0 and args.guidance == 6.0 and args.width == args.height == 1024 \
+    and not args.weights:
    ref_image = Tensor(np.array(Image.open(Path(__file__).parent / "sdxl_seed0.png")))
    distance = (((x.cast(dtypes.float) - ref_image.cast(dtypes.float)) / ref_image.max())**2).mean().item()
    assert distance < 4e-3, colored(f"validation failed with {distance=}", "red")
@@ -109,7 +109,7 @@ class TextDecoder:

  def forward(self, x:Tensor, pos:Union[Variable, Literal[0]], encoded_audio:Tensor):
    seqlen = x.shape[-1]
-    x = self.token_embedding(x) + self.positional_embedding.shrink(((pos, pos+seqlen), None))
+    x = self.token_embedding(x) + self.positional_embedding.shrink(((pos, pos+seqlen), None, None))
    for block in self.blocks: x = block(x, xa=encoded_audio, mask=self.mask, len=pos)
    return self.output_tok(x)

@@ -9,9 +9,6 @@ from PIL import Image
 import numpy as np
 import re, gzip

-# Allow for monkeypatching for mlperf.
-gelu = Tensor.gelu
-
@lru_cache()
 def default_bpe():
  # Clip tokenizer, taken from https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py (MIT license)
@@ -56,8 +53,8 @@ class Tokenizer:
    cs = [chr(n) for n in cs]
    return dict(zip(bs, cs))
  class ClipTokenizer:
-    def __init__(self, version=None):
-      self.byte_encoder, self.version = Tokenizer.bytes_to_unicode(), version
+    def __init__(self):
+      self.byte_encoder = Tokenizer.bytes_to_unicode()
      merges = gzip.open(default_bpe()).read().decode("utf-8").split('\n')
      merges = merges[1:49152-256-2+1]
      merges = [tuple(merge.split()) for merge in merges]
@@ -65,17 +62,11 @@ class Tokenizer:
      vocab = vocab + [v+'</w>' for v in vocab]
      for merge in merges:
        vocab.append(''.join(merge))
-      if self.version == "sd_mlperf_v5_0":
-        import regex
-        vocab.extend(['<start_of_text>', '<end_of_text>'])
-        self.cache = {'<start_of_text>': '<start_of_text>', '<end_of_text>': '<end_of_text>'}
-        self.pat = regex.compile(r"""<start_of_text>|<end_of_text>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", regex.IGNORECASE)
-      else:
-        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
-        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
-        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[^\s]+""", re.IGNORECASE)
+      vocab.extend(['<|startoftext|>', '<|endoftext|>'])
      self.encoder = dict(zip(vocab, range(len(vocab))))
      self.bpe_ranks = dict(zip(merges, range(len(merges))))
+      self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+      self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[^\s]+""", re.IGNORECASE)

    def bpe(self, token):
      if token in self.cache:
@@ -119,17 +110,8 @@ class Tokenizer:

    def encode(self, text:str, pad_with_zeros:bool=False) -> List[int]:
      bpe_tokens: List[int] = []
-      if self.version == "sd_mlperf_v5_0":
-        import regex, ftfy, html
-        text = ftfy.fix_text(text)
-        text = html.unescape(html.unescape(text)).strip()
-        text = Tokenizer.whitespace_clean(text).lower()
-        re_module = regex
-      else:
-        text = Tokenizer.whitespace_clean(text.strip()).lower()
-        re_module = re
-
-      for token in re_module.findall(self.pat, text):
+      text = Tokenizer.whitespace_clean(text.strip()).lower()
+      for token in re.findall(self.pat, text):
        token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
        bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
      # Truncation, keeping two slots for start and end tokens.
@@ -270,8 +252,10 @@ class Open:
      q,k,v = [y.reshape(T, B*self.n_heads, self.d_head).transpose(0, 1).reshape(B, self.n_heads, T, self.d_head) for y in proj.chunk(3)]

      attn_output = Tensor.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
-      attn_output = attn_output.permute(2, 0, 1, 3).reshape(T, B, C)
+      attn_output = attn_output.permute(2, 0, 1, 3).reshape(T*B, C)
+
      attn_output = self.out_proj(attn_output)
+      attn_output = attn_output.reshape(T, B, C)

      return attn_output

@@ -279,10 +263,9 @@ class Open:
    def __init__(self, dims, hidden_dims):
      self.c_fc   = Linear(dims, hidden_dims)
      self.c_proj = Linear(hidden_dims, dims)
-      self.gelu = gelu

    def __call__(self, x:Tensor) -> Tensor:
-      return x.sequential([self.c_fc, self.gelu, self.c_proj])
+      return x.sequential([self.c_fc, Tensor.gelu, self.c_proj])

  # https://github.com/mlfoundations/open_clip/blob/58e4e39aaabc6040839b0d2a7e8bf20979e4558a/src/open_clip/transformer.py#L210
  class ResidualAttentionBlock:
@@ -367,15 +350,15 @@ class Open:
 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/encoders/modules.py#L396
 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/encoders/modules.py#L498
 class FrozenOpenClipEmbedder(Embedder):
-  def __init__(self, dims:int, n_heads:int, layers:int, return_pooled:bool, ln_penultimate:bool=False, clip_tokenizer_version=None):
-    self.tokenizer = Tokenizer.ClipTokenizer(version=clip_tokenizer_version)
+  def __init__(self, dims:int, n_heads:int, layers:int, return_pooled:bool, ln_penultimate:bool=False):
+    self.tokenizer = Tokenizer.ClipTokenizer()
    self.model = Open.ClipTextTransformer(dims, n_heads, layers)
    self.return_pooled = return_pooled
    self.input_key = "txt"
    self.ln_penultimate = ln_penultimate

  def tokenize(self, text:str, device:Optional[str]=None) -> Tensor:
-    return Tensor(self.tokenizer.encode(text, pad_with_zeros=True), dtype=dtypes.int32, device=device).reshape(1,-1)
+    return Tensor(self.tokenizer.encode(text, pad_with_zeros=True), dtype=dtypes.int64, device=device).reshape(1,-1)

  def text_transformer_forward(self, x:Tensor, attn_mask:Optional[Tensor]=None):
    for r in self.model.transformer.resblocks:
@@ -466,7 +449,7 @@ class OpenClipEncoder:
    x = x + self.positional_embedding
    x = self.transformer(x, attn_mask=self.attn_mask)
    x = self.ln_final(x)
-    x = x[Tensor.arange(x.shape[0], device=x.device), tokens.argmax(axis=-1)]
+    x = x[:, tokens.argmax(axis=-1)]
    x = x @ self.text_projection
    return x

@@ -270,10 +270,8 @@ class FidInceptionV3:
    self.Mixed_7b = inception.Mixed_7b
    self.Mixed_7c = inception.Mixed_7c

-  def load_from_pretrained(self, path=None):
-    if path is None:
-      path = fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")
-    state_dict = torch_load(str(path))
+  def load_from_pretrained(self):
+    state_dict = torch_load(str(fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")))
    for k,v in state_dict.items():
      if k.endswith(".num_batches_tracked"):
        state_dict[k] = v.reshape(1)
@@ -50,7 +50,7 @@ class TestBeamSearch(unittest.TestCase):
  def test_variable_shrink_prime_number(self):
    v = Variable("v", 1, 400).bind(367)
    a = rand(400, 367)
-    b = (a.shrink(((0,v), None))+1)[:367,:367].realize()
+    b = (a.shrink(((0,v), None))+1).reshape(367,367).realize()
    np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)

  def test_no_mutate_rawbuffers(self):
@@ -930,7 +930,7 @@ impl<'a> Thread<'a> {

            let op = ((instr >> 16) & 0x3ff) as u32;
            match op {
-                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 => {
+                764 | 765 | 288 | 289 | 290 | 766 | 768 | 769 => {
                    let vdst = (instr & 0xff) as usize;
                    let sdst = ((instr >> 8) & 0x7f) as usize;
                    let f = |i: u32| -> usize { ((instr >> i) & 0x1ff) as usize };
@@ -944,16 +944,6 @@ impl<'a> Thread<'a> {
                    assert_eq!(clmp, 0);

                    let vcc = match op {
-                        767 => {
-                            let (s0, s1, s2): (u32, u32, u64) = (self.val(s0), self.val(s1), self.val(s2));
-                            let (mul_result, overflow_mul) = (s0 as i64).overflowing_mul(s1 as i64);
-                            let (ret, overflow_add) = mul_result.overflowing_add(s2 as i64);
-                            let overflowed = overflow_mul || overflow_add;
-                            if self.exec.read() {
-                                self.vec_reg.write64(vdst, ret as u64);
-                            }
-                            overflowed
-                        },
                        766 => {
                            let (s0, s1, s2): (u32, u32, u64) = (self.val(s0), self.val(s1), self.val(s2));
                            let (mul_result, overflow_mul) = (s0 as u64).overflowing_mul(s1 as u64);
@@ -4,7 +4,7 @@

 Only supported on 7900XTX, requires either AM (`rmmod amdgpu`) or disabling power gating on AMD (`ppfeaturemask=0xffff3fff`, don't forget to rebuild initramfs)

-SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.
+SQTT is implemented on top of normal tinygrad PROFILE=1, `PROFILE=1 SQTT=1` to get profile pickle with sqtt data embedded in it.

 `SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.

@@ -35,7 +35,7 @@ def to_movement_ops(st: ShapeTracker) -> List[Tuple[MovementOps, Tuple]]:
  to_apply:List[Tuple[MovementOps, Tuple]] = []
  for i, v in enumerate(st.views):
    real_shape = tuple(y-x for x,y in v.mask) if v.mask else v.shape
-    offset = (v.offset or 0) + sum(st*(s-1) for s,st in zip(real_shape, v.strides) if st<0)
+    offset = v.offset + sum(st*(s-1) for s,st in zip(real_shape, v.strides) if st<0)
    real_offset = offset + (sum(x*st for (x,_),st in zip(v.mask, v.strides)) if v.mask else 0)
    real_real_shape = [s for s,st in zip(real_shape, v.strides) if st]
    strides: List[int] = [abs(st) if isinstance(st,int) else st for st in v.strides if st]
@@ -177,28 +177,22 @@ def cached_to_movement_ops(shape, st) -> list:

 from tinygrad.shape.shapetracker import ShapeTracker, View
 from extra.to_movement_ops import to_movement_ops, apply_mop, MovementOps
-
-@wrap_view_op
-def _as_strided(tensor:Tensor, size, stride, storage_offset=None):
-  # multiple as_strided do not compound
-  base = canonical_base(tensor)
-  # TODO: this is heavyweight
-  st = ShapeTracker(base.uop.st.views + (View.create(tuple(size), tuple(stride), storage_offset),))
-  ret = base
-  if TORCH_DEBUG >= 1: print("**** as_strided", tensor.shape, size, stride, st)
-  if prod(size) == 1: return ret.flatten()[storage_offset].reshape(size)
-  for mo in cached_to_movement_ops(tuple(base.shape), st): ret = apply_mop(ret, mo)
-  return ret
-
@torch.library.impl("aten::as_strided", "privateuseone")
 def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None):
  storage_offset = storage_offset or tensor.storage_offset()
+  @wrap_view_op
+  def _as_strided(tensor:Tensor, size, stride, storage_offset=None):
+    # multiple as_strided do not compound
+    base = canonical_base(tensor)
+    # TODO: this is heavyweight
+    st = ShapeTracker(base.uop.st.views + (View.create(tuple(size), tuple(stride), storage_offset),))
+    ret = base
+    if TORCH_DEBUG >= 1: print("**** as_strided", tensor.shape, size, stride, st)
+    if prod(size) == 1: return ret.flatten()[storage_offset].reshape(size)
+    for mo in cached_to_movement_ops(tuple(base.shape), st): ret = apply_mop(ret, mo)
+    return ret
  return _as_strided(tensor, size, stride, storage_offset)

-@torch.library.impl("aten::_reshape_alias", "privateuseone")
-def _reshape_alias(tensor:torch.Tensor, size, stride):
-  return _as_strided(tensor, size, stride)
-
@torch.library.impl("aten::empty_strided", "privateuseone")
 def empty_strided(size, stride, dtype, layout=None, device=None, pin_memory=False):
  if TORCH_DEBUG: print(f"empty_strided {size=} {stride=} {dtype=} {layout=} {device=} {pin_memory=}")
@@ -1,6 +1,2 @@
 [pytest]
 norecursedirs = extra
-timeout = 180
-timeout_method = thread
-timeout_func_only = true
-testpaths = test
@@ -9,12 +9,12 @@ with open(directory / 'README.md', encoding='utf-8') as f:

 testing_minimal = [
  "numpy",
-  "torch==2.8.0",
+  "torch==2.7.1",
  "pytest",
  "pytest-xdist",
-  "pytest-timeout",
  "hypothesis",
  "z3-solver",
+  "ml_dtypes"
 ]

 setup(name='tinygrad',
@@ -59,7 +59,7 @@ setup(name='tinygrad',
        'triton': ["triton-nightly>=2.1.0.dev20231014192330"],
        'linting': [
            "pylint",
-            "mypy==1.18.1",
+            "mypy==1.13.0",
            "typing-extensions",
            "pre-commit",
            "ruff",
@@ -1,13 +0,0 @@
-import unittest
-from tinygrad import dtypes, Device
-from tinygrad.device import is_dtype_supported
-
-@unittest.skipUnless(Device.DEFAULT=="NULL", "Don't run when testing non-NULL backends")
-class TestNULLSupportsDTypes(unittest.TestCase):
-  def test_null_supports_ints_floats_bool(self):
-    dts = dtypes.ints + dtypes.floats + (dtypes.bool,)
-    not_supported = [dt for dt in dts if not is_dtype_supported(dt, "NULL")]
-    self.assertFalse(not_supported, msg=f"expected these dtypes to be supported by NULL: {not_supported}")
-
-if __name__ == "__main__":
-  unittest.main()
@@ -4,7 +4,7 @@ import numpy as np
 import torch

 from tinygrad import GlobalCounters, Tensor, Device
-from tinygrad.helpers import getenv, Context, RANGEIFY
+from tinygrad.helpers import getenv, Context
 from tinygrad.nn.state import get_parameters
 from tinygrad.engine.realize import capturing
 from tinygrad.tensor import _to_np_dtype
@@ -106,7 +106,7 @@ class TestOptBinOp(unittest.TestCase):
  def test_no_binop_rerun(self): return self._test_no_binop_rerun(lambda a,b: a*b, lambda a,b: (a*b).reshape(16, 16, 1))
  def test_no_binop_rerun_alt(self): return self._test_no_binop_rerun(lambda a,b: (a*b).reshape(16, 16, 1), lambda a,b: a*b)
  def test_no_binop_rerun_reduce_broadcast(self):
-    return self._test_no_binop_rerun(lambda a,b: a.sum()+b, lambda a,b: a.sum().reshape(1,1)+b, allowed=1 if RANGEIFY else 2)
+    return self._test_no_binop_rerun(lambda a,b: a.sum()+b, lambda a,b: a.sum().reshape(1,1)+b, allowed=2)

  @unittest.skip("this test started failing with the new change, based movementop issue")
  def test_no_binop_rerun_transposed(self): return self._test_no_binop_rerun(lambda a,b: (a.T*b.T).T, lambda a,b: a*b)
@@ -164,7 +164,7 @@ class TestOpt(unittest.TestCase):

  def test_permute_was_pushed(self):
    a = Tensor.randn(16, 16, 16)
-    with CLCache(1 if RANGEIFY else 2):
+    with CLCache(2):
      c = a.sum(2)
      d = c.permute(1,0).contiguous()
      d.realize()
@@ -172,7 +172,7 @@ class TestOpt(unittest.TestCase):

  def test_permute_was_pushed_through_contract_reshape(self):
    a = Tensor.randn(4, 4, 4, 4, 4)
-    with CLCache(1 if RANGEIFY else 2):
+    with CLCache(2):
      c = a.sum(-1)
      d = c.reshape(16,16).permute(1,0).contiguous()
      d.realize()
@@ -180,7 +180,7 @@ class TestOpt(unittest.TestCase):

  def test_permute_was_pushed_through_contractw1s_reshape(self):
    a = Tensor.randn(4, 4, 4, 4, 4)
-    with CLCache(1 if RANGEIFY else 2):
+    with CLCache(2):
      c = a.sum(-1)
      d = c.reshape(16,1,16).permute(2,1,0).contiguous()
      d.realize()
@@ -188,7 +188,7 @@ class TestOpt(unittest.TestCase):

  def test_permute_was_pushed_through_expand_reshape(self):
    a = Tensor.randn(16, 16, 16)
-    with CLCache(1 if RANGEIFY else 2):
+    with CLCache(2):
      c = a.sum(2)
      d = c.reshape(4,4,4,4).permute(2,3,0,1).contiguous()
      d.realize()
@@ -20,7 +20,7 @@ class TestLLaMASpeed(unittest.TestCase):
  def test_llama_compile(self):
    backup_program = Device[Device.DEFAULT].runtime
    backup_allocator = Device[Device.DEFAULT].allocator
-    backup_compiler = Device[Device.DEFAULT].compiler.compile_cached
+    backup_compiler = Device[Device.DEFAULT].compiler
    Device[Device.DEFAULT].runtime = FakeProgram
    Device[Device.DEFAULT].allocator = FakeAllocator(Device.default)

@@ -44,14 +44,14 @@ class TestLLaMASpeed(unittest.TestCase):
    run_llama("codegen(1)")

    # test no compiler use for this
-    Device[Device.DEFAULT].compiler.compile_cached = None
+    Device[Device.DEFAULT].compiler = None
    run_llama("methodcache", False)
    with Profiling(sort='time', frac=0.1, fn="/tmp/llama.prof", ts=5):
      run_llama("profile", False)

    Device[Device.DEFAULT].runtime = backup_program
    Device[Device.DEFAULT].allocator = backup_allocator
-    Device[Device.DEFAULT].compiler.compile_cached = backup_compiler
+    Device[Device.DEFAULT].compiler = backup_compiler

 if __name__ == '__main__':
  TestLLaMASpeed().test_llama_compile()
@@ -63,7 +63,6 @@ if __name__ == "__main__":
    views_to_valid_uop.cache_clear()

    new_uops = uops_allocated()
-    print_uops()
    gc.collect()
    new_uops_gc = uops_allocated()
    print(f"{t.__name__:30s}: {new_uops:3d} -> {new_uops_gc:3d}")
@@ -11,7 +11,7 @@ if __name__ == "__main__":
  for i in range(10_000):
    if i % 1000 == 0:
      print(f"Progress: {i}")
-    dt = random.choice(dtypes.ints + tuple(dt.vec(4) for dt in dtypes.ints))
+    dt = random.choice(dtypes.ints)
    u = UOp.variable('x', random.randint(dt.min, 0), random.randint(1, dt.max), dtype=dt)
    d = random.randint(1, max(1, u.arg[2]))
    if d in powers_of_two: continue
@@ -1,53 +0,0 @@
-import unittest
-from tinygrad import Tensor, dtypes, Device
-from tinygrad.nn.state import get_parameters
-from extra.models import clip
-from examples.mlperf.initializers import gelu_erf
-Device.DEFAULT="NULL"
-GPUS = [f"NULL:{i}" for i in range(8)]
-
-clip_params = {"dims": 1024, "n_heads": 16, "layers": 24, "return_pooled": False, "ln_penultimate": True, "clip_tokenizer_version": "sd_mlperf_v5_0"}
-def get_cond_stage_model(GPUS:list[str]|None=None) -> clip.FrozenOpenClipEmbedder:
-  clip.gelu = gelu_erf
-  model = clip.FrozenOpenClipEmbedder(**clip_params)
-  if GPUS and len(GPUS) > 1:
-    for p in get_parameters(model): p.to_(GPUS)
-  return model
-def get_tokens(BS:int) -> Tensor: return Tensor([0] * 77 * BS, dtype=dtypes.int32).reshape(-1, 77)
-
-class TestOpenClip(unittest.TestCase):
-  def test_tokenizer(self):
-    prompt = "Beautiful is better than ugly.\nExplicit is better than implicit.\nSimple is better than complex.\nComplex is better than complicated."
-    model = get_cond_stage_model()
-    tokens = model.tokenizer.encode(prompt, pad_with_zeros=True)
-    expected = [49406, 1215, 533, 1539, 1126, 8159, 269, 33228, 533, 1539, 1126, 15269, 585, 269, 4129, 533, 1539, 1126, 6324, 269, 6324, 533,
-                1539, 1126, 16621, 269, 49407] + [0]*50
-    self.assertEqual(tokens, expected)
-
-  def test_clip_gelu_init(self):
-    for resblock in get_cond_stage_model().model.transformer.resblocks:
-      self.assertEqual(resblock.mlp.gelu, gelu_erf)
-
-  def test_multigpu_clip_embed(self):
-    BS = 304
-    model = get_cond_stage_model(GPUS)
-    tokens = get_tokens(BS)
-    embeds = model.embed_tokens(tokens.shard(GPUS, axis=0)).realize()
-    self.assertEqual(embeds.shape, (BS, 77, 1024))
-    self.assertEqual(embeds.dtype, dtypes.float32)
-
-  def test_multigpu_clip_score(self):
-    BS = 240
-    vision_cfg = {'width': 1280, 'layers': 32, 'd_head': 80, 'image_size': 224, 'patch_size': 14}
-    text_cfg = {'width': 1024, 'n_heads': 16, 'layers': 24, 'vocab_size': 49408, 'ctx_length': 77}
-    clip.gelu = gelu_erf
-    clip_encoder = clip.OpenClipEncoder(1024, text_cfg, vision_cfg)
-    for p in get_parameters(clip_encoder): p.to_(GPUS)
-    tokens = get_tokens(BS)
-    imgs = Tensor.zeros(BS,3,224,224).contiguous()
-    scores = clip_encoder.get_clip_score(tokens.shard(GPUS, axis=0), imgs.shard(GPUS, axis=0)).realize()
-    self.assertEqual(scores.shape, (BS,))
-    self.assertEqual(scores.dtype, dtypes.float32)
-
-if __name__=="__main__":
-  unittest.main()
@@ -1,14 +1,14 @@
 #!/usr/bin/env python
 import unittest
-from tinygrad import Tensor
 import numpy as np
+from tinygrad.tensor import Tensor
 import torch

 def get_question_samp(bsz, seq_len, vocab_size, seed):
  np.random.seed(seed)
  in_ids= np.random.randint(vocab_size, size=(bsz, seq_len))
  mask = np.random.choice([True, False], size=(bsz, seq_len))
-  seg_ids = np.random.randint(2, size=(bsz, seq_len))  # type_vocab_size
+  seg_ids = np.random.randint(1, size=(bsz, seq_len))
  return in_ids, mask, seg_ids

 def set_equal_weights(mdl, torch_mdl):
@@ -45,7 +45,7 @@ class TestBert(unittest.TestCase):

    seeds = (1337, 3141)
    bsz, seq_len = 1, 16
-    for seed in seeds:
+    for _, seed in enumerate(seeds):
      in_ids, mask, seg_ids = get_question_samp(bsz, seq_len, config['vocab_size'], seed)
      out = mdl(Tensor(in_ids), Tensor(mask), Tensor(seg_ids))
      torch_out = torch_mdl.forward(torch.from_numpy(in_ids).long(), torch.from_numpy(mask), torch.from_numpy(seg_ids).long())[:2]
@@ -1,10 +1,12 @@
-import ast, pathlib, unittest
+import ast
+import pathlib
+import unittest

 import numpy as np
 from PIL import Image

-from tinygrad import Tensor
-from tinygrad.helpers import getenv, CI
+from tinygrad.helpers import getenv
+from tinygrad.tensor import Tensor
 from extra.models.efficientnet import EfficientNet
 from extra.models.vit import ViT
 from extra.models.resnet import ResNet50
@@ -38,13 +40,19 @@ def preprocess(img, new=False):
    img /= np.array([0.229, 0.224, 0.225]).reshape((1, -1, 1, 1))
  return img

-def _infer(model: EfficientNet, img):
-  with Tensor.train(False):
-    out = model.forward(Tensor(img)).argmax(axis=-1)
-  return out.tolist()

-chicken_img = preprocess(Image.open(pathlib.Path(__file__).parent / 'efficientnet/Chicken.jpg'))
-car_img = preprocess(Image.open(pathlib.Path(__file__).parent / 'efficientnet/car.jpg'))
+def _infer(model: EfficientNet, img, bs=1):
+  old_training = Tensor.training
+  Tensor.training = False
+  img = preprocess(img)
+  # run the net
+  if bs > 1: img = img.repeat(bs, axis=0)
+  out = model.forward(Tensor(img))
+  Tensor.training = old_training
+  return _LABELS[np.argmax(out.numpy()[0])]
+
+chicken_img = Image.open(pathlib.Path(__file__).parent / 'efficientnet/Chicken.jpg')
+car_img = Image.open(pathlib.Path(__file__).parent / 'efficientnet/car.jpg')

 class TestEfficientNet(unittest.TestCase):
  @classmethod
@@ -56,20 +64,17 @@ class TestEfficientNet(unittest.TestCase):
  def tearDownClass(cls):
    del cls.model

-  @unittest.skipIf(CI, "covered by test_chicken_car")
  def test_chicken(self):
-    labels = _infer(self.model, chicken_img)
-    self.assertEqual(_LABELS[labels[0]], "hen")
+    label = _infer(self.model, chicken_img)
+    self.assertEqual(label, "hen")
+
+  def test_chicken_bigbatch(self):
+    label = _infer(self.model, chicken_img, 2)
+    self.assertEqual(label, "hen")

-  @unittest.skipIf(CI, "covered by test_chicken_car")
  def test_car(self):
-    labels = _infer(self.model, car_img)
-    self.assertEqual(_LABELS[labels[0]], "sports car, sport car")
-
-  def test_chicken_car(self):
-    labels = _infer(self.model, np.concat([chicken_img, car_img], axis=0))
-    self.assertEqual(_LABELS[labels[0]], "hen")
-    self.assertEqual(_LABELS[labels[1]], "sports car, sport car")
+    label = _infer(self.model, car_img)
+    self.assertEqual(label, "sports car, sport car")

 class TestViT(unittest.TestCase):
  @classmethod
@@ -82,12 +87,12 @@ class TestViT(unittest.TestCase):
    del cls.model

  def test_chicken(self):
-    labels = _infer(self.model, chicken_img)
-    self.assertEqual(_LABELS[labels[0]], "cock")
+    label = _infer(self.model, chicken_img)
+    self.assertEqual(label, "cock")

  def test_car(self):
-    labels = _infer(self.model, car_img)
-    self.assertEqual(_LABELS[labels[0]], "racer, race car, racing car")
+    label = _infer(self.model, car_img)
+    self.assertEqual(label, "racer, race car, racing car")

 class TestResNet(unittest.TestCase):
  @classmethod
@@ -100,12 +105,12 @@ class TestResNet(unittest.TestCase):
    del cls.model

  def test_chicken(self):
-    labels = _infer(self.model, chicken_img)
-    self.assertEqual(_LABELS[labels[0]], "hen")
+    label = _infer(self.model, chicken_img)
+    self.assertEqual(label, "hen")

  def test_car(self):
-    labels = _infer(self.model, car_img)
-    self.assertEqual(_LABELS[labels[0]], "sports car, sport car")
+    label = _infer(self.model, car_img)
+    self.assertEqual(label, "sports car, sport car")

 if __name__ == '__main__':
  unittest.main()
@@ -5,8 +5,12 @@ from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.device import Device
 from tinygrad.helpers import fetch, Context

-from extra.onnx_helpers import validate
-from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry
+try:
+  from extra.onnx_helpers import validate
+  from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry
+  HUGGINGFACE_AVAILABLE = True
+except ModuleNotFoundError:
+  HUGGINGFACE_AVAILABLE = False

 def run_onnx_torch(onnx_model, inputs):
  import torch
@@ -58,7 +62,7 @@ class TestOnnxModel(unittest.TestCase):
    print(cls, _LABELS[cls])
    assert "car" in _LABELS[cls] or _LABELS[cls] == "convertible"

-@unittest.skipUnless(Device.DEFAULT == "METAL", "only run on METAL")
+@unittest.skipUnless(HUGGINGFACE_AVAILABLE and Device.DEFAULT == "METAL", "only run on METAL")
 class TestHuggingFaceOnnxModels(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
@@ -53,8 +53,8 @@ class TestRealWorld(unittest.TestCase):
  @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
  def test_stable_diffusion(self):
    params = unet_params
-    params["model_ch"] = 8
-    params["ctx_dim"] = 8
+    params["model_ch"] = 16
+    params["ctx_dim"] = 16
    params["num_res_blocks"] = 1
    params["n_heads"] = 2
    model = UNetModel(**params)
@@ -114,16 +114,6 @@ class TestRealWorld(unittest.TestCase):

      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 93)

-  @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow")
-  def test_forward_cifar(self):
-    BS = 32
-    # with training batchnorm still though
-    with Tensor.train():
-      model = SpeedyResNet(Tensor.ones((12,3,2,2)))
-      @TinyJit
-      def run(X): return model(X)
-      helper_test("forward_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), run, (1.0/48)*BS, 126)
-
  @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow")
  def test_train_cifar(self):
    with Tensor.train():
@@ -154,7 +144,6 @@ class TestRealWorld(unittest.TestCase):
                                final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=4)
      assert not np.isnan(lr_scheduler.min_lr), "lr too small or initial_div_facotr too big for half"

-  @unittest.skipIf(CI and Device.DEFAULT == "CPU", "slow")
  def test_bert(self):
    with Tensor.train():
      args_tiny = {"attention_probs_dropout_prob": 0.0, "hidden_dropout_prob": 0.0, "vocab_size": 30522, "type_vocab_size": 2,
@@ -1,8 +1,8 @@
 #!/usr/bin/env python
 import unittest
-from tinygrad import Tensor
-from extra.models.rnnt import LSTM
 import numpy as np
+from tinygrad.tensor import Tensor
+from extra.models.rnnt import LSTM
 import torch

 class TestRNNT(unittest.TestCase):
@@ -1,8 +1,9 @@
-import unittest, time
+import unittest
+import time
 import numpy as np
-from tinygrad import Device
 from tinygrad.nn.state import get_parameters
 from tinygrad.nn import optim
+from tinygrad.tensor import Device
 from tinygrad.helpers import getenv, CI
 from extra.training import train
 from extra.models.convnext import ConvNeXt
@@ -39,6 +40,7 @@ class TestTrain(unittest.TestCase):
    check_gc()

  @unittest.skipIf(CI, "slow")
+  @unittest.skipIf(Device.DEFAULT in ["METAL", "WEBGPU"], "too many buffers for webgpu and metal")
  def test_efficientnet(self):
    model = EfficientNet(0)
    X = np.zeros((BS,3,224,224), dtype=np.float32)
@@ -47,6 +49,7 @@ class TestTrain(unittest.TestCase):
    check_gc()

  @unittest.skipIf(CI, "slow")
+  @unittest.skipIf(Device.DEFAULT in ["METAL", "WEBGPU"], "too many buffers for webgpu and metal")
  def test_vit(self):
    model = ViT()
    X = np.zeros((BS,3,224,224), dtype=np.float32)
@@ -54,7 +57,7 @@ class TestTrain(unittest.TestCase):
    train_one_step(model,X,Y)
    check_gc()

-  @unittest.skipIf(CI, "slow")
+  @unittest.skipIf(Device.DEFAULT in ["METAL", "WEBGPU"], "too many buffers for webgpu and metal")
  def test_transformer(self):
    # this should be small GPT-2, but the param count is wrong
    # (real ff_dim is 768*4)
@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Device, Tensor, dtypes
-from tinygrad.helpers import CI, RANGEIFY
+from tinygrad.helpers import CI
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError

 # TODO: write a clean version of this
@@ -351,18 +351,5 @@ class TestKernelOpts(unittest.TestCase):
    ] + [[Opt(OptOps.THREAD, 0, 4)] if Device[Device.DEFAULT].renderer.global_max[0] >= 4 else []]
      + [[Opt(OptOps.THREAD, 0, 8)] if Device[Device.DEFAULT].renderer.global_max[0] >= 8 else []])

-  @unittest.skipUnless(RANGEIFY>=1, "Kernel only fuses with rangeify")
-  def test_double_sum_group(self):
-    a = Tensor.rand(4, 4, 4)
-    r = a.sum((1, 2)).sum()
-    with self.assertRaises(KernelOptError):
-      helper_linearizer_opt(r, [[Opt(OptOps.GROUPTOP, 0, 16)],])
-    r = a.sum((1, 2)).sum()
-    with self.assertRaises(KernelOptError):
-      helper_linearizer_opt(r, [[Opt(OptOps.UNROLL, 1, 4), Opt(OptOps.GROUPTOP, 0, 16)],])
-    r = a.sum((1, 2)).sum()
-    with self.assertRaises(KernelOptError):
-      helper_linearizer_opt(r, [[Opt(OptOps.GROUPTOP, 1, 4), Opt(OptOps.GROUPTOP, 0, 16)],])
-
 if __name__ == '__main__':
  unittest.main()
@@ -1,29 +1,55 @@
 import unittest
 import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
-from tinygrad.helpers import CI, Context, getenv, RANGEIFY
+from tinygrad.helpers import CI, Context, getenv
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from tinygrad.uop.ops import Ops
+from tinygrad.renderer.ptx import PTXRenderer

 class TestArange(unittest.TestCase):
-  def _get_flops(self, N):
+  def _get_flops(self, N, opts=None):
    GlobalCounters.reset()
    tt = Tensor.arange(N)
    sched = tt.schedule()
    self.assertEqual(len(sched), 1)
-    p = get_program(sched[-1].ast)
+    p = get_program(sched[-1].ast, opts=opts)
+    print(p.name)
+    #print(p.src)
    ExecItem(CompiledRunner(p), [tt.uop.buffer]).run()
    np.testing.assert_equal(tt.numpy(), np.arange(N))
    return p.estimates.ops

-  def test_complexity(self):
-    self.assertEqual(self._get_flops(256), 0)
-    self.assertEqual(self._get_flops(2560), 0)
+  def test_complexity(self, opts=None, limit=None):
+    f1 = self._get_flops(256, opts)
+    f2 = self._get_flops(2560, opts)
+    print(f"{f1=}, {f2=}")
+    # add 1 to avoid divide by 0. arange is 0 flops now!
+    assert (f1 < 6000 and f2 < 6000) or ((f2+1) / (f1+1) < 16), f"bad complexity, flops {(f2+1) / (f1+1):.1f}X while inputs 10X"
+    if limit is not None and not isinstance(Device[Device.DEFAULT].renderer, PTXRenderer):
+      # PTX counts index ALU in flops
+      assert f1 <= limit, f"{f1=}, {limit=}"

-  def test_arange_cat(self):
-    t = Tensor.arange(2, dtype=dtypes.int)+Tensor([3])
-    self.assertEqual(t.cat(t).tolist(), [3, 4, 3, 4])
+  # reduce collapse now happens before optimizations
+  """
+  from tinygrad.codegen.opt import Opt, OptOps
+  def test_complexity_w_upcast(self): return self.test_complexity([Opt(OptOps.UPCAST, 0, 4)], limit=0)
+  def test_complexity_w_unroll2(self): return self.test_complexity([Opt(OptOps.UNROLL, 0, 2)], limit=0)
+  def test_complexity_w_unroll4(self): return self.test_complexity([Opt(OptOps.UNROLL, 0, 4)], limit=0)
+  def test_complexity_w_unroll8(self): return self.test_complexity([Opt(OptOps.UNROLL, 0, 8)], limit=0)
+  def test_complexity_w_upcast_and_unroll(self): return self.test_complexity([Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)], limit=0)
+
+  if Device.default.renderer.has_local:
+    # TODO: fix limit
+    def test_complexity_w_group(self): return self.test_complexity([Opt(OptOps.GROUP, 0, 16)], limit=81920)
+    def test_complexity_w_group_top(self): return self.test_complexity([Opt(OptOps.GROUPTOP, 0, 16)], limit=106496)
+
+    def test_complexity_w_local(self): return self.test_complexity([Opt(OptOps.LOCAL, 0, 16)], limit=0)
+    @unittest.skip("doesn't work yet. TODO: this absolutely should work")
+    def test_complexity_w_local_unroll4(self): return self.test_complexity([Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UNROLL, 0, 4)], limit=0)
+    @unittest.skip("doesn't work yet")
+    def test_complexity_w_local_and_padto(self): return self.test_complexity([Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.PADTO, axis=1, arg=32)])
+  """

 class TestRand(unittest.TestCase):
  def test_fused_rand_less_ops(self, noopt=1):
@@ -111,7 +137,7 @@ class TestIndexing(unittest.TestCase):
      X = dataset[idxs]
      assert X.shape == (4,DDIM)
      sched = X.schedule()
-      self.assertEqual(len(sched), 1 if RANGEIFY else 2)
+      self.assertEqual(len(sched), 2)
      run_schedule(sched)
      assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops} != {4*DSET}"
    np.testing.assert_allclose(real_index, X.numpy())
@@ -1,10 +1,9 @@
 #!/usr/bin/env python
 import unittest
-import contextlib
 import numpy as np
 from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import temp, RANGEIFY
+from tinygrad.helpers import temp

 N = 200  # has to be bigger than the cache to fail

@@ -255,8 +254,6 @@ class TestAssign(unittest.TestCase):
    b.assign(a.contiguous()).realize()
    assert GlobalCounters.kernel_count - kc == 2

-  # passing in RANGEIFY=1, RANGEIFY=0 asserts permuted assigns it can't fuse
-  def assert_permuted_assign(self): return self.assertRaisesRegex(RuntimeError, "contiguous") if not RANGEIFY else contextlib.nullcontext()
  def test_permuted_assignment(self):
    a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
    b = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
@@ -280,7 +277,7 @@ class TestAssign(unittest.TestCase):
    #GlobalCounters.cache = []
    ba1 = a.uop.base.realized # noqa: F841
    bb1 = b.uop.base.realized # noqa: F841
-    with self.assert_permuted_assign():
+    with self.assertRaisesRegex(RuntimeError, "contiguous"):
      a.assign(a.permute(1,0) + b)   # this should not work!
      a.realize()
      ba2 = a.uop.base.realized # noqa: F841
@@ -288,22 +285,6 @@ class TestAssign(unittest.TestCase):
      #assert ba1 == ba2 and ba1 != bb1
      np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))

-  @unittest.skipUnless(RANGEIFY, "only correct in rangeify")
-  def test_post_permuted_assignment_alt(self):
-    a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
-    b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
-    new_a = (a.T+b).numpy()
-    a.assign(a.T+b)
-    np.testing.assert_allclose(a.numpy(), new_a)
-
-  def test_post_reshape_assignment_fine(self):
-    a = Tensor.arange(N*N).reshape(N, N).contiguous().realize()
-    b = Tensor.arange(N*N).reshape(N, N).contiguous().realize()
-    rhs = a.reshape(-1).reshape(N, N)
-    new_a = (rhs+b).numpy()
-    a.assign(rhs+b)  # self-assign with reshape view is fine
-    np.testing.assert_allclose(a.numpy(), new_a)
-
  @unittest.skip("multi output not supported anymore")
  def test_simple_assignment_multioutput(self):
    a = Tensor.randn(32, 32).realize()
@@ -328,8 +309,8 @@ class TestAssign(unittest.TestCase):
  def test_permuted_assignment_correct(self):
    a = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
    b = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
-    # TODO: swizzler.py limitation, should NOT raise AssertionError from numpy.
-    with self.assert_permuted_assign():
+    # TODO: scheduler limitation, should NOT raise AssertionError from numpy.
+    with self.assertRaisesRegex(RuntimeError, "contiguous"):
      a = a.permute(1, 0)
      new_val = a + b
      a.assign(new_val)
@@ -338,11 +319,10 @@ class TestAssign(unittest.TestCase):
  def test_permuted_reduceop_child_dual_use(self):
    a = Tensor.randn(32, 32, 32).realize()
    b = Tensor.full((32, 32), 1.).contiguous().realize()
-    with self.assert_permuted_assign():
+    with self.assertRaisesRegex(RuntimeError, "contiguous"):
      r = a.sum(axis=1)
      b.assign(r + b.permute(1, 0))
      b.realize()
-      np.testing.assert_allclose(b.numpy(), a.numpy().sum(axis=1)+np.ones((32, 32)).transpose(1, 0), atol=1e-6, rtol=1e-3)

  @unittest.skip("multi output not supported anymore")
  def test_permuted_reduceop_multioutput_dual_use(self):
@@ -379,17 +359,15 @@ class TestAssign(unittest.TestCase):
    a.assign(a + b)
    kc = GlobalCounters.kernel_count
    a.realize()
-    # rangeify makes two kernels
-    assert GlobalCounters.kernel_count - kc == (2 if RANGEIFY else 1)
+    assert GlobalCounters.kernel_count - kc == 1
    np.testing.assert_equal(a.numpy(), np.ones((4, 4))+np.pad(np.ones((4, 4))[:, 0:2], ((0, 0), (0, 2)), constant_values=2))

  def test_permuted_assignment_masked_view_not_contiguous(self):
    a = Tensor.ones(4, 4).contiguous().realize()
-    with self.assert_permuted_assign():
+    with self.assertRaisesRegex(RuntimeError, "contiguous"):
      b = a.shrink((None, (0, 2))).pad((None, (0, 2)), value=2).permute(1, 0)
      a.assign(a + b)
      a.realize()
-      self.assertListEqual(a.tolist(), [[2.,2.,2.,2.],[2.,2.,2.,2.],[3.,3.,3.,3.], [3.,3.,3.,3.]])

  # TODO: is there a way to sneak in a permute such that it returns the wrong answer?

@@ -3,7 +3,6 @@ from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import DType, ConstType
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen import full_rewrite_to_sink
-from tinygrad.helpers import RANGEIFY
 from tinygrad.device import is_dtype_supported
 import numpy as np
 from test.helpers import not_support_multi_device
@@ -156,7 +155,7 @@ class TestMovedConstFolding(unittest.TestCase):

  def test_add_padded_zero(self):
    # TODO: it's 1 now, this might be possible to fold
-    _check_ast_count(0 if RANGEIFY else 1, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(2).pad(((1, 1),)))
+    _check_ast_count(1, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(2).pad(((1, 1),)))

  def test_mul_shrunk_one(self):
    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(6).shrink(((1, 5),)))
@@ -245,7 +244,7 @@ class TestReduceOpsConstFolding(unittest.TestCase):
        t = Tensor.ones(16, dtype=dt).reshape(4, 4)
        assert t.sum().dtype == t.contiguous().sum().dtype

-@unittest.skipIf(not_support_multi_device() or RANGEIFY, "no multi, RANGEIFY doesn't support multi const folding")
+@unittest.skipIf(not_support_multi_device(), "no multi")
 class TestMultiConstFolding(unittest.TestCase):
  def test_multi_const_folding_literal(self):
    ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
@@ -4,12 +4,13 @@ import torch
 from typing import Any, List
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, DEBUG, CI
-from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype, truncate
+from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad import Device, Tensor, dtypes
 from hypothesis import assume, 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
+import ml_dtypes
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")

@@ -25,7 +26,6 @@ def get_available_cast_dtypes(dtype: DType) -> List[DType]:

 def _to_torch_storage_type(dtype:DType):
  if dtype == dtypes.bfloat16: return torch.float32
-  if dtype in dtypes.fp8s: return torch.float32
  return _to_torch_dtype(dtype)

 def _test_to_np(a:Tensor, np_dtype, target):
@@ -48,15 +48,12 @@ def _test_cast(a:Tensor, target_dtype:DType):
    # TODO: struct.pack cannot pack value > 65504 (max of half) into e format
    a = (a > 65504).where(65504, a)

-  expected = list(a.numpy().astype(_to_np_dtype(target_dtype)))
-  if target_dtype in dtypes.fp8s: expected = list(map(lambda x: truncate[target_dtype](x), expected))
-  _test_op(lambda: a.cast(target_dtype), target_dtype, expected)
+  _test_op(lambda: a.cast(target_dtype), target_dtype, list(a.numpy().astype(_to_np_dtype(target_dtype))))
 def _test_bitcast(a:Tensor, target_dtype:DType, target=None):
  if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and a.dtype == dtypes.int8 and target_dtype.itemsize != a.dtype.itemsize:
    raise unittest.SkipTest("shape changing bitcast of int8 broken on PTX")
-  expected = torch.tensor(a.tolist(), dtype=_to_torch_storage_type(a.dtype)).view(_to_torch_dtype(target_dtype)).tolist()
-  if target_dtype in dtypes.fp8s: expected = list(map(lambda x: fp8_to_float(x, target_dtype), expected))
-  _test_op(lambda: a.bitcast(target_dtype), target_dtype, target or expected)
+  expected = torch.tensor(a.tolist(), dtype=_to_torch_storage_type(a.dtype)).view(_to_torch_dtype(target_dtype))
+  _test_op(lambda: a.bitcast(target_dtype), target_dtype, target or expected.tolist())

 class TestDType(unittest.TestCase):
  DTYPE: Any = None
@@ -132,10 +129,11 @@ class TestDType(unittest.TestCase):
      np.testing.assert_allclose(tin, tor, atol=1e-6, rtol=1e-3)

  def test_finfo(self):
-    if self.DTYPE not in [dtypes.float16, dtypes.float32, dtypes.float64]: return
-    info = np.finfo(_to_np_dtype(self.DTYPE))
-    self.assertEqual(info.bits, self.DTYPE.itemsize*8)
-    self.assertEqual((info.nexp, info.nmant), dtypes.finfo(self.DTYPE))
+    if self.DTYPE not in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]: return
+    info = ml_dtypes.finfo(ml_dtypes.bfloat16 if self.DTYPE is dtypes.bfloat16 else _to_np_dtype(self.DTYPE))
+    assert info.bits == self.DTYPE.itemsize*8
+    assert info.nexp == dtypes.finfo(self.DTYPE)[0]
+    assert info.nmant == dtypes.finfo(self.DTYPE)[1]

 def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
  target_dtype = target_dtype or least_upper_dtype(a_dtype, b_dtype)
@@ -153,8 +151,7 @@ class TestFp8s(unittest.TestCase):

 class TestFp8sConversions(unittest.TestCase):
  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
-  def test_float_to_fp8e4m3(self, x):
-    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
+  def test_float_to_fp8e4m3(self, x): np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), ml_dtypes.float8_e4m3fn(x).tobytes()[0])

  def test_float_to_fp8e4m3_extreme_values(self):
    np.testing.assert_equal(float_to_fp8(FP8E4M3_MAX, dtypes.fp8e4m3), 126)
@@ -167,8 +164,7 @@ class TestFp8sConversions(unittest.TestCase):
    np.testing.assert_equal(float_to_fp8(-math.nan, dtypes.fp8e4m3), 255)

  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2_MAX, max_value=FP8E5M2_MAX))
-  def test_float_to_fp8e5m2(self, x):
-    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.float8_e5m2).view(torch.uint8).item())
+  def test_float_to_fp8e5m2(self, x): np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), ml_dtypes.float8_e5m2(x).tobytes()[0])

  def test_float_to_fp8e5m2_extreme_values(self):
    np.testing.assert_equal(float_to_fp8(FP8E5M2_MAX, dtypes.fp8e5m2), 123)
@@ -181,12 +177,10 @@ class TestFp8sConversions(unittest.TestCase):
    np.testing.assert_equal(float_to_fp8(-math.nan, dtypes.fp8e5m2), 254)

  @given(strat.integers(min_value=0, max_value=255))
-  def test_fp8e4m3_to_float(self, x):
-    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fn).float().item())
+  def test_fp8e4m3_to_float(self, x): np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3), np.uint8(x).view(ml_dtypes.float8_e4m3fn).item())

  @given(strat.integers(min_value=0, max_value=255))
-  def test_fp8e5m2_to_float(self, x):
-    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2).float().item())
+  def test_fp8e5m2_to_float(self, x): np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), np.uint8(x).view(ml_dtypes.float8_e5m2).item())

@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
 class TestBFloat16(unittest.TestCase):
@@ -312,8 +306,6 @@ class TestBitCast(unittest.TestCase):
    assume(not (isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and dt1 == dtypes.int8)) # TODO: bitcasting int8 fails in PTX
    data = rand_for_dtype(dt1, 32).reshape(2, 2, 8)
    expected = torch.tensor(data.tolist(), dtype=_to_torch_storage_type(dt1)).view(_to_torch_dtype(dt2))
-    if dt2 in dtypes.fp8s:
-      expected = torch.tensor(list(map(lambda x: fp8_to_float(x, dt2), expected.view(-1).tolist()))).view_as(expected)
    _test_op(lambda: Tensor(data, dtype=dt1).bitcast(dt2), dt2, expected.tolist())

  def test_shape_change_bitcast_exceptions(self):
@@ -356,9 +348,6 @@ class TestBoolDType(TestDType): DTYPE = dtypes.bool

 class TestBFloat16Type(TestDType): DTYPE = dtypes.bfloat16

-class TestFp8e4m3(TestDType): DTYPE = dtypes.fp8e4m3
-class TestFp8e5m2(TestDType): DTYPE = dtypes.fp8e5m2
-
 class TestPtrDType(unittest.TestCase):
  def test_vec_double(self):
    dt1 = dtypes.float.vec(4).ptr().vec(4)
@@ -1,6 +1,6 @@
 import unittest, operator, math
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.dtype import DType, truncate
+from tinygrad.dtype import DType
 from tinygrad.helpers import CI, getenv
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -8,7 +8,7 @@ from tinygrad.runtime.ops_python import from_storage_scalar
 from tinygrad.renderer.ptx import PTXRenderer
 import numpy as np
 import pytest
-from hypothesis import assume, given, strategies as strat, settings, HealthCheck
+from hypothesis import given, strategies as strat, settings, HealthCheck

 pytestmark = pytest.mark.filterwarnings("ignore")

@@ -48,8 +48,6 @@ class ht:
  int64 = strat.integers(-9223372036854775808, 9223372036854775807)
  bool = strat.booleans()
 ht.bfloat16 = ht.uint16
-ht.fp8e4m3 = ht.uint8
-ht.fp8e5m2 = ht.uint8

 def universal_test(a, b, dtype, op):
  if not isinstance(op, tuple): op = (op, op)
@@ -59,9 +57,8 @@ def universal_test(a, b, dtype, op):
  ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype)
  tensor_value = (op[0](ta, tb)).numpy()
  numpy_value = op[1](ta.numpy(), tb.numpy())
-  if dtype in dtypes.fp8s: numpy_value = truncate[dtype](numpy_value)
  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))
+    atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2)}.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)

@@ -74,10 +71,8 @@ def universal_test_unary(a, dtype, op):
  out: Tensor = op[0](ta)
  tensor_value = out.numpy()
  numpy_value = op[1](ta.numpy())
-  if dtype in dtypes.fp8s: numpy_value = truncate[dtype](numpy_value)
  if dtype in dtypes.floats:
-    atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2),
-      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1)}.get(dtype, (1e-6, 1e-5))
+    atol, rtol = {dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2)}.get(dtype, (1e-6, 1e-5))
    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
  else: np.testing.assert_equal(tensor_value, numpy_value)

@@ -116,16 +111,6 @@ class TestDTypeALU(unittest.TestCase):
  def test_bfloat16(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)

-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), f"no fp8e4m3 on {Device.DEFAULT}")
-  @given(ht.fp8e4m3, ht.fp8e4m3, strat.sampled_from(binary_operations))
-  def test_fp8e4m3(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2), f"no fp8e5m2 on {Device.DEFAULT}")
-  @given(ht.fp8e5m2, ht.fp8e5m2, strat.sampled_from(binary_operations))
-  def test_fp8e5m2(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)
-
  @given(ht.float32, strat.sampled_from(unary_operations))
  def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)

@@ -137,18 +122,6 @@ class TestDTypeALU(unittest.TestCase):
  @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)

-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), f"no fp8e4m3 on {Device.DEFAULT}")
-  @given(ht.fp8e4m3, strat.sampled_from(unary_operations))
-  def test_fp8e4m3_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3), dtypes.fp8e4m3, op)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2), f"no fp8e5m2 on {Device.DEFAULT}")
-  @given(ht.fp8e5m2, strat.sampled_from(unary_operations))
-  def test_fp8e5m2_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)
-
  @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
  def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)

@@ -4,7 +4,7 @@ from tinygrad import Device, dtypes, Tensor, Context
 from tinygrad.device import LRUAllocator, is_dtype_supported
 from tinygrad.dtype import ImageDType
 from tinygrad.engine.realize import lower_schedule
-from tinygrad.helpers import prod, unwrap, RANGEIFY
+from tinygrad.helpers import prod, unwrap
 from test.helpers import REAL_DEV

 IMAGE_SUPPORTED_DEVICES = ("QCOM", "CL")
@@ -139,7 +139,7 @@ class TestImageDType(unittest.TestCase):
      # NOTE: the w1 grad must realize to a seperate kernel
      assert w1.grad.uop.is_realized, f"never realized {w1.grad}"
      self.assertEqual(w1.grad.uop.base.buffer.dtype, dtypes.float32)
-      self.assertEqual(len(sched), 8 if RANGEIFY else 10)
+      self.assertEqual(len(sched), 10)

@unittest.skipUnless(REAL_DEV in IMAGE_SUPPORTED_DEVICES, "Images not supported")
 class TestImageRealization(unittest.TestCase):
@@ -609,22 +609,21 @@ class TestJitFree(unittest.TestCase):
    ext_tensor = Tensor([1,24,23,45,1])
    @TinyJit
    def fxn(x:Tensor):
-      t1 = (x * 2).contiguous().realize()
-      t2 = (t1 + ext_tensor).contiguous().realize()
-      out = (t2.sum()).contiguous().realize()
-      return out
+      out = (x*2+ext_tensor).reshape(5,1).expand(5, 100).contiguous()
+      return out.sum()
    for i in range(5):
-      out = fxn(inp:=Tensor([i,1,2,3,4]))
-      self.assertEqual(out.item(), 114+2*i)
+      out = fxn(Tensor([i,1,2,3,4]))
+      self.assertEqual(out.item(), 11400+200*i)
    pre_free = GlobalCounters.mem_used
    fxn.captured.free_intermediates()
    savings_after_free = pre_free - GlobalCounters.mem_used

-    expected_savings = (len(inp) * inp.dtype.itemsize * 2) + dtypes.float32.itemsize # (t1 and t2) + out
+    # Different allocator implementations have different savings.
+    expected_savings = 8196 if hasattr(Device[Device.DEFAULT].allocator, '_offset') else 2024

    self.assertEqual(savings_after_free, expected_savings)
    out = fxn(Tensor([11,1,2,3,4]))
-    self.assertEqual(out.item(), 136)
+    self.assertEqual(out.item(), 13600)

    # Try one more time...
    pre_free = GlobalCounters.mem_used
@@ -634,7 +633,7 @@ class TestJitFree(unittest.TestCase):

    self.assertEqual(savings_after_free, expected_savings)
    out = fxn(Tensor([11,1,2,3,4]))
-    self.assertEqual(out.item(), 136)
+    self.assertEqual(out.item(), 13600)

  def test_updated_not_freed(self):
    x = Tensor([1]).realize()
@@ -833,20 +832,5 @@ class TestJitGraphSplit(unittest.TestCase):
      multigraph=[self.ji_graph(2), self.ji_copy(), self.ji_comp()],
      hcqgraph=[self.ji_graph(4)])

-class TestJitRandom(unittest.TestCase):
-  def test_jit_rangeify(self):
-    tst = {0:[], 1:[]}
-    for r in [0,1]:
-      Tensor.manual_seed(1337)
-      with Context(RANGEIFY=r):
-        _ = Tensor.randint(4, high=3)
-        # this second one makes the behavior different
-        _ = Tensor.randint(4, high=3)
-        @TinyJit
-        def f(): return Tensor.randint(20, high=5)
-        for _ in range(5): tst[r].append(f().tolist())
-    for i, (t0, t1) in enumerate(zip(tst[0], tst[1])):
-      self.assertListEqual(t0, t1, msg=f"mismatch at list {i}")
-
 if __name__ == '__main__':
  unittest.main()
@@ -16,14 +16,14 @@ class TestKernelCache(unittest.TestCase):

    a1 = Tensor.rand(4,4).realize()
    b1 = Tensor.rand(4,4).realize()
-    orig_compile_func = Device['CPU'].compiler.compile_cached
-    Device['CPU'].compiler.compile_cached = None # making it not callable
+    orig_compile_func = Device['CPU'].compiler
+    Device['CPU'].compiler = None # making it not callable

    try:
      x1 = a1 + b1 + unique_const
      x1.realize() # Same kernel should be from cache.
    finally:
-      Device['CPU'].compiler.compile_cached = orig_compile_func
+      Device['CPU'].compiler = orig_compile_func

 if __name__ == "__main__":
  unittest.main()
@@ -123,7 +123,6 @@ class TestLinearizer(unittest.TestCase):
    assert num_loads <= 4, "more load uops than needed"
    assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"

-  @unittest.skip("this is handled at higher level now")
  def test_upcast_cse(self):
    # when upcasting, within a subtree, there may be common expressions.

@@ -483,7 +482,7 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
  assert s[-1].ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {s[-1]}"
  # now all input buffers in s[-1] should be realized
  # create fresh buffers for the outputs
-  bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)]
+  bufs = [Buffer((x).device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)]
  return push_views(s[-1].ast), bufs

 def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
@@ -505,7 +504,7 @@ def reset_bufs(bufs:list[Buffer]):

 def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
                               apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
-  outbufs = real_bufs[:len(realized_ast.src)]
+  outbufs = [real_bufs[x.src[0].base.arg] for x in realized_ast.src]
  device = real_bufs[0].device
  wanna_output = [np.array(x).flatten() for x in wanna_output]

@@ -12,21 +12,22 @@ from tinygrad.engine.realize import get_program
 from tinygrad.renderer.ptx import PTXRenderer

 class TestLinearizerFailure(unittest.TestCase):
+  @unittest.expectedFailure
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
  def test_failure_beam_mnist(self):
    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(4014080), arg=0, src=())
-    c1 = UOp.range(UOp.const(dtypes.index, 512), 0, AxisType.GLOBAL)
-    c2 = UOp.range(UOp.const(dtypes.index, 784), 1, AxisType.GLOBAL)
-    c3 = UOp.range(UOp.const(dtypes.index, 10), 3, AxisType.GLOBAL)
+    c1 = UOp.range(UOp.const(dtypes.int, 512), 0, AxisType.GLOBAL)
+    c2 = UOp.range(UOp.const(dtypes.int, 784), 1, AxisType.GLOBAL)
+    c3 = UOp.range(UOp.const(dtypes.int, 10), 3, AxisType.GLOBAL)
    c4 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
-    c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True))).load()
-    c6 = UOp.range(UOp.const(dtypes.index, 6000), 1004, AxisType.REDUCE)
-    c7 = UOp.range(UOp.const(dtypes.index, 3750), 2006, AxisType.REDUCE)
-    c8 = UOp.range(UOp.const(dtypes.index, 16), 2007, AxisType.GROUP_REDUCE)
+    c5 = c4.index(c1, UOp.const(dtypes.bool, True)).load()
+    c6 = UOp.range(UOp.const(dtypes.int, 6000), 1004, AxisType.REDUCE)
+    c7 = UOp.range(UOp.const(dtypes.int, 3750), 2006, AxisType.REDUCE)
+    c8 = UOp.range(UOp.const(dtypes.int, 16), 2007, AxisType.GROUP_REDUCE)
    c9 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(47040000), arg=2, src=())
-    c10 = c9.index((((c3*UOp.const(dtypes.index, 4704000))+c2)+(c6*UOp.const(dtypes.index, 784))).valid(UOp.const(dtypes.bool, True))).load()
-    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.index, 6000))+c6)+((c7*UOp.const(dtypes.index, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.index, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
-    c12 = c0.index((((c1*UOp.const(dtypes.index, 7840))+(c2*UOp.const(dtypes.index, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11, c1, c2, c3)
+    c10 = c9.index((((c3*UOp.const(dtypes.int, 4704000))+c2)+(c6*UOp.const(dtypes.int, 784))), UOp.const(dtypes.bool, True)).load()
+    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.int, 6000))+c6)+((c7*UOp.const(dtypes.int, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.int, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
+    c12 = c0.index((((c1*UOp.const(dtypes.int, 7840))+(c2*UOp.const(dtypes.int, 10)))+c3), UOp.const(dtypes.bool, True)).store(c11, c1, c2, c3)
    ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
    _ = get_program(ast, Device["METAL"].renderer)

@@ -5,9 +5,9 @@ from tinygrad.nn.state import get_state_dict

 class TestMethodCache(unittest.TestCase):
  def setUp(self):
-    self.backup_compiler = Device[Device.DEFAULT].compiler.compile_cached
+    self.backup_compiler = Device[Device.DEFAULT].compiler
  def tearDown(self):
-    Device[Device.DEFAULT].compiler.compile_cached = self.backup_compiler
+    Device[Device.DEFAULT].compiler = self.backup_compiler

  def test_simple_methodcache(self):
    a = Tensor([1])
@@ -15,19 +15,19 @@ class TestMethodCache(unittest.TestCase):
    c = Tensor([3])
    d = Tensor([4])
    (a+b).realize()
-    Device[Device.DEFAULT].compiler.compile_cached = None
+    Device[Device.DEFAULT].compiler = None
    (c+d).realize()

  def test_nested_methodcache(self):
    a,b,c,d = Tensor([1]), Tensor([2]), Tensor([3]), Tensor([4])
    ((a+b)+(a+b)).realize()
-    Device[Device.DEFAULT].compiler.compile_cached = None
+    Device[Device.DEFAULT].compiler = None
    ((c+d)+(c+d)).realize()

  def test_nested_methodcache_swap(self):
    a,b,c,d = Tensor([1]), Tensor([2]), Tensor([3]), Tensor([4])
    ((a+b)+(c+d)).realize()
-    Device[Device.DEFAULT].compiler.compile_cached = None
+    Device[Device.DEFAULT].compiler = None
    ((c+d)+(a+b)).realize()

  @unittest.skip("incorrect use of transformer")
@@ -38,7 +38,7 @@ class TestMethodCache(unittest.TestCase):
    # NOTE: you have to do this twice due to the k-v cache
    for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize()
    for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize()
-    Device[Device.DEFAULT].compiler.compile_cached = None
+    Device[Device.DEFAULT].compiler = None
    for i in range(3): model(Tensor([[1,2,3,4]]), Variable("start_pos", 0, 10).bind(i)).realize()

 if __name__ == '__main__':
@@ -2,7 +2,7 @@ import unittest, functools, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import Ops, UOp
-from tinygrad.helpers import CI, getenv, prod, Context, RANGEIFY
+from tinygrad.helpers import CI, getenv, prod, Context
 from tinygrad.nn.state import get_parameters, get_state_dict
 from tinygrad.engine.realize import lower_schedule, BufferCopy, CompiledRunner, run_schedule
 import numpy as np
@@ -54,17 +54,6 @@ class TestMultiTensor(unittest.TestCase):
      assert lb.shape == (128,)
    (X + X).realize()

-  def _test_shard_op(self, op, out, n=4):
-    t = Tensor.ones(n).contiguous().realize().shard(devices_2, 0)
-    r = op(t).realize()
-    assert t.uop.is_realized, "shard didn't realize"
-    self.assertEqual(r.tolist(), out)
-  def test_shard_reshape(self): self._test_shard_op(lambda t:t.reshape(2, 2), [[1.,1.],[1.,1.]])
-  def test_shard_elementwise(self): self._test_shard_op(lambda t:(t+t).reshape(2, 2), [[2.,2.],[2.,2.]])
-  def test_shard_reduce(self):
-    self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=1), [3.,3.], n=6)
-    self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=0), [2.,2.,2.], n=6)
-
  def test_shard_not_multiple(self):
    X = Tensor.ones(256).contiguous().realize()
    with self.assertRaises(RuntimeError):
@@ -383,7 +372,6 @@ class TestMultiTensor(unittest.TestCase):

  # NOTE: this is failing on LLVM CI, no idea why. Works locally.
  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "CPU", "AMD"), "slow, and flaky on CPU")
-  @unittest.skipIf(RANGEIFY, "TODO: pm_rangeify hangs")
  def test_data_parallel_resnet(self):
    from extra.models.resnet import ResNet18

@@ -420,7 +408,6 @@ class TestMultiTensor(unittest.TestCase):
    np.testing.assert_allclose(grad, shard_grad, atol=1e-5, rtol=1e-5)

  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "CPU", "AMD"), "slow, and flaky on CPU")
-  @unittest.skipIf(RANGEIFY, "TODO: pm_rangeify hangs")
  def test_data_parallel_resnet_train_step(self):
    from extra.models.resnet import ResNet18
    fake_image = Tensor.rand((2, 3, 224//16, 224//16))
@@ -428,7 +415,6 @@ class TestMultiTensor(unittest.TestCase):
    m = ResNet18()
    self._test_model_train_step(m, fake_image, labels)

-  @unittest.skipIf(RANGEIFY, "TODO: pm_rangeify hangs")
  def test_data_parallel_simple_train_step(self):
    class Model:
      def __init__(self): self.conv1 = nn.Linear(128,128)
@@ -793,7 +779,6 @@ class TestMultiTensor(unittest.TestCase):
    t = Tensor.rand(16, 16).shard(devices_2, axis=0)
    np.testing.assert_allclose(t.numpy(), t.clone().numpy())

-  @unittest.skipIf(RANGEIFY, "RANGEIFY doesn't support multi const folding")
  def test_multi_const_folding(self):
    with Context(TRACK_MATCH_STATS=0):
      a = Tensor.arange(3).realize()
@@ -229,8 +229,7 @@ class TestNN(unittest.TestCase):
    torch_z = torch_layer(torch_x)
    torch_z.sum().backward()

-    # TODO: why is torch numbers all 0?
-    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=5e-6)
+    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)

  def test_layernorm(self):
    N, C, H, W = 20, 5, 10, 10
@@ -333,7 +332,7 @@ class TestNN(unittest.TestCase):

      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
      np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
-      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=3e-3, rtol=1e-3)
+      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=2e-3, rtol=1e-3)
      np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3)

  def test_rmsnorm(self):
@@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, CPU_LLVM, AMD_LLVM, RANGEIFY
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, CPU_LLVM, AMD_LLVM
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -234,8 +234,7 @@ class TestOps(unittest.TestCase):
  def test_unfold(self):
    helper_test_op([(8,)], lambda x: x.unfold(0, 2, 1))
    helper_test_op([(8,)], lambda x: x.unfold(0, 2, 2))
-    # TODO: something is wrong with unfold
-    if not getenv("TINY_BACKEND"): helper_test_op([(8,)], lambda x: x.unfold(0, 7, 3))
+    helper_test_op([(8,)], lambda x: x.unfold(0, 7, 3))
    helper_test_op([(3,3,3)], lambda x: x.unfold(2, 2, 8))
    helper_test_op([(3,3,3)], lambda x: x.unfold(1, 0, 8))
    helper_test_op([(3,3,3,3,3)], lambda x: x.unfold(-1, 2, 2))
@@ -312,11 +311,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: torch.nn.functional.pad(torch.ones(256,256), pad=(0,64,0,0)).sum(axis=1),
                       lambda: Tensor.ones(256,256).pad(((0,0), (0,64))).sum(axis=1), forward_only=True)

-  def test_sum_twice(self):
-    helper_test_op([(4, 4, 4)], lambda x: x.sum((0, 1)).sum())
-    helper_test_op([(4, 4, 4)], lambda x: x.sum((0, 2)).sum())
-    helper_test_op([(4, 4, 4)], lambda x: x.sum((1, 2)).sum())
-
  # this is more complex and won't fold for a while
  def test_sum_cat_collapse(self):
    helper_test_op([], lambda: torch.cat([torch.ones(256,256), torch.zeros(256,64)], dim=1).sum(axis=1),
@@ -1413,11 +1407,6 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, lambda x: x.max(), forward_only=True, vals=[[False, True]])
    helper_test_op(None, lambda x: x.max(), forward_only=True, vals=[[True, False]])

-  def test_const_reduce(self):
-    helper_test_op([(3,3)], lambda x: torch.full_like(x, 2).sum(), lambda x: (x.full_like(2)).sum(), forward_only=True)
-    helper_test_op([(3,3)], lambda x: torch.full_like(x, 2).prod(), lambda x: (x.full_like(2)).prod(), forward_only=True)
-    helper_test_op([(3,3)], lambda x: torch.full_like(x, 2).max(), lambda x: (x.full_like(2)).max(), forward_only=True)
-
  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
  def test_any(self):
    helper_test_op([(3,4,5,6)], lambda x: x.any(), forward_only=True)
@@ -1469,7 +1458,6 @@ class TestOps(unittest.TestCase):
  def test_mean_zero_axis(self):
    helper_test_op([(1,0,3,0,5)], lambda x: x.mean(axis=(1,3)))

-  @slow_test
  def test_var(self):
    helper_test_op([(15, 25, 35)], lambda x: x.var())
    helper_test_op([(15, 25, 35)], lambda x: x.var(correction=0))
@@ -1505,7 +1493,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(15, 25, 35)], lambda x: x.var(keepdim=True))
    helper_test_op([(15, 25, 35)], lambda x: x.var(0, keepdim=True, correction=0))

-  @slow_test
  def test_std(self):
    helper_test_op([(15, 25, 35)], lambda x: x.std())
    helper_test_op([(15, 25, 35)], lambda x: x.std(correction=0))
@@ -1538,7 +1525,6 @@ class TestOps(unittest.TestCase):
  def test_std_keepdim(self):
    helper_test_op([(15, 25, 35)], lambda x: x.std(keepdim=True))
    helper_test_op([(15, 25, 35)], lambda x: x.std(0, keepdim=True, correction=0))
-  @slow_test
  def test_std_mean(self):
    helper_test_op([(15,25,35)], lambda x: torch.stack(torch.std_mean(x)),
                                 lambda x: Tensor.stack(*x.std_mean()))
@@ -2054,14 +2040,12 @@ class TestOps(unittest.TestCase):
      lambda x,w,b: torch.nn.functional.conv2d(x,w,b),
      lambda x,w,b: Tensor.conv2d(x,w,b), grad_rtol=1e-5)

-  @slow_test
  @unittest.skipIf(IMAGE>0, "no conv3d on images")
  def test_simple_conv3d(self):
    helper_test_op([(1,4,9,9,9), (4,4,3,3,3)],
      lambda x,w: torch.nn.functional.conv3d(x,w),
      lambda x,w: Tensor.conv2d(x,w), grad_rtol=1e-5)

-  @slow_test
  @unittest.skipIf(IMAGE>0, "no conv3d on images")
  def test_padded_conv3d(self):
    helper_test_op([(1,4,5,5,5), (4,4,3,3,3)],
@@ -2118,7 +2102,6 @@ class TestOps(unittest.TestCase):
      lambda x,w: torch.nn.functional.conv_transpose2d(x,w,groups=2),
      lambda x,w: Tensor.conv_transpose2d(x,w,groups=2), grad_rtol=1e-5)

-  @slow_test
  def test_padded_conv_transpose2d(self):
    for padding in [(1,2), (2,1), 2, 1, 0]:
      helper_test_op([(2,4,9,9), (4,4,3,3)],
@@ -2127,7 +2110,6 @@ class TestOps(unittest.TestCase):
    self.helper_test_exception([(2,16,2,2), (32,16,3,3)], lambda x,w: torch.nn.functional.conv_transpose2d(x,w,padding=(1,1,1)),
                   lambda x,w: Tensor.conv_transpose2d(x,w,padding=(1,1,1)), expected=(RuntimeError, ValueError))

-  @slow_test
  def test_dilated_conv_transpose2d(self):
    for dilation in [(1,2), (2,1), 2, 1]:
      helper_test_op([(2,4,9,9), (4,4,3,3)],
@@ -2140,7 +2122,6 @@ class TestOps(unittest.TestCase):
        lambda x,w: torch.nn.functional.conv_transpose2d(x,w, stride=stride),
        lambda x,w: Tensor.conv_transpose2d(x,w,stride=stride), atol=1e-5, grad_rtol=1e-5)

-  @slow_test
  def test_output_padded_conv_transpose2d(self):
    for output_padding, stride in [((1,1), (2,3)), ((2,1), (3,2))]:
      helper_test_op([(2,4,6,5), (4,4,3,3),(4,)],
@@ -2202,10 +2183,8 @@ class TestOps(unittest.TestCase):
              lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups),
              lambda x,w: Tensor.conv2d(x,w,groups=groups), grad_rtol=1e-5)
  def test_conv2d(self): self._test_conv2d(bs=1, cin=3)
-  @slow_test
  def test_conv2d_bs_4_cin_3(self): self._test_conv2d(bs=4, cin=3, cout=2)
  def test_conv2d_bs_1_cin_1(self): self._test_conv2d(bs=1, cin=1)
-  @slow_test
  def test_conv2d_bs_4_cin_1(self): self._test_conv2d(bs=4, cin=1)

  def test_conv2d_errors(self):
@@ -2277,7 +2256,6 @@ class TestOps(unittest.TestCase):
      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups),
      lambda x,w: Tensor.conv2d(x,w,groups=groups), grad_rtol=1e-5)

-  @slow_test
  def test_strided_conv2d_simple(self):
    bs,H,W = 2,3,1
    helper_test_op([(bs,1,5,1), (1,1,H,W)],
@@ -2288,7 +2266,6 @@ class TestOps(unittest.TestCase):
  def test_strided_conv2d_simple_vec(self):
    with Context(DEVECTORIZE=0): self.test_strided_conv2d_simple()

-  @slow_test
  def test_strided_conv2d(self):
    bs = 4
    cin = 3
@@ -2524,7 +2501,6 @@ class TestOps(unittest.TestCase):
      ),
      forward_only=True)

-  @slow_test
  def test_avg_pool2d(self):
    shape = (32,2,111,28)
    for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]:
@@ -3039,8 +3015,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,y.clip(0,1),
                                                                                                        pos_weight=torch.tensor(pos_weight)),
                                       lambda x,y: x.binary_crossentropy_logits(y.clip(0,1),pos_weight=Tensor(pos_weight)))
-
-  @unittest.skipIf(RANGEIFY > 1, "broken on RANGEIFY > 1, TODO: fix")
  def test_cross_entropy_class_probabilities(self):
    helper_test_op([(32,), (32,)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y))
    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y))
@@ -3164,7 +3138,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(32,10)], lambda x: x.masked_fill((x>0.1).detach(), -math.inf))
    helper_test_op([(32,10)], lambda x: x.masked_fill((x<0.1).detach(), -math.inf))

-  @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "AMD" and RANGEIFY, "very slow on MOCKGPU because reduce does not fold")
  def test_masked_select(self):
    helper_test_op([(32, 10)], lambda x: x.masked_select(x>0.5), lambda x: x.masked_select(x>0.5), forward_only=True)
    helper_test_op([(32, 10)], lambda x: x.masked_select(torch.tensor(True)), lambda x: x.masked_select(Tensor(True)), forward_only=True)
@@ -2,7 +2,7 @@ import unittest, pickle, types
 import numpy as np
 from tinygrad import Tensor, TinyJit, Variable, dtypes
 from tinygrad.helpers import GlobalCounters, ContextVar, Context
-from tinygrad.uop.ops import PatternMatcher, UPat, UOp
+from tinygrad.uop.ops import PatternMatcher, UPat, UOp, Ops

 class TestPickle(unittest.TestCase):
  def test_pickle_code_object(self):
@@ -45,9 +45,10 @@ class TestPickle(unittest.TestCase):
    t_values = t.numpy()
    del t # free buffers
    print("** post pickle")
+    init = GlobalCounters.kernel_count
    t2:Tensor = pickle.loads(st)
-    assert t2.uop.is_realized
    np.testing.assert_equal(t_values, t2.numpy())
+    self.assertEqual(GlobalCounters.kernel_count-init, 0)

  def test_pickle_realized_tensor_alt2(self):
    print("** init")
@@ -69,14 +70,14 @@ class TestPickle(unittest.TestCase):
  def test_pickle_buffer_uop(self):
    t = Tensor.arange(4).realize()
    a = t.uop
-    assert a.is_realized
-    self.assertIsNotNone(buffer:=a.base.realized)
+    assert a.op is Ops.BUFFER
+    self.assertIsNotNone(buffer:=a.realized)
    s = pickle.dumps(a)
    # free buffers
    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.realized.as_buffer().cast("I").tolist(), [0, 1, 2, 3])

  def test_pickle_unrealized_tensor(self):
    t = Tensor.ones(10, 10)
@@ -17,7 +17,7 @@ def helper_collect_profile(*devs):
  cpu_events.clear()

  profile_list = []
-  with Context(VIZ=1):
+  with Context(PROFILE=1):
    yield profile_list
    for dev in devs: dev.synchronize()
    for dev in devs: dev._at_profile_finalize()
@@ -3,7 +3,6 @@ import numpy as np
 import unittest
 from dataclasses import replace
 from tinygrad import Tensor, Context, Device, dtypes
-from tinygrad.helpers import RANGEIFY
 from tinygrad.uop.ops import Ops
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item, get_program
@@ -94,8 +93,7 @@ class TestQuantizeOnnx(unittest.TestCase):
    X = Tensor(np.random.uniform(0, 255, size=(1, 32, 128, 128)).astype(np.uint8))
    W = Tensor(np.random.uniform(0, 255, size=(64, 32, 1, 1)).astype(np.uint8))
    out = X.conv2d(W, dtype=X.dtype)
-    # rangeify merges axis in a different order
-    opts = [Opt(op=OptOps.UPCAST, axis=0 if RANGEIFY else 1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
    sexec(out, opts)

  def test_prequant_gemm(self):
@@ -1,16 +1,15 @@
 import unittest, math
 from functools import partial

-from tinygrad import nn, dtypes, Tensor, Device, TinyJit, Variable
-from tinygrad.helpers import getenv, CI, OSX
-from tinygrad.device import is_dtype_supported
-from tinygrad.engine.realize import lower_schedule, CompiledRunner
-from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import not_support_multi_device
-
 import numpy as np
 import torch
+from tinygrad import nn, dtypes, Tensor, Device, TinyJit
+from tinygrad.helpers import getenv, CI
+from tinygrad.device import is_dtype_supported
+from tinygrad.engine.realize import lower_schedule, CompiledRunner
 from hypothesis import given, settings, strategies as strat
+from test.helpers import not_support_multi_device
+from tinygrad.renderer.ptx import PTXRenderer

 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -361,20 +360,5 @@ class TestRandomness(unittest.TestCase):
    assert equal_distribution(lambda *_: nn.BatchNorm2d(*params).weight, lambda _: torch.nn.BatchNorm2d(*params).weight.detach())
    assert equal_distribution(lambda *_: nn.BatchNorm2d(*params).bias, lambda _: torch.nn.BatchNorm2d(*params).bias.detach())

-# TODO: still fails with MAX_KERNEL_BUFFERS
-@unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
-class TestSample(unittest.TestCase):
-  def test_sample(self):
-    X = Tensor.rand(10000, 50).realize()
-    BS = 16
-    idxs = np.random.randint(0, X.shape[0], size=(BS))
-    # this uncovered a bug with arg sort order
-    batch = [Variable(f'idx{i}', 0, X.shape[0]-1).bind(s) for i,s in enumerate(idxs.tolist())]
-    x = Tensor.cat(*[X.shrink(((batch[i], batch[i]+1), None)) for i in range(BS)])
-    print(idxs)
-    ret = x.numpy()
-    base = X.numpy()[idxs]
-    np.testing.assert_equal(ret, base)
-
 if __name__ == "__main__":
  unittest.main()
@@ -3,48 +3,10 @@ from tinygrad import Tensor, nn
 from tinygrad.helpers import RANGEIFY, Context, GlobalCounters
 from tinygrad.uop.ops import UOp

-@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
-class TestRangeifyAssign(unittest.TestCase):
-  def test_assign_permuted(self):
-    A = Tensor.empty(4, 4, dtype='int')
-    B = Tensor.arange(16).reshape(4,4)
-    ret = A.permute(1,0).assign(B)
-    lst = ret.tolist()
-    lst2 = A.tolist()
-    lst3 = B.tolist()
-    print(lst)
-    print(lst2)
-    print(lst3)
-    self.assertListEqual(lst, lst3)
-    self.assertListEqual(lst2, B.permute(1, 0).tolist())
-
 N = 256

-class TestRangeifyOpt(unittest.TestCase):
-  def test_randperm(self):
-    Tensor.randperm(10000).realize()
-
-  def test_one_getitem(self):
-    X = Tensor.empty(10000)
-    sel = Tensor.arange(1000).contiguous().realize()
-    Xsel = X[sel]
-    Tensor.realize(Xsel)
-
-  def test_two_getitem(self):
-    # this is splitting on the child even when it really shouldn't
-    X = Tensor.empty(10000)
-    Y = Tensor.empty(10000)
-    sel = Tensor.arange(1000).contiguous().realize()
-    Xsel, Ysel = X[sel], Y[sel]
-    Tensor.realize(Xsel, Ysel)
-
@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
 class TestRangeify(unittest.TestCase):
-  def test_groupnorm(self):
-    # ranges 1 and 3 are merging
-    x = nn.GroupNorm(32, 128)
-    x(Tensor.empty(1, 128, 64, 64)).realize()
-
  def test_expand_children(self):
    A = Tensor.empty(N, N).sum(axis=1)
    ba = A.expand(N, N)
@@ -82,14 +44,6 @@ class TestRangeify(unittest.TestCase):
    C = Tensor.empty(N, N)
    (((A@B).exp()@C).exp()).realize()

-  def test_double_gemm_exp_child(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    # A@B is used with exp, and also on the sum. this is two kernels now, is this right?
-    ret = A@B
-    ((ret.exp()@C)+ret).realize()
-
  def test_double_gemm_relu(self):
    A = Tensor.empty(N, N)
    B = Tensor.empty(N, N)
@@ -128,11 +82,6 @@ class TestRangeify(unittest.TestCase):
    w1 = Tensor.empty(8, 4, 3, 3)
    x.conv2d(w1).realize()

-  def test_conv2d_elu(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    x.conv2d(w1).elu().realize()
-
  def test_conv2d_t(self):
    x = Tensor.empty(1, 4, 32, 32)
    w1 = Tensor.empty(8, 4, 3, 3)
@@ -144,13 +93,6 @@ class TestRangeify(unittest.TestCase):
    w2 = Tensor.empty(12, 8, 3, 3)
    x.conv2d(w1).conv2d(w2).realize()

-  def test_xception_conv2d(self):
-    # NOTE: this fusion is bad, it's recomputing the inner many times
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 1, 1)
-    w2 = Tensor.empty(8, 1, 3, 3)
-    x.conv2d(w1).conv2d(w2, groups=8).realize()
-
  def test_conv_maxpool_contig(self): self.test_conv_maxpool(True)
  def test_conv_maxpool(self, contig=False):
    GlobalCounters.reset()
@@ -0,0 +1,22 @@
+import unittest
+import numpy as np
+from tinygrad import Tensor, Variable, Device
+from tinygrad.helpers import OSX
+
+# TODO: still fails with MAX_KERNEL_BUFFERS
+@unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
+class TestSample(unittest.TestCase):
+  def test_sample(self):
+    X = Tensor.rand(10000, 50).realize()
+    BS = 16
+    idxs = np.random.randint(0, X.shape[0], size=(BS))
+    # this uncovered a bug with arg sort order
+    batch = [Variable(f'idx{i}', 0, X.shape[0]-1).bind(s) for i,s in enumerate(idxs.tolist())]
+    x = Tensor.cat(*[X.shrink(((batch[i], batch[i]+1), None)) for i in range(BS)])
+    print(idxs)
+    ret = x.numpy()
+    base = X.numpy()[idxs]
+    np.testing.assert_equal(ret, base)
+
+if __name__ == '__main__':
+  unittest.main()
@@ -8,13 +8,13 @@ import functools
 from typing import cast
 from hypothesis import assume, given, settings, strategies as strat

-from tinygrad import nn, dtypes, Device, Tensor, Variable
+from tinygrad import nn, dtypes, Device, Tensor
 from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType, ImageDType
 from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites
 from tinygrad.uop.symbolic import symbolic_simple
-from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp, RANGEIFY
+from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
 from tinygrad.schedule.kernelize import merge_views, get_kernelize_map, Kernel
 from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule
@@ -33,7 +33,6 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
  # test lowering all the ScheduleItems to ExecItems
  kernel_cnt = len([si for si,ei in lower_schedule(sched.copy()) if isinstance(ei.prg, CompiledRunner) or not filter_sink])
  if kernel_cnt != allowed:
-    if RANGEIFY: return sched # allow different kernel count, TODO: fix the asserts
    print(f"SCHEDULE ISSUE, expecting {allowed} got {len(sched)}")
    if DEBUG >= 3:
      for i,s in enumerate(sched):
@@ -42,9 +41,6 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
    raise KernelCountException(f"{kernel_cnt} != {allowed}")
  return sched

-def expect_rangeify_fails(fxn): return (unittest.expectedFailure if RANGEIFY else (lambda f:f))(fxn)
-def expect_nonrangeify_fails(fxn): return (unittest.expectedFailure if not RANGEIFY else (lambda f:f))(fxn)
-
 def _realize_weights(m):
  for p in nn.state.get_parameters(m): p.realize()

@@ -115,7 +111,6 @@ class TestSchedule(unittest.TestCase):
    self.assertListEqual(a.tolist(), [[15]])

  @unittest.skipIf(Device.DEFAULT == "CPU", "devices must mismatch")
-  @expect_rangeify_fails
  def test_error_on_device_mismatch(self):
    a = Tensor.empty(10)
    b = Tensor.empty(10, device="CPU")
@@ -123,12 +118,11 @@ class TestSchedule(unittest.TestCase):
    with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 1)

  @unittest.skipIf(Device.DEFAULT == "CPU", "devices must mismatch")
-  @expect_rangeify_fails
  def test_error_on_device_mismatch_alt(self):
    a = Tensor.empty(10)
    b = Tensor.empty((1,), device="CPU").expand(10).contiguous()
    c = a+b
-    with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 2 if RANGEIFY else 1)
+    with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 1)

  @unittest.skipUnless(is_dtype_supported(dtypes.half) and getenv("CAST_AFTER_EXPAND"), "need half and CAST_AFTER_EXPAND=1")
  @unittest.skip("CAST_AFTER_EXPAND is not supported")
@@ -146,7 +140,6 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_equal(xt.numpy(), X.numpy()[1][0])

  @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
-  @unittest.skipIf(RANGEIFY, "rangeify doesn't implement input buffer limiting")
  def test_add_chain_buffers(self):
    N = 31
    with Context(TRACK_MATCH_STATS=0, DEBUG=0):
@@ -205,10 +198,9 @@ class TestSchedule(unittest.TestCase):

  def test_simplify_padded_const(self):
    a = Tensor.empty(1022).cummax(axis=0)
-    check_schedule(a, 5)
-    # TODO: what is this testing?
-    #ast = sched[0].ast
-    #self.assertLessEqual(len([u for u in ast.toposort() if u.op is Ops.WHERE]), 6)
+    sched = check_schedule(a, 5)
+    ast = sched[0].ast
+    self.assertLessEqual(len([u for u in ast.toposort() if u.op is Ops.WHERE]), 6)

  def test_basic_binop_fusion(self):
    a = Tensor.empty(10)
@@ -286,7 +278,7 @@ class TestSchedule(unittest.TestCase):
    a = Tensor.empty(10,10,10)
    b = Tensor.empty(10,10,1)
    c = a.sum(axis=0, keepdim=True).permute(2,1,0) + b
-    check_schedule(c, 2)
+    with self.assertRaises(KernelCountException): check_schedule(c, 1)

  def test_allow_push_permutes(self):
    a = Tensor.randn(10,10,10).realize()
@@ -324,7 +316,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.empty(10)
    c = a+b
    d = a.reshape(10,1)+b.reshape(10,1)
-    check_schedule(d, 1, [c])
+    with self.assertRaises(KernelCountException): check_schedule(d, 0, [c])

  # failing in new lazy
  def test_cache_binaryop_transpose(self):
@@ -332,7 +324,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.empty(10,10)
    c = (a.T*b.T).T #.contiguous()
    d = a*b
-    check_schedule(d, 1, [c])
+    with self.assertRaises(KernelCountException): check_schedule(d, 0, [c])

  def test_cache_two_reduceops(self):
    a = Tensor.empty(10)
@@ -347,7 +339,7 @@ class TestSchedule(unittest.TestCase):
    r1 = (x - r0).sum(axis=0).div(2)
    out = r0 + r1
    schedule = check_schedule(out, 2)
-    reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
+    reduceops = [x for si in schedule for x in si.ast.toposort() if x.op is Ops.REDUCE_AXIS]
    assert len(reduceops) == 2

  def test_cache_reduce_multiple_children(self):
@@ -357,9 +349,9 @@ class TestSchedule(unittest.TestCase):
    r1 = (x - r0).sum(axis=0).div(2)
    out0 = r0 + y
    out1 = r1 + y
-    schedule = check_schedule([out0, out1], 2 if RANGEIFY else 4)
-    reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
-    assert len(reduceops) == (3 if RANGEIFY else 2)
+    schedule = check_schedule([out0, out1], 4)
+    reduceops = [x for si in schedule for x in si.ast.toposort() if x.op is Ops.REDUCE_AXIS]
+    assert len(reduceops) == 2

  def test_div_collapse_buffer(self):
    a = Tensor.full((4,), 4.0).contiguous().realize()
@@ -402,7 +394,6 @@ class TestSchedule(unittest.TestCase):
    # a and b share the same underlying device memory
    self.assertIs(a.uop.realized, b.uop.realized)

-  @expect_rangeify_fails
  def test_clone_doesnt_dedup(self):
    src = Tensor.ones(4).contiguous().realize()
    a = src.clone()
@@ -426,11 +417,6 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.full((4, 4), 1.).contiguous().realize()
    check_schedule([a+b, a+b], 1)

-  def test_const_realize(self):
-    t = Tensor.ones(2)
-    check_schedule(t[0], 0)
-    check_schedule(t[1], 0)
-
  def test_fold_double_unary(self):
    y = Tensor.empty(2)
    out = y.sum(keepdim=True).sqrt().neg()
@@ -572,7 +558,7 @@ class TestSchedule(unittest.TestCase):
    c = a+b
    d = a.reshape(10,1)+b.reshape(10,1)
    out = c.sum() + d.sum()
-    check_schedule(out, 2)
+    with self.assertRaises(KernelCountException): check_schedule(out, 1)

  def test_children_dont_push(self):
    a = Tensor.empty(10, 10, 1)
@@ -583,7 +569,6 @@ class TestSchedule(unittest.TestCase):
    check_schedule(f, 2)

  # failing in new lazy
-  @unittest.skip("always fusing elementwise")
  def test_dont_fuse_binops_with_children(self):
    a = Tensor.empty(10)
    b = Tensor.empty(10)
@@ -591,8 +576,8 @@ class TestSchedule(unittest.TestCase):
    keep_me = a+b
    e = keep_me.sum() # noqa: F841 give keep_me a child (NOTE: BinaryOps won't be a child since it will instant fuse)
    d = keep_me+c
-    check_schedule(d, 2)
-    check_schedule(keep_me, 0, [d])
+    with self.assertRaises(KernelCountException): check_schedule(d, 2)
+    with self.assertRaises(KernelCountException): check_schedule(keep_me, 0, [d])

  #@unittest.skip("failing in old lazy")
  def test_permute_breaks_fusion(self):
@@ -642,8 +627,7 @@ class TestSchedule(unittest.TestCase):
    x = x.image_conv2d(w3, b3)

    # NOOP, 3 convs, contiguous
-    #check_schedule(x, 5)
-    check_schedule(x, 8)
+    with self.assertRaises(KernelCountException): check_schedule(x, 5)

  def test_image_conv_fusion_minimal(self):
    b1 = Tensor.empty(16)
@@ -716,12 +700,9 @@ class TestSchedule(unittest.TestCase):
    prev_a = (a+1).contiguous()
    a.assign(Tensor([2]))
    a.kernelize(prev_a)
-    # RANGEIFY doesn't apply the post diamond graph, it's fine since we can always apply the fixup on each kernelize call
-    if not RANGEIFY:
-      assert prev_a.uop in a.uop.src, "contiguous usage must run before assign"
+    assert prev_a.uop in a.uop.src, "contiguous usage must run before assign"
    self.assertEqual((prev_a+a*3).item(), 1+2*3)

-  @expect_rangeify_fails
  def test_multioutput_ast(self):
    a = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize().uop
    b = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize().uop
@@ -803,13 +784,6 @@ class TestSchedule(unittest.TestCase):
    out = x + 1
    check_schedule(out, 0, filter_sink=False)

-  def test_zero_size_assign(self):
-    f = Tensor.full((2,), 0.).contiguous().realize()
-    a = f.shrink_to((0,))
-    a.assign(Tensor.ones_like(a))
-    check_schedule(a, 0)
-    self.assertEqual(a.tolist(), [])
-
  def test_reduce_permute_nofuse(self):
    x = Tensor.empty(32, 32, 32)
    y = Tensor.empty(32, 32)
@@ -914,24 +888,26 @@ class TestSchedule(unittest.TestCase):
    out = x.contiguous() + y.contiguous()
    check_schedule(out, 2, filter_sink=False)

+  @unittest.expectedFailure
  def test_reduce_same_size(self):
    Tensor.manual_seed(0)
    a = Tensor.randn(4, 4).realize()
    out0 = a.sum() + 2
    out1 = a.sum() + 4
    out2 = out0 * out1
-    run_schedule(check_schedule([out0, out1, out2], 1 if RANGEIFY else 4))
+    run_schedule(check_schedule([out0, out1, out2], 1))
    np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-6)
    np.testing.assert_allclose(out1.numpy(), out1_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6)
    np.testing.assert_allclose(out2.numpy(), out0_np*out1_np, atol=1e-4, rtol=1e-6)

+  @unittest.expectedFailure
  def test_reduce_multiple_paths(self):
    Tensor.manual_seed(0)
    a = Tensor.randn(4, 4).realize()
    out0 = a.sum().exp2()
    # out1 has two paths to a.sum()
    out1 = a.sum() + out0
-    run_schedule(check_schedule([out0, out1], 1 if RANGEIFY else 3))
+    run_schedule(check_schedule([out0, out1], 1))
    np.testing.assert_allclose(out0.numpy(), out0_np:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+out0_np, atol=1e-4, rtol=1e-6)

@@ -1007,6 +983,7 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_allclose(e.numpy(), e_np:=b.numpy() + out0_np, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out1.numpy(), r_np + e_np[0][0][0], atol=1e-4, rtol=1e-4)

+  # changed by multireduce
  def test_reduce_expand_child(self):
    Tensor.manual_seed(0)
    a = Tensor.randn((32, 32, 32)).realize()
@@ -1018,12 +995,13 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_allclose(out0.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+b.numpy(), atol=1e-4, rtol=1e-4)

+  @unittest.expectedFailure
  def test_reduce_shrink_child(self):
    a = Tensor.empty(100, 100)
    b = Tensor.empty(10,)
    c = a.sum() + b[0]
    d = a.sum() + 2
-    check_schedule([c, d], 1 if RANGEIFY else 3)
+    check_schedule([c, d], 1)

  def test_reduce_multiple_paths_midshrink(self):
    a = Tensor.empty(4, 4)
@@ -1046,6 +1024,20 @@ class TestSchedule(unittest.TestCase):
    run_schedule(check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), x.numpy().std(axis=-1, ddof=1), atol=1e-4, rtol=1e-4)

+  def test_argmin_multireduce_fusion(self):
+    Tensor.manual_seed(0)
+    x = Tensor.randn(4, 32).realize()
+    out = x.argmin(-1)
+    run_schedule(check_schedule(out, 2))
+    np.testing.assert_equal(out.numpy(), x.numpy().argmin(axis=-1))
+
+  def test_argmax_multireduce_fusion(self):
+    Tensor.manual_seed(0)
+    x = Tensor.randn(4, 32).realize()
+    out = x.argmax(-1)
+    run_schedule(check_schedule(out, 2))
+    np.testing.assert_equal(out.numpy(), x.numpy().argmax(axis=-1))
+
  def test_scaled_dot_product_attention_multireduce_fusion(self):
    Tensor.manual_seed(0)
    q = Tensor.randn(32,8,16,8).realize()
@@ -1187,14 +1179,13 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)

  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  @expect_rangeify_fails
  def test_softmax_upcast(self):
    # input half, softmax in float
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.half).realize()
    out = x.softmax(dtype=dtypes.float)
    sched = out.schedule()
-    self.assertEqual(len(sched), 2 if RANGEIFY else 3)
+    self.assertEqual(len(sched), 3)
    self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)

    # input float, softmax in float
@@ -1211,6 +1202,7 @@ class TestSchedule(unittest.TestCase):
    x.softmax().sum().backward()
    run_schedule(check_schedule(x.grad, 4))

+  # changed by: multireduce spec
  def test_layernorm_onelayer_fusion(self):
    Tensor.manual_seed(0)
    layer = nn.LayerNorm([10, 10])
@@ -1324,7 +1316,6 @@ class TestSchedule(unittest.TestCase):
      with Context(FUSE_CONV_BW=1): check_schedule(opt.schedule_step(), 14)

  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  @expect_rangeify_fails
  def test_prefer_half_buffer(self):
    x = Tensor.ones(4).contiguous().realize()
    # y = Tensor.ones(4).contiguous().realize()
@@ -1442,6 +1433,7 @@ class TestSchedule(unittest.TestCase):
    run_schedule(schedule)
    np.testing.assert_allclose(b.numpy(), a.numpy().sum(0)+a.numpy().max(0) + a.numpy().max(1)+a.numpy().sum(1)+2, atol=1e-4, rtol=1e-4)

+  # changed by: multireduce spec
  # pattern in test_transformer
  def test_partial_fuse1(self):
    Tensor.manual_seed(0)
@@ -1454,6 +1446,7 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_allclose(c.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(d.numpy(), (a.numpy().sum() - b.numpy().sum()) * 4, atol=1e-4, rtol=1e-4)

+  # changed by: multireduce spec
  # pattern in conv
  def test_partial_fuse2(self):
    Tensor.manual_seed(0)
@@ -1466,7 +1459,9 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_allclose(c.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(d.numpy(), b.numpy().sum()-(a.numpy().sum()+2), atol=1e-4, rtol=1e-4)

+  # changed by: multireduce spec
  # pattern in adam
+  @unittest.expectedFailure
  def test_partial_fuse3(self):
    Tensor.manual_seed(0)
    a = Tensor.randn(16, 16).realize()
@@ -1476,12 +1471,14 @@ class TestSchedule(unittest.TestCase):
    e = c * d
    f = b.sum() - e
    # run_schedule(check_schedule([c, d, e, f], 1))
-    run_schedule(check_schedule([c, d, e, f], 2 if RANGEIFY else 5))
+    run_schedule(check_schedule([c, d, e, f], 2))
    np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(f.numpy(), b.numpy().sum() - e_np, atol=1e-4, rtol=1e-4)

+  # changed by: multireduce spec
+  @unittest.expectedFailure
  def test_partial_fuse4(self):
    Tensor.manual_seed(0)
    a = Tensor.randn(16, 16).realize()
@@ -1491,7 +1488,7 @@ class TestSchedule(unittest.TestCase):
    e = c * d
    f = (b - d).sum() - e
    # run_schedule(check_schedule([c, d, e, f], 1))
-    run_schedule(check_schedule([c, d, e, f], 5))
+    run_schedule(check_schedule([c, d, e, f], 3))
    np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
@@ -1626,11 +1623,11 @@ class TestSchedule(unittest.TestCase):
    out = x.argmax(1)
    run_schedule(check_schedule(out, 2))

-  def test_conv2d(self): _test_conv2d(4 if RANGEIFY else 7)
-  def test_conv2d_fused(self): _test_conv2d(4 if RANGEIFY else 5, FUSE_CONV_BW=1)
+  def test_conv2d(self): _test_conv2d(7)
+  def test_conv2d_fused(self): _test_conv2d(5, FUSE_CONV_BW=1)

  @unittest.skipUnless(is_dtype_supported(dtypes.half) and is_dtype_supported(dtypes.ulong), "need half and ulong")
-  def test_conv2d_half(self): _test_conv2d(4 if RANGEIFY else 7, dtype=dtypes.half)
+  def test_conv2d_half(self): _test_conv2d(7, dtype=dtypes.half)
  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Causes other tests to fail")
  @unittest.expectedFailure
@@ -1693,7 +1690,6 @@ class TestSchedule(unittest.TestCase):
  def test_late_fusion_post_expand(self):
    self._test_fusion([(32, 32)], lambda a:a-a.sum(1), 2)

-  @expect_rangeify_fails
  def test_cast_padded_view(self):
    a = Tensor.arange(4).reshape(1, 4)
    casted_view = a.pad(((0, 1), (0, 0))).cast(dtypes.float)
@@ -1723,7 +1719,6 @@ class TestSchedule(unittest.TestCase):
    self.assertListEqual(realized_const_view.tolist(), [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]])

  @given(strat.sampled_from(dtypes.all), strat.sampled_from(dtypes.all))
-  @expect_rangeify_fails
  def test_cast_padded_const(self, dt1, dt2):
    assume(is_dtype_supported(dt1) and is_dtype_supported(dt2))
    a = Tensor(1, dtype=dt1).reshape(1, 1).pad(((1, 1), None))
@@ -1733,41 +1728,53 @@ class TestSchedule(unittest.TestCase):
    run_schedule(check_schedule(realized_const_view, 1))
    np.testing.assert_equal(realized_const_view.numpy(), [[0], [1], [0]])

+class TestIndexing(unittest.TestCase):
+  def check_schedule(self, xt:Tensor|list[Tensor], cnt:int):
+    with Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1)):
+      lst = [xt] if isinstance(xt, Tensor) else xt
+      s = Tensor.schedule(*lst)
+      lowered = [x[1] for x in lower_schedule(s.copy())]
+      kernels = [ei for ei in list(lowered) if isinstance(ei.prg, CompiledRunner)]
+      if FUSE_ARANGE and len(kernels) != cnt:
+        raise KernelCountException(f"{len(kernels)} != {cnt}")
+      for ei in lowered: ei.run(do_update_stats=True)
+    return s
+
  def test_simple_indexing(self):
    X = Tensor.randn(10, 10).realize()
    idxs = Tensor([0, 2]).realize()
    xt = X[idxs]
-    run_schedule(check_schedule(xt, 2))
+    self.check_schedule(xt, 2)
    np.testing.assert_equal(xt.numpy(), X.numpy()[idxs.numpy()])

  def test_simple_indexing_alt(self):
    X = Tensor.arange(16).reshape(4, 4)
    xt = X[[1, 2], [-1, 2]]
-    run_schedule(check_schedule(xt, 1))
+    self.check_schedule(xt, 1)
    np.testing.assert_equal(xt.numpy(), (np.arange(16).reshape(4, 4))[[1, 2], [-1, 2]])

  def test_advanced_indexing(self):
    X = Tensor.arange(10)+1
    xt = X[[0, -1]]
-    run_schedule(check_schedule(xt, 1))
+    self.check_schedule(xt, 1)
    np.testing.assert_equal(xt.numpy(), (np.arange(10)+1)[[0, -1]])

  def test_advanced_indexing_alt(self):
    X = Tensor.arange(6).reshape(3, 2)+1
    xt = X[[Tensor([2]), Tensor([1])]]
-    run_schedule(check_schedule(xt, 3))
+    self.check_schedule(xt, 3)
    np.testing.assert_equal(xt.numpy(), 6)

  def test_advanced_simple_indexing_combined(self):
    X = Tensor.arange(16).reshape(4, 4)
    xt = X[1:2, [-1, 2]]
-    run_schedule(check_schedule(xt, 1))
+    self.check_schedule(xt, 1)

  def test_push_through_reshape(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(10, 20).realize()
    out = x.argmax(1)
-    run_schedule(check_schedule(out, 2))
+    self.check_schedule(out, 2)
    np.testing.assert_allclose(out.numpy(), np.argmax(x.numpy(), 1))

  def test_arange_push_through_expand(self):
@@ -1775,35 +1782,35 @@ class TestSchedule(unittest.TestCase):
    a = Tensor.arange(4,)
    b = Tensor.randn(4, 4).realize()
    out = (a+b).sum()
-    run_schedule(check_schedule(out, 1))
+    self.check_schedule(out, 1)
    np.testing.assert_allclose(out.numpy(), (np.arange(4)+b.numpy()).sum(), atol=1e-5)

  def test_argmin(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 32).realize()
    out = x.argmin(-1)
-    run_schedule(check_schedule(out, 2))
+    self.check_schedule(out, 2)
    np.testing.assert_equal(out.numpy(), x.numpy().argmin(axis=-1))

  def test_argmax(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 32).realize()
    out = x.argmax(-1)
-    run_schedule(check_schedule(out, 2))
+    self.check_schedule(out, 2)
    np.testing.assert_equal(out.numpy(), x.numpy().argmax(axis=-1))

  def test_arange_transposed(self):
    Tensor.manual_seed(0)
    x = Tensor.randint(4, 1).realize()
    a = ((Tensor.arange(4,)*x).T).sum()
-    run_schedule(check_schedule(a, 1))
+    self.check_schedule(a, 1)
    np.testing.assert_equal(a.numpy(), (np.arange(4)*x.numpy()).T.sum())

  def test_div_padded_arange(self):
    x = Tensor.full((2,2), 16)
    y = x.idiv(Tensor.linspace(2, 8, steps=4, dtype=dtypes.int).reshape(2,2)).pad(((1,1), (1,1)))
    out = y.sum(axis=1)
-    run_schedule(check_schedule(out, 2))
+    with Context(FUSE_ARANGE=1): run_schedule(check_schedule(out, 2))
    self.assertListEqual(out.tolist(), [0, 12, 4, 0])

  def test_arange_transposed_descendants(self):
@@ -1812,7 +1819,7 @@ class TestSchedule(unittest.TestCase):
    a = (Tensor.arange(4,)*x).T
    b = Tensor.randint(4, 4).realize()
    out = (a+b).sum()
-    run_schedule(check_schedule(out, 1))
+    self.check_schedule(out, 1)
    np.testing.assert_equal(out.numpy(), ((np.arange(4)*x.numpy()).T+b.numpy()).sum())

  def test_arange_index(self):
@@ -1820,7 +1827,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(5, 2).realize()
    a = Tensor.arange(10)
    out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 1))
+    self.check_schedule(out, 1)
    np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)

  def test_arange_index_shrink(self):
@@ -1829,14 +1836,14 @@ class TestSchedule(unittest.TestCase):
      x = Tensor.randn(11).realize()
    a = Tensor.arange(22)
    out = (x + a[:11]).sum()
-    check_schedule(out, 1)
+    self.check_schedule(out, 1)

  def test_arange_index_contiguous(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(5, 2).realize()
    a = Tensor.arange(10).contiguous()
    out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 3))
+    self.check_schedule(out, 3)
    np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)

  def test_arange_index_child(self):
@@ -1844,24 +1851,62 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(5, 2).realize()
    a = Tensor.arange(10)+1
    out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 1))
+    self.check_schedule(out, 1)
    np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)

-  def test_user_contiguous(self):
+  def test_arange_index_contiguous_child(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(5, 2).realize()
    a = (Tensor.arange(10)+1).contiguous()
    out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 3))
+    self.check_schedule(out, 3)
    np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)

+  def test_arange_childless_base(self):
+    a = Tensor.arange(4)
+    self.check_schedule(a, 1)
+    np.testing.assert_equal(a.numpy(), np.arange(4))
+
+  def test_arange_childless_view(self):
+    a = Tensor.arange(4).reshape(2, 2)
+    a[0] = 4
+    np.testing.assert_equal(a.numpy(), [[4, 4], [2, 3]])
+
+  def test_arange_group_childless_base(self):
+    Tensor.manual_seed(0)
+    x = Tensor.randint(4).realize()
+    a = Tensor.arange(4)+x
+    self.check_schedule(a, 1)
+    np.testing.assert_equal(a.numpy(), np.arange(4)+x.numpy())
+
+  def test_arange_group_childless_view(self):
+    Tensor.manual_seed(0)
+    x = Tensor.ones(4).contiguous().realize()
+    a = Tensor.arange(4)+x
+    a[0] = 6
+    np.testing.assert_equal(a.numpy(), [6., 2., 3., 4.])
+
  @unittest.skip("BUFFER_VIEW no longer supported on non-disk devices")
  def test_arange_view_op(self):
    a = Tensor.arange(12).reshape(4, 3).shrink(((1, 2), (1, 3))).contiguous()
-    sched = run_schedule(check_schedule(a, 1))
+    sched = self.check_schedule(a, 1)
    self.assertIs(sched[1].ast.op, Ops.BUFFER_VIEW)
    np.testing.assert_equal(a.numpy(), [[4, 5]])

+  @unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from ext device")
+  def test_arange_shrink_copy(self):
+    a = Tensor.arange(12).reshape(4, 3).shrink(((1, 2), (1, 3))).to("CPU")
+    sched = self.check_schedule(a, 2) # NOTE: there is a contiguous between REDUCE_AXIS and COPY
+    self.assertIs(sched[-1].ast.op, Ops.COPY)
+    np.testing.assert_equal(a.numpy(), [[4, 5]])
+
+  @unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from ext device")
+  def test_arange_expand_copy(self):
+    a = Tensor.arange(4).reshape(2, 2, 1).expand(2, 2, 2).contiguous().to("CPU")
+    sched = self.check_schedule(a, 2) # NOTE: there is a contiguous between REDUCE_AXIS and COPY
+    self.assertIs(sched[2].ast.op, Ops.COPY)
+    np.testing.assert_equal(a.numpy(), [[[0, 0], [1, 1]], [[2, 2], [3, 3]]])
+
  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  def test_precompute_freqs_cis(self):
    from extra.models.llama import precompute_freqs_cis
@@ -1877,7 +1922,7 @@ class TestSchedule(unittest.TestCase):
  def test_fuse_assign_contiguous(self):
    x = Tensor.zeros(4, 4, dtype=dtypes.int).contiguous().realize()
    a = Tensor.arange(8).reshape(4, 2)
-    run_schedule(check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2))
+    self.check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2)
    np.testing.assert_equal(x.numpy(), [[0, 1, 0, 0], [2, 3, 0, 0], [4, 5, 0, 0], [6, 7, 0, 0]])

  def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True)
@@ -1895,27 +1940,13 @@ class TestSchedule(unittest.TestCase):
      tst = x.shrink((None, (0, 2))).assign(a).realize()
      xref[:, :2] = np.arange(8).reshape(4, 2)+y.numpy()
    np.testing.assert_equal(x.numpy(), xref)
-    if RANGEIFY > 0:
-      # NOTE: this is a bug on non rangeify
-      np.testing.assert_equal(tst.numpy(), a.numpy())
-
-  def test_setitem_sched(self, transpose=False):
-    a = Tensor.arange(16, device="CPU").reshape(4, 4).contiguous().realize()
-    a2 = a.T if transpose else a
-    expected = (a+a2).tolist()
-    a.assign(a+a2)
-    kcount = len(sched:=a.schedule())
-    run_schedule(sched)
-    self.assertListEqual(a.tolist(), expected)
-    self.assertEqual(kcount, 2 if transpose else 1)
-  @unittest.skipUnless(RANGEIFY>0, "this asserts on non rangeify")
-  def test_setitem_permuted_sched(self): self.test_setitem_sched(transpose=True)
+    np.testing.assert_equal(tst.numpy(), a.numpy())

  def test_sparse_categorical_crossentropy_simple(self):
    X = Tensor([[0, 2, 3], [1, 2, 3]]).realize()
    Y = Tensor([1, 2]).realize()
    loss = X.sparse_categorical_crossentropy(Y)
-    run_schedule(check_schedule(loss, 4))
+    self.check_schedule(loss, 4)
    np.testing.assert_allclose(loss.item(), 0.878309, atol=1e-5, rtol=1e-6)

  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Validation error on WebGPU")
@@ -1927,21 +1958,28 @@ class TestSchedule(unittest.TestCase):
    yt = Tensor.randn(BS, 10).realize()
    with Context(SPLIT_REDUCEOP=0):
      loss = yt.sparse_categorical_crossentropy(Y_train[samples])
-      run_schedule(check_schedule(loss, 6))
+      self.check_schedule(loss, 6)
      loss_fused = loss.numpy()
    loss_ref = torch.nn.CrossEntropyLoss()(torch.tensor(yt.numpy()), torch.tensor(Y_train.numpy())[torch.tensor(samples.numpy())])
    np.testing.assert_allclose(loss_fused, loss_ref.numpy(), atol=1e-6, rtol=1e-6)

+  @unittest.expectedFailure
  def test_arange_fuse_grouped_children(self):
    X = Tensor.randn(4, 4).realize()
    r = (X+Tensor.arange(16).reshape(4, 4)).sum()
    out0 = r+2
    out1 = r+3
-    run_schedule(check_schedule([out0, out1], 1 if RANGEIFY else 3))
+    self.check_schedule([out0, out1], 1)
    r_ref = (X.numpy()+np.arange(16).reshape(4, 4)).sum()
    np.testing.assert_allclose(out0.numpy(), r_ref+2, rtol=2e-7)
    np.testing.assert_allclose(out1.numpy(), r_ref+3, rtol=2e-7)

+  def test_dont_fold_arange_contiguous_view(self):
+    X = Tensor.randn(4, 4).realize()
+    r = (X+Tensor.arange(16).reshape(4, 4).contiguous()).sum(1, keepdim=True)
+    self.check_schedule([r], 2)
+    np.testing.assert_allclose(r.numpy(), (X.numpy()+np.arange(16).reshape(4, 4)).sum(1, keepdims=True), atol=1e-5, rtol=1e-6)
+
  @unittest.skip("multi output isn't supported")
  def test_multiview_arange_children(self):
    X = Tensor.randn(2,3,4,4).numpy()
@@ -1958,19 +1996,6 @@ class TestSchedule(unittest.TestCase):
    self.assertEqual(new_uop.st, ShapeTracker.from_shape((4,)).reshape((4, 1)))
    self.assertEqual(swizzle_cnt(new_uop), 0)

-  @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
-  @unittest.skipIf(RANGEIFY, "rangeify doesn't implement input buffer limiting")
-  def test_limit_bufs_with_var(self):
-    N = 31
-    with Context(TRACK_MATCH_STATS=0, DEBUG=0):
-      bufs = [Tensor([1]*10).contiguous().realize() for i in range(N)]
-
-    vi = Variable("i", 0, 9).bind(1)
-    vj = Variable("j", 0, 9).bind(2)
-    root = bufs[0][vi] + bufs[0][vj]
-    for X in range(1,N): root = root + bufs[X][vi] + bufs[X][vj]
-    self.assertEqual(root.item(), N * 2)
-
 def swizzle_cnt(u:UOp) -> int:
  return len([x for x in u.toposort() if x.op is Ops.VIEW and len(x.src) != 0 and x.src[0].op not in {Ops.BUFFER, Ops.DEFINE_GLOBAL, Ops.ASSIGN}])

@@ -2084,7 +2109,6 @@ class TestView(unittest.TestCase):
    run_schedule(sched)
    np.testing.assert_equal(b.numpy(), 0)

-  @expect_rangeify_fails
  def test_mask_dim_1(self):
    # mask out dim = 1 works too
    a = Tensor.rand(10, 10).realize()
@@ -2111,7 +2135,6 @@ class TestView(unittest.TestCase):

  # a*VIEW(x), where VIEW(x) = 0
  # x collapses along with its children
-  @unittest.skipIf(RANGEIFY, "this only fails if you run all of TestSchedule, some global tensor map bug?")
  def test_parent_view_collapses(self):
    a = Tensor([1, 2])
    b = Tensor.arange(3).contiguous()
@@ -2193,6 +2216,84 @@ class TestSimplifier(unittest.TestCase):
    assert UPat(Ops.CONST, arg=False).match(sink, {}), f"expected {sink} to collapse to a const False"
    assert sink.shape == a.shape

+tensor_const_pm = PatternMatcher([
+  (UPat(Ops.CONST, src=(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),)),)), lambda: True),
+  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR, src=(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),)))), UPat(Ops.CONST))), lambda: True),
+])
+class TestConst(unittest.TestCase):
+  # ** part 1: basic functionality of a tensor directly created from CONST
+
+  def test_tensor_const(self):
+    a = Tensor(1)
+    print(a.uop)
+    self.assertTrue(tensor_const_pm.rewrite(a.uop))
+
+  def test_tensor_variable(self):
+    vv = UOp.variable("a", 0, 10).bind(1)
+    a = Tensor(vv)
+    print(a.uop)
+    self.assertTrue(tensor_const_pm.rewrite(a.uop))
+
+  def test_const_schedule(self):
+    a = Tensor.ones((4, 4))
+    sched = a.schedule()
+    self.assertEqual(len(sched), 0)
+
+  def test_const_contiguous_schedule(self):
+    # this ends up in the big graph
+    a = Tensor.ones((4,)).contiguous()
+    sched = a.schedule()
+    self.assertEqual(len(sched), 1)
+
+  # ** part 2: scheduler behavior when const folding happens later
+
+  def test_const_folding_no_realize(self):
+    a = Tensor([1, 2, 3, 4])*0
+    sched = a.schedule()
+    self.assertEqual(len(sched), 0)
+
+  def test_src_const_folding(self):
+    with Context(TRACK_MATCH_STATS=0):
+      a = Tensor.full((4,), 1).contiguous().realize()
+      b = Tensor.full((4,), 2).contiguous().realize()
+    mul0 = a*0
+    add = b+mul0
+    sched = add.schedule()
+    self.assertEqual(len(sched), 0)
+    # b+0 and b share the same underlying device memory
+    self.assertIs(add.uop.buffer, b.uop.buffer)
+    self.assertListEqual(add.tolist(), [2, 2, 2, 2])
+
+  def test_src_masked_const_folding(self):
+    with Context(TRACK_MATCH_STATS=0):
+      a = Tensor.full((4,), 1).contiguous().realize()
+      b = Tensor.full((6,), 2).contiguous().realize()
+    mul0 = a*0
+    add = b+mul0.pad((1, 1), value=2)
+    sched = add.schedule()
+    self.assertEqual(len(sched), 1)
+    run_schedule(sched)
+    # add gets assigned to a new buffer
+    self.assertIsNot(add.uop.base.realized, b.uop.base.realized)
+    self.assertListEqual(add.tolist(), [4, 2, 2, 2, 2, 4])
+
+  # ** part 3: Tensor variable bindings
+
+  #@unittest.expectedFailure # TODO: should schedule assert if you try to realize a Variable?
+  def test_var_schedule(self):
+    vv = UOp.variable("a", 0, 10).bind(1)
+    a = Tensor(vv)
+    sched = a.schedule()
+    self.assertEqual(len(sched), 0)
+
+  def test_add_tvar(self):
+    vv = UOp.variable("a", 0, 10).bind(1)
+    a = Tensor(vv)+2
+    sched, var_vals = a.schedule_with_vars()
+    self.assertEqual(len(sched), 1)
+    run_schedule(sched, var_vals)
+    self.assertEqual(a.tolist(), 3)
+
@unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from another device to cpu")
 class TestCopyFolding(unittest.TestCase):
  def test_const_copy_is_free(self):
@@ -2206,7 +2307,6 @@ class TestCopyFolding(unittest.TestCase):
    b = (a*zeros).to("CPU")
    run_schedule(check_schedule(b, 0, filter_sink=False))
    self.assertListEqual(b.tolist(), [0, 0, 0])
-    self.assertEqual(b.device, "CPU")

  def test_alu_after_copy(self):
    a = Tensor.ones((4,)).to("CPU")
@@ -2215,12 +2315,6 @@ class TestCopyFolding(unittest.TestCase):
    add.kernelize()
    assert all_same([x.device for x in add.uop.src]), f"ALU has different devices! {[x.device for x in add.src]}"

-  def test_alu_before_copy(self):
-    buf = Tensor.ones(1).contiguous().realize()
-    a = buf+1
-    b = a.to("CPU")
-    self.assertListEqual(b.tolist(), [2.])
-
  def test_copy_to_same_device(self):
    a = Tensor.empty(4).uop
    b = a.copy_to_device(a.device)
@@ -2237,15 +2331,6 @@ class TestCopyFolding(unittest.TestCase):
    b = schedule_graph_rewrite(b)
    self.assertIs(b.base, a.base)

-  def test_copy_to_same_device_sched(self):
-    a = Tensor.ones(4).contiguous().realize().uop.as_buf()
-    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
-    run_schedule(sched)
-    assert t.uop.is_realized, f"didn't realize Tensor {t}"
-    self.assertListEqual(t.tolist(), [1.,1.,1.,1.])
-
  def test_clone(self):
    a = Tensor.empty(4)
    check_schedule(a.clone(), 1, filter_sink=False)
@@ -2283,14 +2368,6 @@ class TestCopyFolding(unittest.TestCase):
    b.realize()
    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])

-  @expect_nonrangeify_fails
-  def test_permute_on_disk_contiguous(self):
-    with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_buffer())
-    a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}")
-    b = a.reshape(2, 2).permute(1, 0).contiguous().to("CPU")
-    b.realize()
-    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
-
  def test_permute_after_shrink(self):
    a = Tensor.arange(5)
    b = a.shrink(((0, 4),)).reshape(2, 2).permute(1, 0).to("CPU")
@@ -2299,7 +2376,7 @@ class TestCopyFolding(unittest.TestCase):

  # NOTE: disk permute must come after COPY
  # TODO: this is wrong because of the permute
-  @expect_nonrangeify_fails
+  @unittest.expectedFailure
  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())
    a = Tensor.empty(5, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_5_permute')}")
@@ -2430,7 +2507,6 @@ class TestUOpBecome(unittest.TestCase):
    self.assertEqual(add.uop.shape, (8, 2))
    assert add.uop is not add.uop.base

-  @expect_rangeify_fails
  def test_new_flat_buffer(self):
    a = Tensor.empty(4,)
    b = Tensor.empty(4,)
@@ -2442,7 +2518,6 @@ class TestUOpBecome(unittest.TestCase):
  # sometimes we prefer to perform an op before movement ops, in this case we should stack the mops on top of the new buffer

  # NOTE: this expand is not reordered because there's before it to fuse
-  @expect_rangeify_fails
  def test_reorder_expand(self):
    a = Tensor.empty(4, 1)
    b = a.expand(4, 4).reciprocal()
@@ -2457,7 +2532,6 @@ class TestUOpBecome(unittest.TestCase):
    z = (img*x) / y
    check_schedule(z, 1)

-  @expect_rangeify_fails
  def test_become_existing_buffer(self):
    a = Tensor.empty(4, 4)
    b = a*1
@@ -2485,7 +2559,6 @@ class TestUOpBecome(unittest.TestCase):
    check_schedule(b, 0)
    assert UPat(Ops.CONST, arg=0).match(b.uop.base, {}) # scheduling replaces the tensor uop with a VIEW(BUFFER)

-  @expect_rangeify_fails
  def test_become_const_in_view(self):
    # if we shrink the base down to a size 0, only the VIEW becomes CONST, base is unchanged.
    add = Tensor.empty(2, 2)+Tensor.empty(2, 2)
@@ -2503,7 +2576,6 @@ class TestUOpBecome(unittest.TestCase):
    assert UPat(Ops.CONST, arg=3).match(const_add.uop.base, {})

  # tensors can become another realized tensor source
-  @expect_rangeify_fails
  def test_become_existing_buf_simple(self):
    a = Tensor.empty(4, 4)
    b = a+0
@@ -2512,14 +2584,12 @@ class TestUOpBecome(unittest.TestCase):
    self.assertIs(a.uop, b.uop)

  # they can also chain other movement ops on top of the tensor source
-  @expect_rangeify_fails
  def test_become_existing_buf_view(self):
    a = Tensor.empty(4, 4)
    b = a.permute((1, 0))+0
    check_schedule(b, 0)
    self.assertEqual(b.uop.st, a.uop.permute((1, 0)).st)

-  @expect_rangeify_fails
  def test_become_existing_buf_view_alt(self):
    a = Tensor.empty(4, 4)
    b = a.permute((1, 0)).reshape((8, 2))+0
@@ -2527,7 +2597,6 @@ class TestUOpBecome(unittest.TestCase):
    self.assertEqual(b.uop.st, a.uop.permute((1, 0)).reshape((8, 2)).st)

  # they can also have other base parents that simplified, in that case we just backtrack to the chained mops
-  @expect_rangeify_fails
  def test_become_existing_buf_complex(self):
    a = Tensor.empty(4, 4)
    b = (a.permute((1, 0))+0).reshape((8, 2))+0
@@ -2535,7 +2604,6 @@ class TestUOpBecome(unittest.TestCase):
    self.assertEqual(b.uop.st, a.uop.permute((1, 0)).reshape((8, 2)).st)
    assert b.uop.base.op is Ops.BUFFER

-  @expect_rangeify_fails
  def test_become_multiple_choices(self):
    a = Tensor.empty(16)
    b = (a.reshape(1, 1, 4, 1, 4)+0).reshape(1, 1, 4, 4).shrink(((0, 1), (0, 1), (0, 3), (0, 3)))+0
@@ -2547,7 +2615,6 @@ class TestUOpBecome(unittest.TestCase):
    assert b.uop is c.uop
    assert UPat(Ops.VIEW, src=(UPat(Ops.BUFFER),)).match(c.uop, {})

-  @expect_rangeify_fails
  def test_setitem_becomes_subbuffer(self):
    a = Tensor.full((4,), 2.).contiguous().realize()
    b = a.shrink(((0, 2),)).assign(Tensor.full((2,), 1.0))
@@ -1,6 +1,4 @@
 import unittest
-import random
-from os import getenv
 from tinygrad import Tensor, TinyJit, Variable, dtypes
 from tinygrad.helpers import Context
 import numpy as np
@@ -167,41 +165,6 @@ class TestSetitem(unittest.TestCase):
    t[idx] = val
    self.assertEqual(t.tolist(), [val]*idx_size+[idx_size])

-  def test_setitem_advanced_indexing(self):
-    # Example from https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing
-    t = Tensor.zeros(10,20,30,40,50).contiguous()
-    ind_1 = Tensor([5,3,7,8])
-    ind_2 = Tensor([[[0],[1],[2]],[[3],[4],[5]]])
-    v = Tensor.arange(2*3*4*10*30*50).reshape(2,3,4,10,30,50)
-    t[:, ind_1, :, ind_2, :] = v
-    n = np.zeros((10,20,30,40,50))
-    n[:, ind_1.numpy(), :, ind_2.numpy(), :] = v.numpy()
-    np.testing.assert_allclose(t.numpy(), n)
-
-  def test_setitem_2d_tensor_indexing(self):
-    t = Tensor.zeros(2).contiguous()
-    index = Tensor([[0, 1], [1,0]])
-    v = Tensor.arange(2*2).reshape(2, 2).contiguous()
-    t[index] = v
-    n = np.zeros((2,))
-    n[index.numpy()] = v.numpy()
-    np.testing.assert_allclose(t.numpy(), n)
-
-  @unittest.skip("slow")
-  def test_setitem_tensor_indexing_fuzz(self):
-    random.seed(getenv("SEED", 42))
-    for _ in range(getenv("ITERS", 100)):
-      size = random.randint(5, 10)
-      d0, d1, d2 = random.randint(1,5), random.randint(1,5), random.randint(1,5)
-      t = Tensor.zeros(size).contiguous()
-      n = np.zeros((size,))
-      index = Tensor.randint((d0, d1, d2), low=0, high=size)
-      v = Tensor.arange(d0*d1*d2).reshape(d0, d1, d2)
-      t[index] = v
-      n[index.numpy()] = v.numpy()
-      np.testing.assert_allclose(t.numpy(), n, err_msg=f"failed with index={index.numpy().tolist()} and v={v.numpy().tolist()}")
-
-
 class TestWithGrad(unittest.TestCase):
  def test_no_requires_grad_works(self):
    z = Tensor.rand(8, 8)
@@ -2,7 +2,7 @@ import unittest
 import numpy as np
 from tinygrad import Tensor, GlobalCounters, Context, Device
 from tinygrad.dtype import DTypeLike, dtypes
-from tinygrad.helpers import DEBUG, get_single_element, RANGEIFY
+from tinygrad.helpers import DEBUG, get_single_element
 from tinygrad.engine.realize import lower_schedule_item
 from tinygrad.device import is_dtype_supported

@@ -30,7 +30,7 @@ def single_kernel_softmax(x_in:Tensor, axis=-1, dtype:DTypeLike|None=None) -> Te
 def run_one_schedule_item(out): lower_schedule_item(get_single_element(out.schedule())).run()

 class TestFuse(unittest.TestCase):
-  def _test_fuse(self, fxn, *args, atol=1e-6, allow_multiple=False, **kwargs):
+  def _test_fuse(self, fxn, *args, atol=1e-7, allow_multiple=False, **kwargs):
    GlobalCounters.reset()
    out_single = fxn(*args, **kwargs).fuse()
    if not allow_multiple: run_one_schedule_item(out_single)
@@ -39,17 +39,14 @@ class TestFuse(unittest.TestCase):
    np_multi = fxn(*args, **kwargs).numpy()
    np.testing.assert_allclose(np_single, np_multi, atol=atol)

-  @unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
  def test_fuse_norm(self):
    a = Tensor.rand(50,50).realize()
    self._test_fuse(lambda a: a / a.mean(axis=1), a)

-  @unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
  def test_fuse_argmax(self):
    a = Tensor.rand(50,50).realize()
    self._test_fuse(lambda a: a.argmax(axis=-1), a)

-  @unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
  def test_fuse_softmax(self):
    a = Tensor.rand(50,50).realize()
    self._test_fuse(lambda a: a.softmax(axis=-1), a)
@@ -60,7 +57,6 @@ class TestFuse(unittest.TestCase):
    self._test_fuse(lambda a,b: ((a@b).relu()+a).contiguous().softmax(axis=-1), a,b, allow_multiple=True)

  @unittest.skipUnless(is_dtype_supported(dtypes.float16, Device.DEFAULT), f"no float16 on {Device.DEFAULT}")
-  @unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
  def test_fuse_softmax_dtype(self):
    a = Tensor.rand(50,50).realize()
    self._test_fuse(lambda a: a.softmax(axis=-1, dtype='half'), a, atol=3e-4)
@@ -68,7 +64,6 @@ class TestFuse(unittest.TestCase):
  def test_fuse_arange_eye(self):
    self._test_fuse(lambda: Tensor.arange(10).reshape(10,1).expand(10,10) == Tensor.arange(10).reshape(1,10).expand(10,10))

-  @unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
  def test_double_gemm(self):
    N = 32
    with Context(TRACK_MATCH_STATS=0, DEBUG=0):
@@ -91,7 +86,6 @@ class TestFuse(unittest.TestCase):
      return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
    self._test_fuse(embedding, a, atol=1e-5)

-  @unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
  def test_attention_kernel_count(self):
    wq = Tensor.empty(32, 32)
    wk = Tensor.empty(32, 32)
@@ -104,7 +98,6 @@ class TestFuse(unittest.TestCase):
    s = attn.schedule()
    self.assertEqual(len(s), 4) # 3 matmul and 1 attention

-  @unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
  def test_flash_attention(self):
    BS = 4
    HEADS = 2
@@ -172,7 +165,6 @@ class TestSoftmaxFusion(unittest.TestCase):

    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)

-  @unittest.skipIf(0<RANGEIFY<2, "needs RANGEIFY>1")
  def test_auto_softmax(self):
    print("*** softmax ***")
    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
@@ -2,7 +2,6 @@ import unittest

 from test.helpers import assert_jit_cache_len
 from tinygrad import Variable, Tensor, TinyJit
-from tinygrad.helpers import RANGEIFY
 import numpy as np

 class TestSymbolicJit(unittest.TestCase):
@@ -12,23 +11,11 @@ class TestSymbolicJit(unittest.TestCase):
    a = Tensor.rand(3, 10)
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(a[:, :vi])[:3, :i].numpy()
+      symbolic = jf(a[:, :vi]).reshape(3, i).numpy()
      expected = f(a[:, :i]).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
    assert_jit_cache_len(jf, 1)

-  def test_plus1_pad(self):
-    # TODO: without contiguous, the pad is not captured in jit
-    def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).contiguous().realize()
-    jf = TinyJit(f)
-    a = Tensor.rand(3, 10)
-    for i in range(1, 5):
-      vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(a[:, :vi]).numpy()
-      expected = f(a[:, :i]).numpy()
-      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-    assert_jit_cache_len(jf, 1 if RANGEIFY else 2) # one add and one pad, can be one kernel?
-
  def test_add(self):
    def f(a, b): return (a+b).realize()
    jf = TinyJit(f)
@@ -36,8 +23,7 @@ class TestSymbolicJit(unittest.TestCase):
    b = Tensor.rand(3, 10)
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(a[:, :vi], b[:, :vi])
-      symbolic = symbolic[:3, :i].numpy()
+      symbolic = jf(a[:, :vi], b[:, :vi]).reshape(3, i).numpy()
      expected = f(a[:, :i], b[:, :i]).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
    assert_jit_cache_len(jf, 1)
@@ -77,10 +63,10 @@ class TestSymbolicJit(unittest.TestCase):
    v = Tensor.rand(2, 10, 4, 8)
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(q, k[:, :vi], v[:, :vi])[:2, :4, :1, :8].numpy()
+      symbolic = jf(q, k[:, :vi], v[:, :vi]).reshape(2, 4, 1, 8).numpy()
      expected = f(q, k[:, :i], v[:, :i]).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-    assert_jit_cache_len(jf, 4 if RANGEIFY else 5)
+    assert_jit_cache_len(jf, 5)

  def test_cat_dim0(self):
    def f(a, b): return a.cat(b, dim=0).realize()
@@ -89,7 +75,7 @@ class TestSymbolicJit(unittest.TestCase):
    b = Tensor.rand(2, 3)
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(a[:vi], b)[:i+2, :3].numpy()
+      symbolic = jf(a[:vi], b).reshape(i+2, 3).numpy()
      expected = f(a[:i], b).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
    assert_jit_cache_len(jf, 1)
@@ -101,7 +87,7 @@ class TestSymbolicJit(unittest.TestCase):
    b = Tensor.rand(3, 2)
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(a[:, :vi], b)[:3, :i+2].numpy()
+      symbolic = jf(a[:, :vi], b).reshape(3, i+2).numpy()
      expected = f(a[:, :i], b).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
    assert_jit_cache_len(jf, 1)
@@ -115,7 +101,7 @@ class TestSymbolicJit(unittest.TestCase):
      for j in range(2, 5):
        vi = Variable("i", 1, 10).bind(i)
        vj = Variable("j", 1, 10).bind(j)
-        symbolic = jf(a[:vi], b[:vj])[:i+j, :3].numpy()
+        symbolic = jf(a[:vi], b[:vj]).reshape(i+j, 3).numpy()
        expected = f(a[:i], b[:j]).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
    assert_jit_cache_len(jf, 1)
@@ -129,7 +115,7 @@ class TestSymbolicJit(unittest.TestCase):
      for j in range(2, 5):
        vi = Variable("i", 1, 10).bind(i)
        vj = Variable("j", 1, 10).bind(j)
-        symbolic = jf(a[:, :vi], b[:, :vj])[:3, :i+j].numpy()
+        symbolic = jf(a[:, :vi], b[:, :vj]).reshape(3, i+j).numpy()
        expected = f(a[:, :i], b[:, :j]).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
    assert_jit_cache_len(jf, 1)
@@ -143,7 +129,7 @@ class TestSymbolicJit(unittest.TestCase):
      for j in range(2, 5):
        vi = Variable("i", 1, 10).bind(i)
        vj = Variable("j", 1, 10).bind(j)
-        symbolic = jf(a[:vi, :], b[:, :vj])[:i, :j].numpy()
+        symbolic = jf(a[:vi, :], b[:, :vj]).reshape(i, j).numpy()
        expected = f(a[:i, :], b[:, :j]).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
    assert_jit_cache_len(jf, 1)
@@ -157,7 +143,7 @@ class TestSymbolicJit(unittest.TestCase):
      for j in range(2, 5):
        vi = Variable("i", 1, 10).bind(i)
        vj = Variable("j", 1, 10).bind(j)
-        symbolic = jf(a[:vj, :], b[:, :vi])[:j, :i].numpy()
+        symbolic = jf(a[:vj, :], b[:, :vi]).reshape(j, i).numpy()
        expected = f(a[:j, :], b[:, :i]).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
    assert_jit_cache_len(jf, 1)
@@ -209,8 +195,8 @@ class TestSymbolicJit(unittest.TestCase):
      vi = Variable("i", 1, 10).bind(i)
      a = Tensor.ones(vi, 11).contiguous()
      symbolic = a[:, 1:2]
-      symbolic = jf(symbolic)[:i, :1].numpy()
-      expected = f(a[:i, :][:, 1:2]).numpy()
+      symbolic = jf(symbolic).reshape(i, 1).numpy()
+      expected = f(a.reshape(i, 11)[:, 1:2]).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
    assert_jit_cache_len(jf, 1)

@@ -245,7 +231,7 @@ class TestSymbolicJit(unittest.TestCase):
      expected = b[:i].mean(0).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
      # axis = 1
-      symbolic = jf1(c[:vi])[:i].numpy()
+      symbolic = jf1(c[:vi]).reshape(i).numpy()
      expected = c[:i].mean(1).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

@@ -268,11 +254,11 @@ class TestSymbolicJit(unittest.TestCase):
        expected = a[:i, :j].mean().numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
        # axis = 0
-        symbolic = jf0(b[:vi, :vj])[:j].numpy()
+        symbolic = jf0(b[:vi, :vj]).reshape(j).numpy()
        expected = b[:i, :j].mean(0).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
        # axis = 1
-        symbolic = jf1(c[:vi, :vj])[:i].numpy()
+        symbolic = jf1(c[:vi, :vj]).reshape(i).numpy()
        expected = c[:i, :j].mean(1).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

@@ -297,7 +283,7 @@ class TestSymbolicJit(unittest.TestCase):
      expected = b[:i].var(0).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
      # axis = 1
-      symbolic = jf1(c[:vi])[:i].numpy()
+      symbolic = jf1(c[:vi]).reshape(i).numpy()
      expected = c[:i].var(1).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

@@ -320,11 +306,11 @@ class TestSymbolicJit(unittest.TestCase):
        expected = a[:i, :j].var().numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
        # axis = 0
-        symbolic = jf0(b[:vi, :vj])[:j].numpy()
+        symbolic = jf0(b[:vi, :vj]).reshape(j).numpy()
        expected = b[:i, :j].var(0).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
        # axis = 1
-        symbolic = jf1(c[:vi, :vj])[:i].numpy()
+        symbolic = jf1(c[:vi, :vj]).reshape(i).numpy()
        expected = c[:i, :j].var(1).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

@@ -13,16 +13,7 @@ class TestSymbolicOps(unittest.TestCase):
    a = Tensor.rand(3, 10)
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
-      symbolic = f(a[:, :vi])[:3, :i].numpy()
-      expected = f(a[:, :i]).numpy()
-      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-
-  def test_plus1_pad(self):
-    def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
-    a = Tensor.rand(3, 10)
-    for i in range(1, 5):
-      vi = Variable("i", 1, 10).bind(i)
-      symbolic = f(a[:, :vi]).numpy()
+      symbolic = f(a[:, :vi]).reshape(3, i).numpy()
      expected = f(a[:, :i]).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

@@ -32,7 +23,7 @@ class TestSymbolicOps(unittest.TestCase):
    b = Tensor.rand(3, 10)
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
-      symbolic = f(a[:, :vi], b[:, :vi])[:, :i].numpy()
+      symbolic = f(a[:, :vi], b[:, :vi]).reshape(3, i).numpy()
      expected = f(a[:, :i], b[:, :i]).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

@@ -55,7 +46,7 @@ class TestSymbolicOps(unittest.TestCase):
      vi = Variable("i", 1, 10).bind(i) if use_symbolic else i
      Tensor.realize(q, k, v)
      GlobalCounters.reset()
-      symbolic = f(q, k[:, :vi, :, :], v[:, :vi, :, :])[:2, :4, :1, :8].numpy()
+      symbolic = f(q, k[:, :vi, :, :], v[:, :vi, :, :]).reshape(2, 4, 1, 8).numpy()
      expected = f(q, k[:, :i, :, :], v[:, :i, :, :]).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

@@ -94,7 +85,7 @@ class TestSymbolicOps(unittest.TestCase):
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
      b = Tensor.rand(2, 3)
-      symbolic = f(a[:vi, :], b)[:i+2, :3].numpy()
+      symbolic = f(a[:vi, :], b).reshape(i+2, 3).numpy()
      expected = f(a[:i, :], b).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

@@ -104,7 +95,7 @@ class TestSymbolicOps(unittest.TestCase):
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
      b = Tensor.rand(3, 2)
-      symbolic = f(a[:, :vi], b)[:3, :i+2].numpy()
+      symbolic = f(a[:, :vi], b).reshape(3, i+2).numpy()
      expected = f(a[:, :i], b).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

@@ -116,7 +107,7 @@ class TestSymbolicOps(unittest.TestCase):
      for j in range(2, 5):
        vi = Variable("i", 1, 10).bind(i)
        vj = Variable("j", 1, 10).bind(j)
-        symbolic = f(a[:vi, :], b[:vj, :])[:i+j, :3].numpy()
+        symbolic = f(a[:vi, :], b[:vj, :]).reshape(i+j, 3).numpy()
        expected = f(a[:i, :], b[:j, :]).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

@@ -128,41 +119,50 @@ class TestSymbolicOps(unittest.TestCase):
      for j in range(2, 5):
        vi = Variable("i", 1, 10).bind(i)
        vj = Variable("j", 1, 10).bind(j)
-        symbolic = f(a[:, :vi], b[:, :vj])[:3, :i+j].numpy()
+        symbolic = f(a[:, :vi], b[:, :vj]).reshape(3, i+j).numpy()
        expected = f(a[:, :i], b[:, :j]).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

  def test_two_vars_plus1_ij(self):
    def f(a, b): return (a@b+1).realize()
-    a = Tensor.rand(10, 3).realize()
-    b = Tensor.rand(3, 10).realize()
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(3, 10)
    for i in range(2, 5):
      for j in range(2, 5):
        vi = Variable("i", 1, 10).bind(i)
        vj = Variable("j", 1, 10).bind(j)
-        symbolic = f(a[:vi, :], b[:, :vj])[:i, :j].numpy()
+        symbolic = f(a[:vi, :], b[:, :vj]).reshape(i, j).numpy()
        expected = f(a[:i, :], b[:, :j]).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

  def test_two_vars_plus1_ji(self):
    # reverse the order of variables
    def f(a, b): return (a@b+1).realize()
-    a = Tensor.rand(10, 3).realize()
-    b = Tensor.rand(3, 10).realize()
+    a = Tensor.rand(10, 3)
+    b = Tensor.rand(3, 10)
    for i in range(2, 5):
      for j in range(2, 5):
        vi = Variable("i", 1, 10).bind(i)
        vj = Variable("j", 1, 10).bind(j)
-        symbolic = f(a[:vj, :], b[:, :vi])[:j, :i].numpy()
+        symbolic = f(a[:vj, :], b[:, :vi]).reshape(j, i).numpy()
        expected = f(a[:j, :], b[:, :i]).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

+  def test_reshape_from_symbolic(self):
+    a = Tensor.rand(30)
+    for i in range(3, 5):
+      vi = Variable("i", 3, 10).bind(i)
+      symbolic = a[:vi*3].reshape((3, 3)).numpy()
+      # To match symbolic reshape (potential implicit shrink), we need a shrink
+      expected = a[:i*3].shrink(((0, 9),)).reshape((3, 3)).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
  def test_invalid_symbolic_reshape(self):
    a = Tensor.rand(30)
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
      # Cannot reshape into symbolic from non-symbolic
-      with self.assertRaises(ValueError): a.reshape((3, vi))
+      with self.assertRaises(AssertionError): a.reshape((3, vi))

  def test_shrink(self):
    for i in range(1, 5):
@@ -178,7 +178,6 @@ class TestSymbolicOps(unittest.TestCase):
      vi = Variable("i", 1, 10).bind(i)
      a = Tensor.rand(7, 11)
      symbolic = a[3:5, vi:vi+2]
-      print(symbolic.shape)
      symbolic = symbolic.numpy()
      expected = a[3:5, i:i+2].numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
@@ -187,7 +186,7 @@ class TestSymbolicOps(unittest.TestCase):
    a = Tensor.rand(7, 11)
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
-      symbolic = a[3:5, :vi:1][:2, :i].numpy()
+      symbolic = a[3:5, :vi:1].reshape(2, i).numpy()
      expected = a[3:5, :i:1].numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

@@ -195,7 +194,7 @@ class TestSymbolicOps(unittest.TestCase):
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
      a = Tensor(1).unsqueeze(0).pad((0, 1)).unsqueeze(0)
-      symbolic = a.expand(vi, 2)[:i, :2].numpy()
+      symbolic = a.expand(vi, 2).reshape(i, 2).numpy()
      expected = a.expand(i, 2).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

@@ -203,8 +202,8 @@ class TestSymbolicOps(unittest.TestCase):
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
      a = Tensor.ones(vi, 11).contiguous()
-      symbolic = a[:, 1:2][:i, :1].numpy()
-      expected = Tensor.ones(i, 11)[:, 1:2].numpy()
+      symbolic = a[:, 1:2].reshape(i, 1).numpy()
+      expected = a.reshape(i, 11)[:, 1:2].numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

  def test_ones_sum(self):
@@ -221,11 +220,7 @@ class TestSymbolicOps(unittest.TestCase):
      vi = Variable("i", 1, 10).bind(i)
      for axis in [None, 0, 1]:
        expected = a[:i].mean(axis).numpy()
-        symbolic = a[:vi].mean(axis)
-        if axis is None:
-          symbolic = symbolic.numpy()
-        else:
-          symbolic = symbolic[:expected.shape[0]].numpy()
+        symbolic = a[:vi].mean(axis).reshape(expected.shape).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

  def test_mean_2d(self):
@@ -236,11 +231,7 @@ class TestSymbolicOps(unittest.TestCase):
        vj = Variable("j", 1, 10).bind(j)
        for axis in [None, 0, 1]:
          expected = a[:i, :j].mean(axis).numpy()
-          symbolic = a[:vi, :vj].mean(axis)
-          if axis is None:
-            symbolic = symbolic.numpy()
-          else:
-            symbolic = symbolic[:expected.shape[0]].numpy()
+          symbolic = a[:vi, :vj].mean(axis).reshape(expected.shape).numpy()
          np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

  def test_var(self):
@@ -249,11 +240,7 @@ class TestSymbolicOps(unittest.TestCase):
      vi = Variable("i", 1, 10).bind(i)
      for axis in [None, 0, 1]:
        expected = a[:i].var(axis).numpy()
-        symbolic = a[:vi].var(axis)
-        if axis is None:
-          symbolic = symbolic.numpy()
-        else:
-          symbolic = symbolic[:expected.shape[0]].numpy()
+        symbolic = a[:vi].var(axis).reshape(expected.shape).numpy()
        np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

  def test_var_2d(self):
@@ -264,11 +251,7 @@ class TestSymbolicOps(unittest.TestCase):
        vj = Variable("j", 1, 10).bind(j)
        for axis in [None, 0, 1]:
          expected = a[:i, :j].var(axis).numpy()
-          symbolic_result = a[:vi, :vj].var(axis)
-          if axis is None:
-            symbolic = symbolic_result.numpy()
-          else:
-            symbolic = symbolic_result[:expected.shape[0]].numpy()
+          symbolic = a[:vi, :vj].var(axis).reshape(expected.shape).numpy()
          np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)

  def test_bitcast_down(self):
@@ -276,11 +259,7 @@ class TestSymbolicOps(unittest.TestCase):
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
      expected = a[:i].bitcast(dtypes.uint8).numpy()
-      symbolic_result = a[:vi].bitcast(dtypes.uint8)
-      if len(expected.shape) == 2:
-        symbolic = symbolic_result[:expected.shape[0], :expected.shape[1]].numpy()
-      else:
-        symbolic = symbolic_result[:].numpy()
+      symbolic = a[:vi].bitcast(dtypes.uint8).reshape(expected.shape).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)

  @unittest.skipUnless(is_dtype_supported(dtypes.uint64), "no uint64")
@@ -289,11 +268,7 @@ class TestSymbolicOps(unittest.TestCase):
    for i in range(1, 5):
      vi = Variable("i", 1, 10).bind(i)
      expected = a[:i].bitcast(dtypes.uint64).numpy()
-      symbolic_result = a[:vi].bitcast(dtypes.uint64)
-      if len(expected.shape) == 2:
-        symbolic = symbolic_result[:expected.shape[0], :expected.shape[1]].numpy()
-      else:
-        symbolic = symbolic_result[:].numpy()
+      symbolic = a[:vi].bitcast(dtypes.uint64).reshape(expected.shape).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)

  @unittest.expectedFailure
@@ -4,7 +4,7 @@ import torch
 import unittest, copy, mmap, random, math, array
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _METADATA
-from tinygrad.helpers import getenv, temp, mv_address, RANGEIFY
+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
@@ -516,6 +516,10 @@ class TestTinygrad(unittest.TestCase):
    print(c)

  def test_env_overwrite_default_device(self):
+    subprocess.run(['DISK=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT != \\"DISK\\""'],
+                    shell=True, check=True)
+    subprocess.run(['NPY=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT != \\"NPY\\""'],
+                    shell=True, check=True)
    subprocess.run([f'{Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
                    shell=True, check=True)
    subprocess.run([f'DISK=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
@@ -550,11 +554,6 @@ class TestTinygrad(unittest.TestCase):
  def test_shrink(self):
    t = Tensor.arange(32).contiguous().realize()
    self.assertListEqual(t[16:20].tolist(), [16,17,18,19])
-    self.assertListEqual(t.shrink_to(16).tolist(), list(range(16)))
-    t = t.reshape(4, 8).contiguous().realize()
-    self.assertListEqual(t.shrink_to(2, 2).tolist(), [[0, 1], [8, 9]])
-    with self.assertRaises(ValueError): t.shrink_to(2)
-    with self.assertRaises(ValueError): t.shrink_to(2, 2, 2)

@unittest.skip("this test is just flaky, sync issue")
 class TestMoveTensor(unittest.TestCase):
@@ -649,22 +648,17 @@ class TestZeroShapeTensor(unittest.TestCase):

  def test_pad(self):
    t = Tensor.rand(3, 2, 0).pad((None, None, (1, 1)), value=1)
-    self.assertEqual(t.shape, (3, 2, 2))
+    assert t.shape == (3, 2, 2)
    np.testing.assert_equal(t.numpy(), np.ones((3, 2, 2)))

    t = Tensor.rand(3, 2, 0).pad((None, (1, 1), None), value=1)
-    self.assertEqual(t.shape, (3, 4, 0))
+    assert t.shape == (3, 4, 0)
    np.testing.assert_equal(t.numpy(), np.ones((3, 4, 0)))

    t = Tensor.rand(3, 2, 0).pad(((1, 1), None, None), value=1)
-    self.assertEqual(t.shape, (5, 2, 0))
+    assert t.shape == (5, 2, 0)
    np.testing.assert_equal(t.numpy(), np.ones((5, 2, 0)))

-    np.testing.assert_equal(Tensor([1, 2]).pad_to(4).numpy(), [1, 2, 0, 0])
-    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3).numpy(), [[1, 2, 0], [0, 0, 0]])
-    with self.assertRaises(TypeError): Tensor([1, 2]).pad_to(2, 3)
-    with self.assertRaises(TypeError): Tensor([[1, 2]]).pad_to(3)
-
  def test_shrink_into_zero(self):
    t = Tensor.rand(3, 4).realize()
    assert t.shrink((None, (2, 2))).realize().shape == (3, 0)
@@ -871,18 +865,11 @@ class TestTensorMetadata(unittest.TestCase):
    self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid")
    self.assertTrue(y.grad.uop.metadata[0].backward)
    si = Tensor.schedule(out, x.grad, y.grad)[-1]
-    if not RANGEIFY:
-      self.assertEqual(len(si.metadata), 4, f"failed with {si.metadata}")
-      self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "__mul__", "relu"})
-      bw = [m for m in si.metadata if m.backward]
-      self.assertEqual(len(bw), 2)
-      self.assertEqual(bw[0].name, "sigmoid")
-    else:
-      self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}")
-      self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "relu"})
-      bw = [m for m in si.metadata if m.backward]
-      self.assertEqual(len(bw), 1)
-      self.assertEqual(bw[0].name, "sigmoid")
+    self.assertEqual(len(si.metadata), 4, f"failed with {si.metadata}")
+    self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "__mul__", "relu"})
+    bw = [m for m in si.metadata if m.backward]
+    self.assertEqual(len(bw), 2)
+    self.assertEqual(bw[0].name, "sigmoid")

 class TestIdxUpcast(unittest.TestCase):
  def _find_op(self, ast: UOp, op: Ops):
@@ -935,10 +922,6 @@ class TestIdxUpcast(unittest.TestCase):
    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):
@@ -38,7 +38,7 @@ class TestTensorVariable(unittest.TestCase):
    vv = Variable("a", 1, 10).bind(2)
    vv2 = Variable("b", 1, 10).bind(2)
    t = Tensor.ones(10, 10).contiguous()[:vv2, :vv]
-    ret = t.mean(axis=1)[:2].reshape(2, 1).numpy()
+    ret = t.mean(axis=1).reshape(2, 1).numpy()
    assert np.all(ret == 1)

  def test_symbolic_mean_2d_add(self):
@@ -66,25 +66,25 @@ class TestTensorVariable(unittest.TestCase):
  def test_symbolic_arange(self):
    vv = Variable("a", 1, 10)
    ret = Tensor.arange(0, vv.bind(4))
-    self.assertListEqual(ret[:4].tolist(), [0,1,2,3])
+    self.assertListEqual(ret.reshape(4).tolist(), [0,1,2,3])

  def test_symbolic_arange_sym_start(self):
    vv = Variable("a", 1, 6)
    ret = Tensor.arange(vv.bind(4), 7)
-    self.assertListEqual(ret[:3].tolist(), [4,5,6])
+    self.assertListEqual(ret.reshape(3).tolist(), [4,5,6])

  # TODO: add vmin/vmax pattern for symbolic denominator
  @unittest.expectedFailure
  def test_symbolic_arange_sym_step(self):
    vv = Variable("step", 1, 3)
    ret = Tensor.arange(0, 10, vv.bind(2))
-    self.assertListEqual(ret[:5].tolist(), [0,2,4,6,8])
+    self.assertListEqual(ret.reshape(5).tolist(), [0,2,4,6,8])

  def test_symbolic_arange_two_vars(self):
    begin = Variable("b", 1, 5)
    end = Variable("e", 6, 10)
    ret = Tensor.arange(begin.bind(4), end.bind(7))
-    self.assertListEqual(ret[:3].tolist(), [4,5,6])
+    self.assertListEqual(ret.reshape(3).tolist(), [4,5,6])

  def test_variable_empty(self):
    v = Variable("i", 1, 10)
@@ -95,7 +95,7 @@ class TestTiny(unittest.TestCase):
    ones = Tensor.ones(10).contiguous()
    for s in [2,5]:
      ret = ones[:i.bind(s)] + 1
-      self.assertListEqual(ret.contiguous()[:s].tolist(), [2.0]*s)
+      self.assertListEqual(ret.contiguous().reshape(s).tolist(), [2.0]*s)

  def test_symbolic_reduce(self):
    i = Variable('i', 1, 10)
@@ -419,7 +419,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_fold(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
+    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 128)
    ld = d0.index(ridx0, ridx0<50).load()
    w = (ridx0<50).where(ld, 5)
    uops = to_uops_list([w])
@@ -429,7 +429,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_folds_swapped_branches(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
+    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 128)
    ld = d0.index(ridx0, (ridx0<50).logical_not()).load()
    w = (ridx0<50).where(5, ld)
    uops = to_uops_list([w])
@@ -437,25 +437,13 @@ class TestUOpGraph(unittest.TestCase):
      assert u.op is not Ops.WHERE
      if u.op is Ops.LOAD: assert u.src[1].arg==5

-  def test_where_in_store_becomes_gate(self):
-    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
-    idx = d0.index(ridx0)
-    ld = idx.load()
-    val = (ridx0<50).where(5, ld)
-    st = idx.store(val, ridx0)
-    uops = to_uops_list([st])
-    for u in uops:
-      assert u.op is not Ops.WHERE
-      if u.op is Ops.STORE: assert u.src[1].arg==5
-
  def test_load_idx_becomes_int(self):
    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 1)
-    l0 = UOp(Ops.LOAD, dtypes.long, (d0.index(UOp.const(dtypes.int, 0)),)).cast(dtypes.index)
+    l0 = UOp(Ops.LOAD, dtypes.long, (d0.index(UOp.const(dtypes.int, 0)),))
    idx = l0 * 600
    valid = (l0<-1).ne(True)&(l0<3000)
-    l1 = UOp(Ops.LOAD, dtypes.long, (d1.index(idx.valid(valid)),))
+    l1 = UOp(Ops.LOAD, dtypes.long, (d1.index(idx, valid),))
    uops = to_uops_list([l1])
    for u in uops:
      if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int)
@@ -581,20 +569,12 @@ class TestUOpGraph(unittest.TestCase):
      glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 0)
      gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 42),), "gidx0")
      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, gidx0<8),)).cast(dtypes.index)
-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<32)),))
+      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<32)),)).cast(dtypes.index)
      to_uops_list([ld1])

      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<64)),))
      with self.assertRaises(RuntimeError): to_uops_list([ld1])

-  def test_bounds_with_loaded_bool(self):
-    glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
-    glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(8), (), 0)
-    gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 16),), "gidx0")
-    ld0 = glbl0.index(gidx0).load()
-    ld1 = glbl1.index(gidx0.valid(ld0)).load()
-    with self.assertRaises(RuntimeError): to_uops_list([ld1])
-
  def test_fold_gated_load(self):
    glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
    glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 1)
@@ -544,6 +544,86 @@ class TestUopsObject(unittest.TestCase):
    with Timing("create 10k uops:"): ret = [UOp(Ops.CONST, dtypes.int, arg=10000000+i) for i in range(10000)]
    assert len(ret) == 10000

+
+class TestShapeSpec(unittest.TestCase):
+  # ** CONST is CONST(VIEW(DEVICE)) -> RESHPAE -> EXPAND
+
+  def test_expanded_const(self):
+    a = Tensor(1).uop
+    self.assertEqual(a.st, ShapeTracker.from_shape(()))
+    a = Tensor.ones((4, 4)).uop
+    self.assertEqual(a.st, ShapeTracker.from_shape(()).reshape((1,1)).expand((4,4)))
+
+  # NOTE: CONST ShapeTracker comes from its source
+  def test_scalar_const(self):
+    a = Tensor(0).uop
+    self.assertEqual(a.st, ShapeTracker.from_shape(()))
+
+  def test_scalar_var(self):
+    vv = UOp.variable("a", 1, 4).bind(2)
+    t = Tensor(vv).uop
+    self.assertEqual(t.st, ShapeTracker.from_shape(()))
+
+  # ** ASSIGN is ASSIGN(VIEW(BUFFER), new_val)
+
+  def test_assign_flat(self):
+    buffer = Tensor.arange(4).realize()
+    a = buffer.assign(Tensor.zeros((4,), dtype=dtypes.int))
+    assign_pattern = UPat(Ops.ASSIGN, src=(UPat(Ops.BUFFER), UPat()))
+    assert assign_pattern.match(a.uop, {})
+    a.realize()
+    self.assertEqual(buffer.tolist(), [0, 0, 0, 0])
+
+  def test_assign_permuted(self):
+    buffer = Tensor.arange(4).reshape(2, 1, 2).contiguous().realize()
+    a = buffer.permute((1, 2, 0)).assign(Tensor.arange(4).reshape(1, 2, 2).contiguous())
+    a.realize()
+    self.assertEqual(buffer.tolist(), [[[0, 2]], [[1, 3]]])
+
+  def test_assign_reshaped(self):
+    buffer = Tensor.ones((4,)).contiguous().realize()
+    a = buffer.reshape((2, 2)).assign(Tensor.zeros((2, 2)))
+    assign_pattern = UPat(Ops.ASSIGN, src=(UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER))), UPat()))
+    assert assign_pattern.match(a.uop, {})
+    a.realize()
+    self.assertEqual(buffer.tolist(), [0, 0, 0, 0])
+
+  # setitem is a partial assign
+  def test_setitem(self):
+    a = Tensor.ones((4,)).contiguous().realize()
+    assign = a.shrink(((1, 2),)).assign(Tensor.zeros((1,)))
+    # the ASSIGN UOp has size=1
+    self.assertEqual(assign.uop.size, 1)
+    # the ASSIGN views the buffer with a shrunk st
+    self.assertEqual(assign.uop.src[0].st, ShapeTracker.from_shape((4,)).shrink(((1, 2),)))
+    # the underlying BUFFER has a size=4
+    self.assertEqual(assign.uop.buf_uop.size, 4)
+    # NOTE: output shape is different from the BUFFER shape
+    self.assertNotEqual(assign.uop.shape, a.uop.shape)
+    assign.realize()
+    self.assertEqual(a.tolist(), [1, 0, 1, 1])
+
+  def test_buffer_st(self):
+    a = UOp.new_buffer(Device.DEFAULT, 10, dtypes.float)
+    self.assertEqual(a.st, ShapeTracker.from_shape((10,)))
+
+  def test_ops_st(self):
+    # view / mop
+    a = Tensor.empty(4, 2, 1).permute((1, 2, 0)).uop
+    self.assertEqual(a.st, ShapeTracker.from_shape((4, 2, 1)).permute((1, 2, 0)))
+    # alu / reduce
+    alu = a*2
+    self.assertEqual(alu.st, ShapeTracker.from_shape((2, 1, 4)))
+    r = Tensor.empty(4, 4).sum(axis=1)
+    self.assertEqual(r.uop.st, ShapeTracker.from_shape((4,)))
+
+  def test_st_wmma_none(self):
+    A = UOp(Ops.DEFINE_VAR, dtypes.float.vec(16), arg=('a', UOp.const(dtypes.float, 0), UOp.const(dtypes.float, 1)))
+    B = UOp(Ops.DEFINE_VAR, dtypes.float.vec(16), arg=('b', UOp.const(dtypes.float, 0), UOp.const(dtypes.float, 2)))
+    C = UOp(Ops.DEFINE_VAR, dtypes.float.vec(16), arg=('c', UOp.const(dtypes.float, 0), UOp.const(dtypes.float, 3)))
+    wmma = UOp(Ops.WMMA, dtypes.float.vec(16), (A, B, C))
+    assert wmma.st is None
+
 class TestUOpChildren(unittest.TestCase):
  def test_children_exist(self):
    a = UOp.variable("weird_name_234", 0, 10)
@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor
-from tinygrad.helpers import getenv, GlobalCounters, EMULATE, RANGEIFY
+from tinygrad.helpers import getenv, GlobalCounters, EMULATE
 from tinygrad.engine.realize import lower_schedule_item, ProgramSpec, get_program
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import full_rewrite
@@ -51,11 +51,7 @@ class TestMemoryCount(unittest.TestCase):
    a = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
    b = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024)
    _, mem = get_stats(a+b)
-    if RANGEIFY:
-      # rangeify is smart!
-      self.assertEqual(mem, 1024 + 2*1024)  # 2 lil reads + 1 lil write
-    else:
-      self.assertEqual(mem, 1024*1024 + 2*1024)  # 2 lil reads + 1 write
+    self.assertEqual(mem, 1024*1024 + 2*1024)  # 2 lil reads + 1 write

  def test_self_add(self):
    a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)
@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Tensor, dtypes, TinyJit, UOp
-from tinygrad.helpers import RANGEIFY
 from tinygrad.apps.llm import apply_rope

 # TODO: test_scheduler, but just in uint
@@ -13,7 +12,7 @@ class TestAttention(unittest.TestCase):
    attn = q.scaled_dot_product_attention(k, v)
    sched = attn.schedule()
    # attention has 5 kernels now
-    self.assertEqual(len(sched), 4 if RANGEIFY else 5)
+    self.assertEqual(len(sched), 5)
    softmax_inputs = sched[1:4]
    for si in softmax_inputs:
      assert all(b.dtype == dtypes.half for b in si.bufs), f"non half {si.bufs=}"
@@ -43,4 +42,4 @@ class TestAttention(unittest.TestCase):
    self.assertEqual(prune_size, 1)

 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()
@@ -64,15 +64,6 @@ class TestDevice(unittest.TestCase):
                        shell=True, check=True, env={**os.environ, "DEV": "AMD", "AMD_HIP": "1", "AMD_LLVM": "1"})
    else: self.skipTest("only run on CPU/AMD")

-  def test_compiler_envvar(self):
-    d = Device[Device.DEFAULT]
-    dname = Device.DEFAULT.split(':')[0].upper()
-    assert d._get_compiler_envvar(type("Compiler", (), {})) == f"{dname}_COMPILER"
-    assert d._get_compiler_envvar(type("LLVMCompiler", (), {})) == f"{dname}_LLVM"
-    assert d._get_compiler_envvar(type("RandomCompiler", (), {})) == f"{dname}_RANDOM"
-    assert d._get_compiler_envvar(type(f"{dname}Compiler", (), {})) == f"{dname}_{dname}COMPILER" # do not repeat device name alone
-    assert d._get_compiler_envvar(type(f"{dname}LLVMCompiler", (), {})) == f"{dname}_LLVM" # do not repeat device name
-
 class MockCompiler(Compiler):
  def __init__(self, key): super().__init__(key)
  def compile(self, src) -> bytes: return src.encode()
@@ -101,7 +92,7 @@ class TestCompiler(unittest.TestCase):
 class TestRunAsModule(unittest.TestCase):
  def test_module_runs(self):
    p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE,
-      env={**os.environ, "DEBUG": "1"}, timeout=30,)
+      env={**os.environ, "DEBUG": "1"}, timeout=10,)
    out = (p.stdout + p.stderr).decode()
    self.assertEqual(p.returncode, 0, msg=out)
    self.assertIn("CPU", out) # for sanity check
@@ -307,7 +307,7 @@ class TestDiskTensor(unittest.TestCase):
    ret = t.bitcast(dtypes.uint16).to("CPU") + 1
    assert ret.tolist() == [2827, 3341, 3855, 4369]

-  @unittest.skipIf(OSX or Device.DEFAULT == "CL", "new LLVM has an issue on OSX, CL=1 gives the wrong output")
+  @unittest.skipIf(OSX, "new LLVM has an issue on OSX")
  def test_bf16_disk_write_read(self):
    t = Tensor([10000, -1, -1000, -10000, 20], dtype=dtypes.float32)
    t.to(f"disk:{temp('dt_bf16_disk_write_read_f32')}").realize()
@@ -1,11 +1,12 @@
 import unittest, math, operator, subprocess, struct
 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, truncate, truncate_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, CI, DEBUG
 from hypothesis import given, settings, strategies as strat
 import numpy as np
 import torch
+import ml_dtypes

 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -21,9 +22,7 @@ def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float
  if DEBUG >= 2: print(tensor.numpy())
  try:
    assert tensor.dtype == target_dtype
-    np.testing.assert_allclose(tensor.numpy(), target, rtol={dtypes.float16:1e-3, dtypes.bfloat16:1e-2,
-      dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1}.get(target_dtype, tol_target_dtype))
-
+    np.testing.assert_allclose(tensor.numpy(), target, rtol={dtypes.float16:1e-3, dtypes.bfloat16:1e-2}.get(target_dtype, tol_target_dtype))
  except AssertionError as e:
    raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e

@@ -106,16 +105,16 @@ class TestHelpers(unittest.TestCase):
      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_truncate_fp16(self):
+    self.assertEqual(truncate_fp16(1), 1)
+    self.assertEqual(truncate_fp16(65504), 65504)
+    self.assertEqual(truncate_fp16(65519.999), 65504)
+    self.assertEqual(truncate_fp16(65520), math.inf)
+    self.assertEqual(truncate_fp16(1e-8), 0.0)
+    self.assertEqual(truncate_fp16(-65504), -65504)
+    self.assertEqual(truncate_fp16(-65519.999), -65504)
+    self.assertEqual(truncate_fp16(-65520), -math.inf)
+    self.assertTrue(math.isnan(truncate_fp16(math.nan)))

  def test_float_to_bf16(self):
    # TODO: fuzz this better
@@ -191,7 +190,7 @@ class TestHelpers(unittest.TestCase):
    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())
+    else: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), ml_dtypes.float8_e4m3fn(x))

  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
  def test_truncate_fp8e5m2(self, x):
@@ -199,7 +198,7 @@ class TestHelpers(unittest.TestCase):
    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())
+    else: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), ml_dtypes.float8_e5m2(x))

 class TestTypeSpec(unittest.TestCase):
  def setUp(self):
@@ -578,10 +577,10 @@ class TestAutoCastType(unittest.TestCase):
  def test_gradient_dtype(self):
    old_default_float = dtypes.default_float

-    for default_dtype in dtypes.floats:
+    for default_dtype in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
      if not is_dtype_supported(default_dtype): continue
      dtypes.default_float = default_dtype
-      for dtype in  dtypes.floats:
+      for dtype in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
        if not is_dtype_supported(dtype): continue
        if DEBUG >= 2:
          print(f"testing {default_dtype=}, {dtype=}")
@@ -593,6 +592,14 @@ class TestAutoCastType(unittest.TestCase):

    dtypes.default_float = old_default_float

+  @unittest.skipIf(CI, "TODO: broken RuntimeError: Attempting to relocate against an undefined symbol 'fmaxf'")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
+  def test_backward_sum_acc_dtype(self):
+    # test acc of sum in the backward is upcasted to float
+    t = Tensor([5, -5], dtype=dtypes.half, requires_grad=True)
+    t.reshape(2, 1).expand(2, 10001).max().backward()
+    np.testing.assert_allclose(t.grad.numpy(), [1, 0])
+
  @unittest.skipIf(Device.DEFAULT == "PYTHON", "very slow")
  @unittest.skipIf(CI and Device.DEFAULT == "AMD", "very slow")
  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Binding size is larger than the maximum storage buffer binding size")
@@ -603,7 +610,6 @@ class TestAutoCastType(unittest.TestCase):
    t = Tensor([[x]], dtype=dtypes.half, requires_grad=True).expand(N, N).contiguous()
    np.testing.assert_allclose(t.mean(axis=1).numpy(), np.array([x] * N, dtype=np.float16), rtol=1e-3)

-  @unittest.skip("this test only works with SPLIT_REDUCEOP=1")
  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  def test_mean_half_precision_overflow(self):
    N = 256
@@ -1,7 +1,6 @@
 import unittest
 from tinygrad import Tensor
 from tinygrad.uop import Ops
-from tinygrad.helpers import RANGEIFY

 class TestKernelize(unittest.TestCase):
  def test_add_reshaped(self):
@@ -18,11 +17,7 @@ class TestKernelize(unittest.TestCase):
    a1 = a.sum(axis=1)
    a0 = a1.sum(axis=0)
    a0.kernelize()
-    self.assertEqual(len([s for s in a0.uop.toposort() if s.op is Ops.KERNEL]), 2 if RANGEIFY else 3)
-    self.assertIs(a1.uop.base.op, Ops.REDUCE_AXIS if RANGEIFY else Ops.ASSIGN)
-    # input Tensor and user contiguous kernelize
-    self.assertIs(a0.uop.base.op, Ops.ASSIGN)
-    self.assertIs(a.uop.base.op, Ops.ASSIGN)
+    self.assertIs(a1.uop.base.op, Ops.ASSIGN)

  def test_two_reduce_w_add(self):
    a = Tensor.ones(16,16).contiguous()
@@ -1,26 +1,29 @@
-import unittest, functools
-from tinygrad import Tensor
 import numpy as np
+import unittest
+from tinygrad import Tensor
+from typing import List
+import functools

-def orthogonality_helper(A:Tensor, tolerance=1e-5):
+def orthogonality_helper(A:Tensor,tolerance=1.0e-5):
  b_shape,m = A.shape[0:-2],A.shape[-2]  #outer dimension should be the dim along orthogonality
-  A_identity = (Tensor.eye(m).reshape((1,)*len(b_shape)+(m,m)).expand(b_shape+(m,m)))
+  A_identity = (Tensor.eye(m).reshape((1,) * len(b_shape)+(m,m)).expand(b_shape+(m,m)))
  np.testing.assert_allclose((A @ A.transpose(-2,-1)).numpy(),A_identity.numpy(),atol=tolerance,rtol=tolerance)

-def reconstruction_helper(A:list[Tensor],B:Tensor, tolerance=1e-5):
+def reconstruction_helper(A:List[Tensor],B:Tensor, tolerance=1.0e-5):
  reconstructed_tensor = functools.reduce(Tensor.matmul, A)
  np.testing.assert_allclose(reconstructed_tensor.numpy(),B.numpy(),atol=tolerance,rtol=tolerance)

 class TestLinAlg(unittest.TestCase):
+
  def test_svd_general(self):
    sizes = [(2,2),(5,3),(3,5),(3,4,4),(2,2,2,2,3)]
    for size in sizes:
      a = Tensor.randn(size).realize()
-      U,S,V = a.svd()
+      U,S,V = Tensor.svd(a)
      b_shape,m,n = size[0:-2],size[-2],size[-1]
      k = min(m,n)
      s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
-      s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([None]*len(b_shape) + [(0,m-k), (0,n-k)]))
+      s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([(0,0) for _ in range(len(size)-2)] + [(0,m-k), (0,n-k)]))
      orthogonality_helper(U)
      orthogonality_helper(V)
      reconstruction_helper([U,s_diag,V],a)
@@ -29,7 +32,7 @@ class TestLinAlg(unittest.TestCase):
    sizes = [(2,2),(5,3),(3,5),(2,2,2,2,3)]
    for size in sizes:
      a = Tensor.randn(size).realize()
-      U,S,V = a.svd(full_matrices=False)
+      U,S,V = Tensor.svd(a,full_matrices=False)
      b_shape,m,n = size[0:-2],size[-2],size[-1]
      k = min(m,n)
      s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)).expand(b_shape + (k,k)))
@@ -42,20 +45,20 @@ class TestLinAlg(unittest.TestCase):
  def test_svd_large(self):
    size = (1024,1024)
    a = Tensor.randn(size).realize()
-    U,S,V = a.svd()
+    U,S,V = Tensor.svd(a)
    b_shape,m,n = size[0:-2],size[-2],size[-1]
    k = min(m,n)
    s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
-    s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([None]*len(b_shape) + [(0,m-k), (0,n-k)]))
-    orthogonality_helper(U,tolerance=1e-3)
-    orthogonality_helper(V,tolerance=1e-3)
-    reconstruction_helper([U,s_diag,V],a,tolerance=1e-3)
+    s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([(0,0) for _ in range(len(size)-2)] + [(0,m-k), (0,n-k)]))
+    orthogonality_helper(U,tolerance=1.0e-3)
+    orthogonality_helper(V,tolerance=1.0e-3)
+    reconstruction_helper([U,s_diag,V],a,tolerance=1.0e-3)

  def test_qr_general(self):
    sizes = [(3,3),(3,6),(6,3),(2,2,2,2,2)]
    for size in sizes:
      a = Tensor.randn(size).realize()
-      Q,R = a.qr()
+      Q,R = Tensor.qr(a)
      orthogonality_helper(Q)
      reconstruction_helper([Q,R],a)

@@ -65,9 +68,9 @@ class TestLinAlg(unittest.TestCase):
    for coefs in coefficients:
      for size in sizes:
        a = Tensor.randn(size)
-        b = a.newton_schulz(steps=20, params=coefs, eps=0.0)
+        b = Tensor.newton_schulz(a, steps=20, params=coefs, eps=0.0)
        # ns(A) = U @ Vt -> (U @ Vt) @ (U @ Vt)t = I
-        orthogonality_helper(b if size[-1] > size[-2] else b.transpose(-2, -1), tolerance=1e-3)
+        orthogonality_helper(b if size[-1] > size[-2] else b.transpose(-2, -1), tolerance=1e-1)

 if __name__ == "__main__":
  unittest.main()
@@ -10,11 +10,10 @@ class TestQcom(unittest.TestCase):

    def __validate(imgdt, expected_pitch):
      img = dev.allocator.alloc(imgdt.shape[0] * imgdt.shape[1] * 16, options:=BufferSpec(image=imgdt))
-      pitch = img.texture_info.pitch
+      pitch = (img.descriptor[2] & 0x1fffff80) >> 7
      assert pitch == expected_pitch, f"Failed pitch for image: {imgdt}. Got 0x{pitch:X}, expected 0x{expected_pitch:X}"
      dev.allocator.free(img, imgdt.shape[0] * imgdt.shape[1] * 16, options)

-    # Match opencl pitches for perf
    __validate(dtypes.imageh((1, 201)), 0x680)
    __validate(dtypes.imageh((16, 216)), 0x700)
    __validate(dtypes.imageh((16, 9)), 0x80)
@@ -814,6 +814,20 @@ class TestShapeTrackerSize(unittest.TestCase):
    st = ShapeTracker.from_shape((10,10)).pad(((2,4), (3,1))).flip((True, True))
    self.assertEqual(st.real_size(), 100)

+class TestRender(unittest.TestCase):
+  def test_render(self):
+    st = ShapeTracker.from_shape((2, 3))
+    valid_idx = st.to_valid_uop()
+    idx, valid = valid_idx.get_idx(), valid_idx.get_valid()
+    self.assertEqual(idx.render(), "((ridx0*3)+ridx1)")
+    self.assertEqual(valid.render(), "True")
+
+    st = st.pad(((0, 1), (0, 0)))
+    valid_idx = st.to_valid_uop()
+    idx, valid = valid_idx.get_idx(), valid_idx.get_valid()
+    self.assertEqual(idx.render(), "((ridx0*3)+ridx1)")
+    self.assertEqual(valid.render(), "(ridx0<2)")
+
 class TestVariableShrink(unittest.TestCase):
  def test_shrink(self):
    st = ShapeTracker.from_shape((10,))
@@ -8,13 +8,13 @@ from tinygrad.helpers import Context

 def get_gated_load_uop(valid:UOp, idx:UOp):
  return UOp(Ops.LOAD, dtypes.float, (
-    UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0).index(idx.valid(valid)),
+    UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0).index(idx, valid),
    UOp.const(dtypes.float, 0.0)
  ))

 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
  return UOp(Ops.LOAD, dtypes.float.vec(4), (
-    UOp(Ops.DEFINE_GLOBAL, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.VECTORIZE, dtypes.index.vec(2), idx).valid(valid)),
+    UOp(Ops.DEFINE_GLOBAL, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.VECTORIZE, dtypes.int.vec(2), idx), valid),
    UOp(Ops.VECTORIZE, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4)
  ))

@@ -71,8 +71,8 @@ class TestValidIdxSimplification(unittest.TestCase):
    idx = ridx0+ridx1+ridx2+ridx3
    load = get_gated_load_uop(valid, idx)
    self.check(load,
-      "(((r0+r1)+r2)+r3)",
-      "((((r0*3)+r1)<8)&((((r2*3)+r3)%4)<2))")
+      "(((ridx0+ridx1)+ridx2)+ridx3)",
+      "((((ridx0*3)+ridx1)<8)&((((ridx2*3)+ridx3)%4)<2))")

  def test_simplify_within_valid2(self):
    gidx0 = Special("gidx0", 56)
@@ -85,8 +85,8 @@ class TestValidIdxSimplification(unittest.TestCase):
    ridx0 = Range(0, 2)
    v0 = ridx0<1
    v1 = ((ridx0*5+1)%6)<5
-    self.assertEqual(simplify_valid(v0&v1).render(), "(r0<1)")
-    self.assertEqual(simplify_valid(v1&v0).render(), "(r0<1)")
+    self.assertEqual(simplify_valid(v0&v1).render(), "(ridx0<1)")
+    self.assertEqual(simplify_valid(v1&v0).render(), "(ridx0<1)")

  def test_valid_order_matters2(self):
    gidx0 = Special("gidx0", 13)
@@ -128,8 +128,8 @@ class TestValidIdxSimplification(unittest.TestCase):
    valid = ((((((ridx2*2)+(ridx3*3))+3)%4)<2)!=True) # noqa: E712
    load = get_gated_load_uop(valid, idx)
    self.check(load,
-      "(((r0*2)+(r3*-1))+1)",
-      "(r2<1)")
+      "(((ridx0*2)+(ridx3*-1))+1)",
+      "(ridx2<1)")

  def test_load_in_valid(self):
    # from FUSE_ARANGE=1 python test/test_ops.py TestOps.test_scatter_add
@@ -154,8 +154,8 @@ class TestValidIdxSimplification(unittest.TestCase):
    valid = (ridx2<1)&(ridx1<6)
    load = get_gated_load_uop(valid, idx)
    self.check(load,
-      "(r0*1568)",
-      "((r2<1)&(r1<6))")
+      "(ridx0*1568)",
+      "((ridx2<1)&(ridx1<6))")

  def test_valid_becomes_const1_z3(self):
    from z3 import Ints, Solver, And, If, Not, unsat
@@ -195,7 +195,7 @@ class TestValidIdxSimplification(unittest.TestCase):
    load = get_gated_load_uop(valid, idx)
    self.check(load,
      "1",
-      "((((r0+r1)<1)!=True)&(((r2+r3)<1)!=True))")
+      "((((ridx0+ridx1)<1)!=True)&(((ridx2+ridx3)<1)!=True))")

  def test_valid_with_non_const_rhs(self):
    ridx0 = Range(0, 2**16)
@@ -205,8 +205,8 @@ class TestValidIdxSimplification(unittest.TestCase):
    idx = ridx0%1024
    load = get_gated_load_uop(valid, idx)
    self.check(load,
-      "r0",
-      "(r0<((r1*4)+r2))")
+      "ridx0",
+      "(ridx0<((ridx1*4)+ridx2))")

 class TestImageSimplification(unittest.TestCase):
  def check(self, load, svalid, sidx0, sidx1):
@@ -269,6 +269,7 @@ class TestImageSimplification(unittest.TestCase):
    load = get_load_image_uop(shape, (gidx1<5), (gidx0, gidx1+5))
    self.check(load, None, "gidx0", "(gidx1+5)")

+  @unittest.skip("this should be constructed with an invalid gate")
  def test_valid_empty_set(self):
    gidx0 = Special("gidx0", 32)
    gidx1 = Special("gidx1", 32)
@@ -304,7 +305,7 @@ class TestImageSimplification(unittest.TestCase):
    idx = ((alu4+1530)%1536, alu1+((idx1+((ridx2+7)//8)+31)//32)+(-2))

    load = get_load_image_uop(shape, valid, idx)
-    self.check(load, None, "((((idx1*48)+(r2*6))+r0)+-6)", "(((idx2*2)+r1)+-1)")
+    self.check(load, None, "((((idx1*48)+(ridx2*6))+ridx0)+-6)", "(((idx2*2)+ridx1)+-1)")

  def test_openpilot_conv2(self):
    # conv in test/external/external_test_valid_remove.py
@@ -325,7 +326,7 @@ class TestImageSimplification(unittest.TestCase):
    idx = ((alu3+765)%768, alu1+((idx1+((ridx2+7)//8)+31)//32)+(-2))
    load = get_load_image_uop(shape, valid, idx)

-    self.check(load, None, "((((idx1*24)+(r2*3))+r0)+-3)", "(((idx2*2)+r1)+-1)")
+    self.check(load, None, "((((idx1*24)+(ridx2*3))+ridx0)+-3)", "(((idx2*2)+ridx1)+-1)")

  def test_openpilot_conv3(self):
    # in openpilot 0.9.7
@@ -346,9 +347,9 @@ class TestImageSimplification(unittest.TestCase):
    load = get_load_image_uop(shape, valid, idx)

    self.check(load,
-               "((((idx2*2)+r0)<11)&((((idx1*8)+r1)<3)!=True))",
-               "(((idx0+((idx1*512)+(r1*64)))+832)%1024)",
-               "((((idx2*2)+r0)+(((idx1+((r1+5)//8))+1)//2))+-4)")
+               "((((idx2*2)+ridx0)<11)&((((idx1*8)+ridx1)<3)!=True))",
+               "(((idx0+((idx1*512)+(ridx1*64)))+832)%1024)",
+               "((((idx2*2)+ridx0)+(((idx1+((ridx1+5)//8))+1)//2))+-4)")

  def test_simplify1(self):
    # idx has the form (A % m, A // m + k) and valid has (c0 < A) and (A < c1)
@@ -424,7 +425,7 @@ class TestImageSimplification(unittest.TestCase):
    alu1 = ((idx2*1536)+(ridx4*768)+ridx3+(idx1*24)+(ridx5*3)+-771)//768
    valid = (((idx2+ridx4)<1)!=1)&(((idx1+ridx5)<1)!=1)
    load = get_load_image_uop((128, 768, 4), valid, (alu0, alu1))
-    self.check(load, None, "((((idx1*24)+r3)+(r5*3))+-3)", "(((idx2*2)+r4)+-1)")
+    self.check(load, None, "((((idx1*24)+ridx3)+(ridx5*3))+-3)", "(((idx2*2)+ridx4)+-1)")

 if __name__ == '__main__':
  unittest.main()
@@ -13,6 +13,7 @@ class TestSymbolic(unittest.TestCase):
    assert st.shape == (x, 3)
    assert st.real_strides() == (3, 1)

+  @unittest.expectedFailure
  def test_real_strides_0(self):
    st = ShapeTracker(views=(View(shape=(2, (Variable('start_pos', 1, 8)+1), 1, 1), strides=(8, 1, 0, 0), offset=0, mask=((0, 2), (0, Variable('start_pos', 1, 8)), (0, 1), (0, 1)), contiguous=False), View(shape=(2, (Variable('start_pos', 1, 8)+1)), strides=((Variable('start_pos', 1, 8)+1), 1), offset=0, mask=None, contiguous=True)))   # noqa: E501
    self.assertEqual(st.real_strides(), (8, None))
@@ -197,7 +198,7 @@ class TestSymbolicPad(unittest.TestCase):
  def test_pad(self):
    v = Variable("v", 1, 100).bind(5)
    t = Tensor.ones(100)[:v].pad(((4, 0),))
-    t = t[:9]
+    t = t.reshape(9)
    assert t.tolist() == [0,0,0,0,1,1,1,1,1]


@@ -32,8 +32,9 @@ class TestTensorMutates(unittest.TestCase):
    d.realize()
    is_pattern_uop(d.uop.base, realized_pattern)
    is_pattern_uop(c.uop.base, realized_pattern)
-    is_pattern_uop(c.uop.base, realized_pattern)
-    assert d.uop is not d.uop.base
+    # NOTE: we keep movement ops on top of the buffer view
+    is_pattern_uop(c.uop, UPat(Ops.BUFFER))
+    is_pattern_uop(d.uop, UPat(Ops.VIEW, src=(realized_pattern,)))

  def test_reshape_is_same_child(self):
    a = Tensor([1,2,3])
@@ -55,8 +56,41 @@ class TestTensorUopRepresentation(unittest.TestCase):
    b = Tensor([4.,5,6]).realize()
    c = a+b
    print(c.uop)
-    is_pattern(c, UPat(Ops.ADD))
-    for s in c.uop.src: is_pattern_uop(s.base, realized_pattern)
+    is_pattern(c, UPat(Ops.ADD, src=(realized_pattern, realized_pattern)))
+
+  def test_const_pattern(self):
+    a = Tensor(1)
+    print(a.uop)
+    is_pattern(a, const_pattern) # const in tensor has a DEVICE and VIEW src
+    is_pattern(a, UPat.cvar("x")) # even cvar works!
+
+  def test_consts_do_not_realize(self):
+    a = Tensor(1)
+    print(a.uop)
+    pre_realize = a.uop
+    a.realize()
+    assert a.uop is pre_realize
+
+  def test_viewed_consts_do_not_realize(self):
+    a = Tensor.ones(10, 10)
+    print(a.uop)
+    a.realize()
+    is_pattern(a, const_pattern)
+    self.assertEqual(a.uop.shape, (10, 10))
+
+  # CONST is EXPAND -> RESHAPE -> CONST -> DEVICE
+  def test_consts_dont_have_buffers(self):
+    a = Tensor.ones(10, 10)
+    buffers_in_parents = [x.op for x in a.uop.toposort() if x.op is Ops.BUFFER]
+    self.assertEqual(len(buffers_in_parents), 0)
+    is_pattern(a, UPat(Ops.EXPAND, src=(UPat(Ops.RESHAPE, src=(const_pattern,)),)))
+
+  # COPY has a copyin source and a device.
+  def test_copyin(self):
+    a = Tensor([1.,2,3]).realize()
+    c = a.to("TEST")   # NOTE: this isn't checked
+    print(c.uop)
+    is_pattern(c, UPat(Ops.COPY, src=(realized_pattern, UPat(Ops.DEVICE)), arg=None))

  def test_empty_buf(self):
    a = Tensor.empty(3, 3)
@@ -2,7 +2,7 @@ from __future__ import annotations
 import unittest

 from tinygrad import Tensor
-from tinygrad.helpers import DEBUG, RANGEIFY
+from tinygrad.helpers import DEBUG
 from tinygrad.uop.ops import UOp, Ops, print_uops
 from tinygrad.uop.spec import type_verify, ast_spec, tensor_uop_spec
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -75,7 +75,6 @@ class TestUOpSpec(unittest.TestCase):
    st = UOp.store(buf.view(ShapeTracker.from_shape(())), a.cast(dtypes.float))
    helper_test_verify_ast(st)

-  @unittest.skipIf(RANGEIFY, "RANGEIFY does not push views")
  def test_assert_masked_view_in_const(self):
    t = Tensor(6).uop
    a = t.replace(src=(t.src[0].replace(arg=t.st.reshape((1,)).pad(((0, 1),))),))
@@ -93,37 +93,6 @@ class TestSymbolic(unittest.TestCase):
    assert idx1+idx2 is not idx2
    assert idx1*idx2 is not idx2*idx1

-  def test_uop_gcd_method(self):
-    a = Variable("a", 0, 8)
-    b = Variable("b", 0, 8)
-    self.assertEqual(UOp.gcd(a, a*b, a*3).simplify(), a)
-    self.assertEqual(UOp.gcd(a*a*a, a*b*a, a*3*a).simplify(), a*a)
-    self.assertEqual(UOp.gcd(a*a*10, b*a*5, a*a*5).simplify(), a*5)
-    self.assertEqual(UOp.gcd(a*10, b*5, a*5).simplify(), a.const_like(5))
-    self.assertEqual(UOp.gcd(a, b*5, a*5).simplify(), a.const_like(1))
-
-  def test_divides_exact(self):
-    a = Variable("a", 1, 8)
-    b = Variable("b", 1, 8)
-    self.assertEqual((a*a*3).divide_exact(a).simplify(), a*3)
-    self.assertEqual((a*a*3).divide_exact(a*a*3).simplify(), a.const_like(1))
-    self.assertEqual((a*b*3).divide_exact(a.const_like(3)).simplify(), a*b)
-    self.assertEqual((a*a*3).divide_exact(a*a.const_like(-3)).simplify(), a*-1)
-    self.assertEqual((a*a*b*3).divide_exact(a*b).simplify(), a*3)
-    self.assertEqual((a*3+a*b).divide_exact(a).simplify(), b+3)
-    self.assertEqual((a*b*3+a*b*b).divide_exact(a*b).simplify(), b+3)
-    self.assertEqual((((a*-2)+14)*b).divide_exact(((a*-2)+14)).simplify(), b)
-
-  def test_divide_exact_not(self):
-    a = Variable("a", 1, 8)
-    b = Variable("b", 1, 8)
-    x = Variable("x", -20, 0)
-    self.assertEqual((a).divide_exact(b), None)
-    self.assertEqual((a+2).divide_exact(a), None)
-    self.assertEqual((x*-1).divide_exact(a), None)
-    self.assertEqual((a*5).divide_exact(a*10), None)
-    self.assertEqual((a*10-1).divide_exact(a*10), None)
-
  def test_factorize(self):
    a = Variable("a", 0, 8)
    b = Variable("b", 0, 8)
@@ -141,7 +110,7 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(-Variable("a", 0, 8), -8, 0, "(a*-1)")

  def test_xor_0(self):
-    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) ^ 0, 0, 8, "a", test_z3=False)
+    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) ^ 0, 0, 8, "a")

  def test_add_1(self):
    self.helper_test_variable(Variable("a", 0, 8)+1, 1, 9, "(a+1)")
@@ -246,7 +215,7 @@ class TestSymbolic(unittest.TestCase):
  def test_range_mod_its_symbolic_bound(self):
    a = Variable("a", 1, 10, dtypes.index)
    ridx = UOp.range(a+2, 0)
-    self.helper_test_variable(ridx%(a+2), 0, 11, "r0")
+    self.helper_test_variable(ridx%(a+2), 0, 11, "ridx0")

  def test_div_min_max(self):
    self.helper_test_variable(Variable("a", 2, 7) // 2, 1, 3, "(a//2)")
@@ -481,33 +450,6 @@ class TestSymbolic(unittest.TestCase):
  def test_mul_div_factor_div_neg(self):
    self.helper_test_variable((Variable("a", 0, 10)*-4+4)//8, -4, 0, "(((a*-1)+1)//2)")

-  def test_div_symbolic_const_gcd(self):
-    a = Variable("a", -10, 10)
-    b = Variable("b", -10, 10)
-    d = Variable("d", 1, 10)
-    self.helper_test_variable((3*a+9*b)//(3*d), -40, 40, "((a+(b*3))//d)")
-
-  def test_symbolic_gcd_div(self):
-    a = Variable("a", -10, 10)
-    b = Variable("b", -10, 10)
-    c = Variable("c", -10, 10)
-    d1 = Variable("d1", 1, 10)
-    d2 = Variable("d2", -10, -1)
-    self.helper_test_variable((d1*a*b*d1)//(d1), -1000, 1000, "(a*(b*d1))")
-    self.helper_test_variable((d1*a*d2*b*d1)//(d1*d2),  -1000, 1000, "(a*(b*d1))")
-    self.helper_test_variable((d1*a + b*d1)//(d1), -20, 20, "(a+b)")
-    self.helper_test_variable((d1*a + b*d1 + c*d1)//(d1), -30, 30, "(c+(a+b))")
-    self.helper_test_variable((3*a*d1 + 9*b*d1)//(3*d1*d2), -40, 40, "(((a+(b*3))//(d2*-1))*-1)")
-    self.helper_test_variable((3*a*d1 + 9*b*d1+3)//(3*d1*d2), -401, 399, "(((((a*d1)+((b*d1)*3))+1)//((d1*d2)*-1))*-1)")
-
-  def test_symbolic_factor_remainder_div(self):
-    a = Variable("a", 0, 10)
-    b = Variable("b", 0, 10)
-    d = Variable("d", 1, 10)
-    self.helper_test_variable((d*a+b)//d, 0, 20, "(a+(b//d))")
-    self.helper_test_variable((d*a*20+b)//(5*d), 0, 42, "((a*4)+(b//(d*5)))")
-    self.helper_test_variable((d*a*20+b*d*5+10)//(5*d), 0, 52, "((b+(a*4))+(2//d))")
-
  def test_mod_gcd_factor_neg(self):
    self.helper_test_variable((Variable("a", 0, 10)*-4+4)%8, -4, 4, "((((a*-1)+1)%2)*4)")

@@ -578,13 +520,6 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//12, 0, 4, ("(((lidx2//2)+gidx0)//3)", "((gidx0+(lidx2//2))//3)"))
    self.helper_test_variable((lidx2*2+gidx0*4+lidx3)//12, 0, 4, ("(((lidx2//2)+gidx0)//3)", "((gidx0+(lidx2//2))//3)"))

-  @unittest.expectedFailure  # TODO: improve nest_div_by_smallest_factor
-  def test_sum_div_complex4(self):
-    gidx0 = Variable("gidx0", 0, 2)
-    lidx2 = Variable("lidx2", 0, 12)
-    lidx3 = Variable("lidx3", 0, 12)
-    self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, ("((lidx2+(lidx3*2))//3)"))
-
  def test_sum_mul_distribute(self):
    gidx0 = Variable("gidx0", 0, 7)
    lidx2 = Variable("lidx2", 0, 12)
@@ -656,6 +591,45 @@ class TestSymbolic(unittest.TestCase):
    with self.assertRaises(AssertionError):
      self.helper_test_variable((30 * b + 1) % 18 + ((30 * b + 1) // 18) * 18, 1, 3001, "((b*30)+1)")

+  def test_arange_unrolled4(self):
+    gidx = Variable("gidx", 0, 2559)
+    unrolled_div = (gidx+2561)//4+(gidx+2562)//4+(gidx+2560)//4+(gidx+2559)//4
+    self.helper_test_variable(unrolled_div, 2559, 5118, "(gidx+2559)")
+
+  def test_arange_unrolled4_with_cast(self):
+    gidx = Variable("gidx", 0, 2559, dtypes.index)
+    dt = dtypes.int
+    unrolled_div = ((gidx+2561)//4 + 2).cast(dt)+((gidx+2562)//4).cast(dt)+((gidx+2560)//4).cast(dt)+((gidx+2559)//4).cast(dt)
+    self.helper_test_variable(unrolled_div, 2561, 5120, "((int)(gidx)+2561)")
+
+  def test_arange_unrolled4_mul(self):
+    gidx = Variable("gidx", 0, 2559)
+    unrolled_div = 2*((gidx+2561)//4)+2*((gidx+2562)//4)+2*((gidx+2560)//4)+2*((gidx+2559)//4)
+    self.helper_test_variable(unrolled_div, 5118, 10236, "((gidx*2)+5118)")
+
+  def test_arange_unrolled4_small(self):
+    gidx = Variable("gidx", 0, 3)
+    unrolled_div = (gidx)//4+(gidx+2)//4+(gidx+3)//4+(gidx+1)//4
+    self.helper_test_variable(unrolled_div, 0, 3, "gidx")
+
+    gidx = Variable("gidx", 0, 2)
+    unrolled_div = (gidx)//4+(gidx+2)//4+(gidx+3)//4+(gidx+1)//4
+    self.helper_test_variable(unrolled_div, 0, 2, "gidx")
+
+    gidx = Variable("gidx", 0, 1)
+    unrolled_div = (gidx)//4+(gidx+2)//4+(gidx+3)//4+(gidx+1)//4
+    self.helper_test_variable(unrolled_div, 0, 1, "gidx")
+
+  def test_arange_unrolled2(self):
+    gidx = Variable("gidx", 0, 2559)
+    unrolled_div = (gidx+2559)//2+(gidx+2560)//2+3
+    self.helper_test_variable(unrolled_div, 2562, 5121, "(gidx+2562)")
+
+  def test_arange_unrolled2_neg(self):
+    ridx = Variable("ridx", 0, 255)
+    unrolled_div = -((255-ridx)//2) - ((256-ridx)//2)
+    self.helper_test_variable(unrolled_div, -255, 0, "(ridx+-255)")
+
  def test_gated_load(self):
    idx = Variable("idx", 0, 24)
    self.helper_test_variable(idx//4, 0, 6, "(idx//4)")
@@ -408,7 +408,7 @@ class TestVizProfiler(unittest.TestCase):
      get_profile(prof)

  def test_python_marker(self):
-    with Context(VIZ=1):
+    with Context(PROFILE=1):
      a = Tensor.empty(1, device="NULL")
      b = Tensor.empty(1, device="NULL")
      (a+b).realize()
@@ -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, RANGEIFY
+from tinygrad.helpers import CI, Profiling, WINO

@unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
 class TestWinogradClose(unittest.TestCase):
@@ -35,35 +35,32 @@ class TestWinograd(unittest.TestCase):
  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 if RANGEIFY else 4)
+    self.assertEqual(len(out.schedule()), 4)

  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), 3 if RANGEIFY else 9)
+    self.assertEqual(len(backward_schedule), 9)

  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()
+    Tensor.conv2d(x,w).realize()
    ops_wino, mem_wino = GlobalCounters.global_ops, GlobalCounters.global_mem
+    WINO.value = 0
    GlobalCounters.reset()
-    with Context(WINO=0):
-      Tensor.conv2d(x,w).realize()
+    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}")
-
-    if not RANGEIFY:
-      self.assertLess(ops_ratio, 2.6)  # TODO: there's issues with factorization now
-      self.assertLess(mem_ratio, 10)
+    self.assertLess(ops_ratio, 2.6)  # TODO: there's issues with factorization now
+    self.assertLess(mem_ratio, 10)

  def test_dtype(self):
    IC, OC, X, Y = 4,4,9,9
@@ -118,7 +118,7 @@ class TransformerBlock:
    return h + self.ffn_down(gated)

  def __call__(self, x: Tensor, start_pos: int|UOp):
-    return self._feed_forward(self._attention(x, start_pos)).contiguous()
+    return self._feed_forward(self._attention(x, start_pos))

 class Transformer:
  def __init__(self, *, num_blocks, dim, hidden_dim, n_heads, n_kv_heads, norm_eps, vocab_size, max_context):
@@ -156,8 +156,6 @@ class Transformer:
                        n_heads=kv[f'{arch}.attention.head_count'], n_kv_heads=kv[f'{arch}.attention.head_count_kv'],
                        norm_eps=kv[f'{arch}.attention.layer_norm_rms_epsilon'], vocab_size=len(kv['tokenizer.ggml.tokens']), max_context=max_context)
    nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused
-    # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
-    for s in nn.state.get_parameters(model): s.replace(s.contiguous())
    return model, kv

  def generate(self, tokens:list[int], start_pos=0):
@@ -19,7 +19,7 @@ from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, blo
 from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
 from tinygrad.codegen.opt.postrange import pm_postrange_opt
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_reduce_simplify, pm_flatten_range
-from tinygrad.schedule.rangeify import pm_add_buffers, rangeify_codegen
+from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen

@dataclass
 class RewriteStep:
@@ -76,7 +76,7 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
  ret.append(RewriteStep(sym+pm_pre_expander+expander, name="expander"))

  # add locals
-  ret.append(RewriteStep(pm_add_buffers+rangeify_codegen, name="add local buffers"))
+  ret.append(RewriteStep(pm_add_buffers_local+rangeify_codegen, name="add local buffers"))

  # ** devectorizer (full_graph_rewrite) **
  # remove reduce
@@ -12,7 +12,7 @@ from tinygrad.renderer import Renderer

 def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
  if (idx:=uop_given_valid(valid, start_idx)) is None: return buf.index(UOp.invalid())
-  if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx.valid(valid))
+  if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx, valid)

  # wait for it to be image indexed before running simplification
  if start_idx.dtype.count != 2: return None
@@ -43,7 +43,7 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:

  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
-  return buf.index(idx.valid(new_valid) if new_valid is not None else idx)
+  return buf.index(idx, new_valid)

 def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp, cast:UOp|None=None) -> UOp|None:
  if store_gate not in [gate.src[0] for gate in val.toposort() if gate.op is Ops.IF]: return None
@@ -52,11 +52,14 @@ def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp,

 load_store_indexing = PatternMatcher([
  # image load valid idx simplification
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
-  # simplify away long after index has been lowered
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("x", dtypes.long), UPat.var("c", dtypes.bool))), lambda buf,x,c: simplify_valid_load(buf, x, c)),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("start_idx"), UPat.var("valid"))), simplify_valid_load),
+  # lower turn the invalid into a gate, must come before index dtype lowering
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate,),), lambda buf,x,cond,i: buf.index(x, cond)),
  # drop true gate
  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("x"), UPat.const(dtypes.bool, True)),), lambda buf,x: buf.index(x)),
+  # remove hanging cast
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.int).cast()),), lambda buf,idx: buf.index(idx)),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.int).cast(), UPat.var("valid"))), lambda buf,idx,valid: buf.index(idx, valid)),
  # delete_redundant_gates (after expand)
  (UPat(Ops.STORE, src=(UPat.any(stidx:=UPat.var("buf").index(UPat.var("idx"), UPat.var("store_gate")), stidx.cast().named("cast")),
                                  UPat.var("val")), name="store", allow_any_len=True), delete_redundant_gates),
@@ -64,21 +67,21 @@ load_store_indexing = PatternMatcher([

 # ***** load/store grouping *****

-def expand_index(buf:UOp, vec:UOp):
-  if getenv("UNSAFE_DISABLE_MASK", 0): vec = vec.get_idx()
+def expand_index(buf:UOp, vec:UOp, mask:UOp|None=None):
+  if getenv("UNSAFE_DISABLE_MASK", 0): mask = None
  # generate the individual indexes
-  midx = graph_rewrite(UOp.sink(*[buf.index(vec.gep(i)) for i in range(vec.dtype.count)]),
+  midx = graph_rewrite(UOp.sink(*[buf.index(vec.gep(i), mask.gep(i) if mask is not None else None) for i in range(vec.dtype.count)]),
                       symbolic_flat+load_store_indexing, name=f"index_buf_{buf.arg}")
  # extract all the relevant offsets
  offsets_rootsrc: defaultdict[Any, dict[int, list[int]]] = defaultdict(dict)
  for i in range(vec.dtype.count):
-    idx: Any = midx.src[i].src[1].get_idx()
+    idx: Any = midx.src[i].src[1]
    if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].arg
    elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].arg
    elif idx.op is Ops.CONST and idx.arg is Invalid: root_src, arg = "INVALID", 0
    elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.arg
    else: root_src, arg = idx, 0
-    root_src = (midx.src[i].src[1].get_valid(), root_src)
+    if len(midx.src[i].src) == 3: root_src = (midx.src[i].src[2], root_src)
    offsets_rootsrc[root_src].setdefault(arg, []).append(i)

  # then rewrite everything we can into groups
@@ -99,7 +102,7 @@ def expand_index(buf:UOp, vec:UOp):
      global_offset += len(grp)
  assert None not in idxs, f"some idxs are missing {idxs}"
  # this base thing is for image, we want the CAT to be a normal pointer
-  post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(global_offset), tuple(ret))
+  post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(vec.dtype.count), tuple(ret))
  return post_cat.gep(tuple(cast(list[int], idxs)))

 def cat_after_store(cat:UOp, data:UOp, sto:UOp):
@@ -121,6 +124,8 @@ def gep_on_store(gep:UOp, st:UOp, sto:UOp):

 load_store_folding = PatternMatcher([
  (UPat(Ops.INDEX, src=(UPat(Ops.VECTORIZE, src=UPat(GroupOp.Defines, name="buf")), UPat.var("vec"))), expand_index),
+  (UPat(Ops.INDEX, src=(UPat(Ops.VECTORIZE, src=UPat(GroupOp.Defines, name="buf")), UPat.var("vec"),
+                        UPat.var("mask"))), expand_index),
  # GEP after LOAD
  (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True),
   lambda gep, ld: ld.replace(dtype=ld.dtype.scalar().vec(gep.dtype.count), src=(gep.src[0],)+ld.src[1:]).gep(gep.arg)),
@@ -128,7 +133,7 @@ load_store_folding = PatternMatcher([
  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st")), allow_any_len=True, name="sto"), gep_on_store),
  # put PTRCAT after LOAD
  (UPat(Ops.LOAD, src=(UPat(Ops.PTRCAT, name="cat"),), name="ld", allow_any_len=True),
-   lambda cat,ld: UOp(Ops.CAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
+   lambda cat,ld: UOp(Ops.CAT, ld.dtype, tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
  # put PTRCAT after STORE
  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), allow_any_len=True, name="sto"), cat_after_store),
 ])
@@ -160,8 +165,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
  lengths.append(1)  # worst case, it's not folded

  # filter fold lengths that don't divide
-  offset, mask = idx.src[1].get_idx(), idx.src[1].get_valid()
-  if must_divide: lengths = [x for x in lengths if offset.divides(x) is not None]
+  if must_divide: lengths = [x for x in lengths if idx.src[1].divides(x) is not None]

  # split based on the fold lengths
  global_offset = 0
@@ -170,7 +174,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
    # with 1 at the end of the lengths list, this will always hit
    for fold_length in lengths:
      if global_offset+fold_length > sz: continue
-      lidx = buf.index((offset + global_offset).valid(mask))
+      lidx = buf.index(idx.src[1] + global_offset, idx.src[2] if len(idx.src) > 2 else None)
      if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
      if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length))))+ls.src[2:]))
      else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar().vec(fold_length)))
@@ -186,20 +190,19 @@ def image_fixup(ls:UOp):
  if ls.src[0].op is Ops.CAST and isinstance(image_dtype:=ls.src[0].src[0].dtype, ImageDType):
    assert ls.src[0].dtype.count == 4, "image must be casted to 4"
    idx = ls.src[0].src[0]
-    x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
-    oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), ((x // 4) % image_dtype.shape[1], (x // (4*image_dtype.shape[1]))))
-    idx = idx.replace(src=(idx.src[0], oidx.valid(valid)))
+    oidx = UOp(Ops.VECTORIZE, dtypes.int.vec(2), ((idx.src[1] // 4) % image_dtype.shape[1], (idx.src[1] // (4*image_dtype.shape[1]))))
+    idx = idx.replace(src=(idx.src[0], oidx)+idx.src[2:])
    return ls.replace(src=(idx,)+ls.src[1:])

  # this is an unprocessed image without a cast, aka unfoldable image load. this doesn't work for stores
-  if isinstance(image_dtype:=ls.src[0].dtype, ImageDType) and ls.src[0].src[1].get_idx().dtype != dtypes.index.vec(2):
+  if isinstance(image_dtype:=ls.src[0].dtype, ImageDType) and ls.src[0].src[1].dtype != dtypes.int.vec(2):
    assert ls.op is Ops.LOAD, "if an image store isn't upcasted to 4, we can't store it"
    idx = ls.src[0]
-    x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
-    oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), ((x // 4) % image_dtype.shape[1], (x // (4*image_dtype.shape[1]))))
-    idx = idx.replace(src=(idx.src[0], oidx.valid(valid)))
+    id4 = idx.src[1] % 4
+    oidx = UOp(Ops.VECTORIZE, dtypes.int.vec(2), ((idx.src[1] // 4) % image_dtype.shape[1], (idx.src[1] // (4*image_dtype.shape[1]))))
+    idx = idx.replace(src=(idx.src[0], oidx)+idx.src[2:])
    vec_load = ls.replace(dtype=ls.dtype.vec(4), src=(idx,)+ls.src[1:])
-    return functools.reduce(lambda ret, i: (x % 4).ne(i).where(ret, vec_load.gep(i)), range(4), ls.const_like(float('nan')))
+    return functools.reduce(lambda ret, i: id4.ne(i).where(ret, vec_load.gep(i)), range(4), ls.const_like(float('nan')))

  return None

@@ -226,7 +229,6 @@ def no_vectorized_wmma(wmma:UOp):

 def no_vectorized_alu(alu:UOp):
  if alu.dtype.vcount == 1: return None
-  if alu.op is Ops.WHERE and alu.src[2].arg is Invalid: return None  # image load/store has cond.where(idx.vec(2), Invalid) as the index
  alus = tuple(UOp(alu.op, alu.dtype.scalar(), tuple(s.gep(i) for s in alu.src), alu.arg) for i in range(alu.dtype.vcount))
  return UOp(Ops.VECTORIZE, alu.dtype, alus)

@@ -236,7 +238,7 @@ def no_vectorized_buf(buf:UOp):
 def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp):
  cnt = cast.dtype.count
  assert idx.dtype.count == 1, f"idx dtype must be 1 {idx.dtype}"
-  return buf.broadcast(cnt).index(idx.broadcast(cnt)*cnt+UOp.const(dtypes.index.vec(cnt), tuple(range(cnt))))
+  return buf.broadcast(cnt).index(idx.broadcast(cnt)*cnt+UOp.const(dtypes.int.vec(cnt), tuple(range(cnt))))

 devectorize = PatternMatcher([
  # no ALU on vectorized dtypes
@@ -256,8 +258,7 @@ pm_render = PatternMatcher([
  (UPat(Ops.VECTORIZE, src=(UPat(name='x'),)), lambda x: x),
  # give any loads that are masked an alt value
  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), allow_any_len=True, name="x"),
-    lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:])
-      if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None),
+    lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:]) if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE) else None),
  # gate any stores that aren't gated with ifs
  (UPat(Ops.STORE, src=(UPat(src=(UPat(), UPat(), UPat(dtype=dtypes.bool)), name="idx").or_casted(), UPat()), name="store", allow_any_len=True),
    lambda store,idx: UOp(Ops.STORE, dtype=store.dtype, src=store.src[:2]+(UOp(Ops.IF, src=(idx.src[2],)),)+store.src[2:]) if \
@@ -2,7 +2,7 @@
 import functools, itertools, operator
 from tinygrad.dtype import dtypes, PtrDType, AddrSpace
 from tinygrad.helpers import AMX, dedup, flatten, all_same, prod, partition
-from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, AxisType, range_start
+from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, AxisType
 from tinygrad.schedule.rangeify import BufferizeOpts

 def _expand_arg_to_idx(args:tuple[tuple[int, int], ...], rpk:dict[int, int]) -> int:
@@ -50,7 +50,7 @@ def do_expand(root:UOp):
      if root.op is Ops.IF or src.op is Ops.IF:
        # for the first arg of IF, just pass them through ignoring UNROLLS
        new_srcs.append(src)
-      elif root.op in range_start and i >= range_start[root.op]:
+      elif (root.op is Ops.STORE and i >= 2) or (root.op in {Ops.REDUCE, Ops.BUFFERIZE} and i >= 1) or (root.op is Ops.WMMA and i >= 3):
        # for any range args of STORE/REDUCE, pass them through
        new_srcs.append(src)
      elif root.op is Ops.INDEX and i >= 1 and not isinstance(root.dtype, PtrDType):
@@ -222,8 +222,6 @@ def remove_blockend(x:UOp):
      if late_ops[i].op is Ops.BARRIER and late_ops[i+1].op is Ops.BARRIER: late_ops[i+1] = UOp(Ops.NOOP)
    arg = BasicBlock(parent_block.arg.lst+tuple(late_ops), tuple([y for y in x.arg.ctx if y is not x.arg.end]), cnt=x.arg.cnt)
    return UOp(Ops.BLOCK, src=tuple(y for y in x.src if y is not parent_block)+parent_block.src, arg=arg)
-  # else the whole context ended by the blockend is already in this block and we can safely turn it into a block
-  return UOp(Ops.BLOCK, src=x.src, arg=BasicBlock(x.arg.lst, tuple([y for y in x.arg.ctx if y is not x.arg.end]), cnt=x.arg.cnt))

 block_merge = PatternMatcher([
  (UPat((Ops.BLOCK, Ops.BLOCKEND), name="x"), merge_block),
@@ -48,7 +48,32 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
  # make a copy so it does not mutate the input
  k = k.copy()

-  # upcast float4 images, this must be early so we don't accidentally add locals before the upcast
+  # should use matvec - TODO: adjust/tune based on the wide vs tall/large vs small mat
+  MV_BLOCKSIZE, MV_THREADS_PER_ROW, MV_ROWS_PER_THREAD = getenv("MV_BLOCKSIZE", 4), getenv("MV_THREADS_PER_ROW", 8), getenv("MV_ROWS_PER_THREAD", 4)
+  if k.opts.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and  \
+    k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.opts.has_shared and \
+    (mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.LOAD and mulop.src[1].op is Ops.LOAD:
+    idx0, idx1 = mulop.src[0].src[0].src[1].get_idx(), mulop.src[1].src[0].src[1].get_idx()
+    first_reduce_rng = k.ranges_of(AxisType.REDUCE)[0]
+    if any(u is first_reduce_rng for u in idx0.split_uop(Ops.ADD)) and all(r in idx1.ranges for r in idx0.ranges):
+      for global_idx in k.axes_of(AxisType.GLOBAL):
+        if first_reduce_rng.src[0].divides(MV_THREADS_PER_ROW) is not None and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
+          if DEBUG >= 3:
+            print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
+          if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
+          if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
+          if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
+          return k
+
+  # are we grouping? (requires local shape support)
+  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False):
+    for sz in [16]:
+      try:
+        k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz))
+        break
+      except KernelOptError: pass
+
+  # upcast float4 images
  for buf_index,buf in enumerate(k.bufs):
    if isinstance(buf.src[0].dtype, ImageDType):
      # part of real_strides
@@ -60,32 +85,6 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
        elif axis in k.unrollable_dims:
          k.apply_opt(Opt(OptOps.UNROLL, k.unrollable_dims.index(axis), 4))

-  # should use matvec - TODO: adjust/tune based on the wide vs tall/large vs small mat
-  MV_BLOCKSIZE, MV_THREADS_PER_ROW, MV_ROWS_PER_THREAD = getenv("MV_BLOCKSIZE", 4), getenv("MV_THREADS_PER_ROW", 8), getenv("MV_ROWS_PER_THREAD", 4)
-  if k.opts.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and  \
-    k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.opts.has_shared and \
-    (mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.LOAD and mulop.src[1].op is Ops.LOAD:
-    idx0, idx1 = mulop.src[0].src[0].src[1].get_idx(), mulop.src[1].src[0].src[1].get_idx()
-    if k.ranges_of(AxisType.REDUCE):
-      first_reduce_rng = k.ranges_of(AxisType.REDUCE)[0]
-      if any(u is first_reduce_rng for u in idx0.split_uop(Ops.ADD)) and all(r in idx1.ranges for r in idx0.ranges):
-        for global_idx in k.axes_of(AxisType.GLOBAL):
-          if first_reduce_rng.src[0].divides(MV_THREADS_PER_ROW) is not None and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
-            if DEBUG >= 3:
-              print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
-            if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
-            if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
-            if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
-            return k
-
-  # are we grouping? (requires local shape support)
-  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False):
-    for sz in [16]:
-      try:
-        k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz))
-        break
-      except KernelOptError: pass
-
  # no more opt if we are grouping
  if k.group_for_reduces: return k

@@ -2,15 +2,14 @@ from __future__ import annotations
 import math, itertools
 from collections import defaultdict
 from typing import cast, Final
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, can_pad, GroupOp
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, can_pad
 from tinygrad.device import Buffer
 from tinygrad.dtype import AddrSpace, dtypes, ImageDType
-from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts
+from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod
 from tinygrad.codegen.opt import axis_colors, Opt, OptOps, KernelOptError, check, axis_letters
 from tinygrad.codegen.simplify import pm_flatten_range
 from tinygrad.renderer import Renderer
-
-remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
+from tinygrad.schedule.rangeify import remove_tags

 # NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
 axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
@@ -71,13 +70,6 @@ class Scheduler:
                        or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
    for ls in local_store_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])

-    # filter any not in reduces
-    # TODO: enable this
-    """
-    reduce_rngs = [x.ranges for x in self.ast.toposort() if x.op is Ops.REDUCE]
-    for ls in reduce_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])
-    """
-
    return [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE and x.arg[1] == AxisType.LOOP] if store_rngs else []

  def convert_loop_to_global(self):
@@ -147,11 +139,6 @@ class Scheduler:
        upcast_local_sz = prod([self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)])
        smem_sz = amt*upcast_local_sz*self.reduceop.dtype.itemsize
        check(smem_sz <= self.opts.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.opts.shared_max}")
-      if self.reduceop is not None and (opt.op in {OptOps.GROUP, OptOps.GROUPTOP}):
-        # We currently dont support a group within another rudece, TODO: fix if-contexts
-        reduce = [u for u in self.ast.parents if u.op is Ops.REDUCE and rng in merge_dicts([r.ranges for r in u.src[1:]])][0]
-        check(not any(u.arg[-1] in (AxisType.REDUCE, AxisType.UNROLL, AxisType.GROUP_REDUCE) for u in reduce.ranges),
-          "cannot have a GROUP_REDUCE inside another reduce")

      if opt.op is OptOps.UNROLL:
        check(amt <= 32, "don't unroll more than 32")
@@ -328,12 +315,12 @@ def apply_opts(ctx:Renderer, ast:UOp):
  if ast.tag is not None: return None
  k = Scheduler(ast, ctx)
  k.convert_loop_to_global()
-  if ast.arg is not None and ast.arg.opts_to_apply is not None:
-    for opt in ast.arg.opts_to_apply: k.apply_opt(opt)
-  elif BEAM >= 1:
+  if BEAM >= 1:
    from tinygrad.codegen.opt.search import beam_search
    rawbufs = bufs_from_ast(ast, ctx.device)
    k = beam_search(k, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
+  elif ast.arg is not None and ast.arg.opts_to_apply is not None:
+    for opt in ast.arg.opts_to_apply: k.apply_opt(opt)
  elif not NOOPT and (ast.arg is None or ast.arg.applied_opts == ()):
    from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
    # NOTE: hand_coded_optimizations doesn't support multiblock opts yet
@@ -1,10 +1,10 @@
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start
-from tinygrad.uop.symbolic import symbolic_flat, sym, invalid_pat
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute
+from tinygrad.uop.symbolic import symbolic_flat, sym
 from tinygrad.helpers import partition
 from tinygrad.dtype import dtypes

 def flatten_range(r:UOp):
-  off = range_start[r.op]
+  off = 2 if r.op is Ops.STORE else 1
  rngs = r.src[off:]
  if not len(rngs): return None
  new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
@@ -17,24 +17,20 @@ pm_flatten_range = PatternMatcher([

 def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
 def simplify_merge_adjacent(u:UOp) -> UOp|None:
-  reduce_ranges = [x.ranges for x in u.sparents if x.op is Ops.REDUCE]
-  i = range_start[u.op]
+  i = 2 if u.op is Ops.STORE else 1
  while i < len(u.src)-1:
    r0, r1 = u.src[i], u.src[i+1]
    # check same type
    if r0.arg[-1] == r1.arg[-1]:
-      # check if the ranges to merge are in the same reduces
-      if all((r0 in rngs) == (r1 in rngs) for rngs in reduce_ranges):
-        s0, s1 = r0.src[0], r1.src[0]
-        # do the merge
-        new_range = r0.replace(src=(s0*s1,))
-        nidx = graph_rewrite(u, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
-                             name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
-
-        # check if it simplifies
-        if count_divmod(nidx) <= count_divmod(u):
-          u = nidx
-          continue
+      s0, s1 = r0.src[0], r1.src[0]
+      # do the merge
+      new_range = r0.replace(src=(s0*s1,))
+      nidx = graph_rewrite(u, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
+                           name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
+      # check if it simplifies
+      if count_divmod(nidx) <= count_divmod(u):
+        u = nidx
+        continue
    i += 1
  return u

@@ -44,19 +40,19 @@ pm_simplify_ranges = PatternMatcher([

 # **** reduce simplification ****

-def no_range(u:UOp) -> bool: return not any(x.op is Ops.RANGE for x in u.sparents)
-
 def reduce_rangeless(red:UOp):
  # TODO: share code with reduce_unparented
  if red.arg not in {Ops.ADD, Ops.MAX}: return None
  if red.src[0].dtype != red.dtype: return None
-  if not no_range(red.src[0]): return None
+  if any(x.op in {Ops.RANGE} for x in red.src[0].toposort()): return None
  ret = red.src[0]
  if red.arg is Ops.ADD:
    for r in red.src[1:]:
      ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
  return ret

+def no_range(u:UOp) -> bool: return not any(x.op is Ops.RANGE for x in u.sparents)
+
 pm_reduce_collapse = PatternMatcher([
  # lift x+y out of reduce on lt
  ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
@@ -78,12 +74,12 @@ pm_reduce_collapse = PatternMatcher([
   lambda x,gate,b=None: gate.broadcast(x.dtype.count).where(x, 0) if b is not None else gate.where(x, 0)),
  # WHERE on LOAD (works on max too)
  (UPat.var("gate").where(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).load(), 0).reduce(arg=Ops.ADD, allow_any_len=True),
-   lambda buf,idx,gate: buf.index(idx.valid(gate)).load()),
+   lambda buf,idx,gate: buf.index(idx, gate).load()),
  (UPat.var("gate").where(0, UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).load()).reduce(arg=Ops.ADD, allow_any_len=True),
-   lambda buf,idx,gate: buf.index(idx.valid(gate.logical_not())).load()),
+   lambda buf,idx,gate: buf.index(idx, gate.logical_not()).load()),
  # INDEX on RANGE / gated RANGE
-  (UPat.var("buf").index(UPat.var("idx").eq(UPat(Ops.RANGE, name="r").or_casted()).where(UPat.var("expr"), invalid_pat)),
-   lambda buf,r,idx,expr,i: buf.index(expr.substitute({r:idx.cast(r.dtype)}).valid((idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])))),
+  (UPat.var("buf").index(UPat.var("expr"), UPat.var("idx").eq(UPat(Ops.RANGE, name="r").or_casted())),
+   lambda buf,r,idx,expr: buf.index(expr.substitute({r:idx.cast(r.dtype)}), (idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0]))),
  # AND on WHERE
  ((UPat.any(UPat(Ops.DEFINE_VAR, name="x"), UPat(Ops.DEFINE_VAR).gep(name="x")) & UPat.var("y")) \
   .where(UPat.cvar("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
@@ -102,17 +98,16 @@ def reduce_collapse(red:UOp):
        replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
  collapse_fxn = red.substitute(replaces)
  sink = graph_rewrite(collapse_fxn, pm_reduce_collapse, name="reduce_collapse")
-  return sink.substitute({v:k for k,v in replaces.items()}) if no_range(sink) else None
+  if any(x.op is Ops.RANGE for x in sink.toposort()): return None
+  return sink.substitute({v:k for k,v in replaces.items()})

 def reduce_unparented(red:UOp):
-  if red.arg not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None
+  if red.arg not in {Ops.ADD, Ops.MAX}: return None
  reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].sparents)
  if len(reduce_unparented) == 0: return None
  ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
  if red.arg is Ops.ADD:
    for r in reduce_unparented: ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
-  if red.arg is Ops.MUL:
-    for r in reduce_unparented: ret = ret ** r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
  return ret

 pm_reduce_simplify = PatternMatcher([
@@ -279,24 +279,22 @@ class Compiled:

  def __init__(self, device:str, allocator:Allocator, compilers:Sequence[CompilerPairT]|None, runtime, graph=None, group_id=None):
    self.device, self.allocator, self.runtime, self.graph, self.group_id = device, allocator, runtime, graph, group_id
-    self.compilers = cast(list[CompilerPairT], compilers or [(Renderer, Compiler)])
+    compilers = cast(list[CompilerPairT], compilers or [(Renderer, Compiler)])

-    envnames = [self._get_compiler_envvar(c) for r,c in self.compilers]
-    enable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, self.compilers) if en is not None and getenv(en, -1) == 1)
-    disable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, self.compilers) if en is not None and getenv(en, -1) == 0)
+    devname = device.split(':')[0].upper()
+    envnames = [f"{devname}_{unwrap_class_type(c).__name__.removesuffix('Compiler').removeprefix(devname).upper()}" for r,c in compilers]
+
+    enable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, compilers) if en is not None and getenv(en, -1) == 1)
+    disable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, compilers) if en is not None and getenv(en, -1) == 0)

    if len(enable_comps) > 1: raise RuntimeError(f"{self.device}: multiple compilers set in env {enable_comps}")
-    for _, comp_pair in disable_comps: self.compilers.remove(comp_pair)
+    for _, comp_pair in disable_comps: compilers.remove(comp_pair)

-    try: self.renderer, self.compiler = next(self._get_available_compilers([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers))
+    try: self.renderer, self.compiler = next(self._get_available_compilers([list(enable_comps)[0][1]] if len(enable_comps) == 1 else compilers))
    except StopIteration as exc: raise RuntimeError(f"no usable compilers for {self.device}") from exc

    if DEBUG >= 1: print(f"{self.device}: using {self.compiler.__class__.__name__}")

-  def _get_compiler_envvar(self, c):
-    compiler_name = f"{unwrap_class_type(c).__name__.upper().removesuffix('COMPILER').removeprefix(devname:=self.device.split(':')[0].upper())}"
-    return f"{devname}_{compiler_name if len(compiler_name) > 0 else unwrap_class_type(c).__name__.upper()}"
-
  def _get_available_compilers(self, compilers) -> Iterator[tuple[Renderer, Compiler]]:
    for renderer, compiler in compilers:
      with contextlib.suppress(Exception): yield renderer(), compiler()
@@ -327,8 +325,10 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
    if device == "METAL": return not CI
    if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX")
    if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"}
-    return device in {"AMD", "PYTHON", "NULL"}
-  if dtype in dtypes.fp8s: return device in {"PYTHON", "NULL"}
+    return device in {"AMD", "PYTHON"}
+  if dtype in dtypes.fp8s:
+    # not supported yet - in progress
+    return False
  if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,
                                          dtypes.ushort, dtypes.float, dtypes.int32, dtypes.uint32, dtypes.half]
  # for CI GPU and OSX, cl_khr_fp16 isn't supported
@@ -352,26 +352,21 @@ if PROFILE:

    with open(fn:=temp("profile.pkl", append_user=True), "wb") as f: pickle.dump(cpu_events+Compiled.profile_events+Buffer.profile_events, f)

-    from tinygrad.uop.ops import launch_viz
-    launch_viz("PROFILE", fn)
+    if not getenv("SQTT", 0):
+      from tinygrad.uop.ops import launch_viz
+      launch_viz(PROFILE, fn)

 if __name__ == "__main__":
-  from tinygrad import Tensor, Device
-
  for device in ALL_DEVICES:
-    compilers_results, any_works = [], False
    try:
-      default_compiler = (d:=Device[device]).compiler
-      for i,(r,c) in enumerate(d.compilers):
-        try:
-          d.renderer, d.compiler = r(), c()
-          with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist()
-          if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
-          default_text = '(default)' if type(default_compiler) is type(d.compiler) else f'({d._get_compiler_envvar(c)}=1 to make default)'
-          compilers_results.append(f"{colored('+', 'green')} {unwrap_class_type(c).__name__} {default_text}")
-          any_works = True
-        except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {unwrap_class_type(c).__name__}: {e}")
-      result = (colored('PASS', 'green') if any_works else f"{colored('FAIL', 'yellow')}") + ''.join([f'\n{" "*16} {x}' for x in compilers_results])
+      _ = Device[device].device
+      try:
+        from tinygrad import Tensor
+        with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist()
+        if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
+        result = colored("PASS", "green")
+      except Exception as e:
+        result = f"{colored('FAIL', 'yellow')} {e}"
    except Exception as e:
      result = f"{colored('FAIL', 'red')} {e}"
    print(f"{'*' if device == Device.DEFAULT else ' '} {device:10s}: {result}")
@@ -32,9 +32,7 @@ class DTypeMetaClass(type):
    DTypeMetaClass.dcache[args] = ret = super().__call__(*args)
    return ret

-class AddrSpace(Enum):
-  def __repr__(self): return str(self)
-  GLOBAL = auto(); LOCAL = auto(); REG = auto()  # noqa: E702
+class AddrSpace(Enum): GLOBAL = auto(); LOCAL = auto(); REG = auto()  # noqa: E702

@dataclass(frozen=True, eq=False)
 class DType(metaclass=DTypeMetaClass):
@@ -235,7 +233,7 @@ def sum_acc_dtype(dt:DType):
  if dtypes.is_int(dt) or dt == dtypes.bool: return least_upper_dtype(dt, dtypes.int)
  return least_upper_dtype(dt, to_dtype(getenv("SUM_DTYPE", "float32")))

-def float_to_fp16(x):
+def truncate_fp16(x):
  try: return struct.unpack('e', struct.pack('e', float(x)))[0]
  except OverflowError: return math.copysign(math.inf, x)

@@ -312,7 +310,7 @@ def fp8_to_float(x: int, dtype: DType) -> float:
  return float(float32_val)

 truncate: dict[DType, Callable] = {dtypes.bool: bool,
-  dtypes.float16: float_to_fp16, dtypes.bfloat16: lambda x: float_to_bf16(float(x)),
+  dtypes.float16: truncate_fp16, dtypes.bfloat16: lambda x: float_to_bf16(float(x)),
  **{fp8: (lambda x, dtype=fp8: fp8_to_float(float_to_fp8(x, dtype), dtype)) for fp8 in dtypes.fp8s},
  dtypes.float32: lambda x: ctypes.c_float(x).value, dtypes.float64: lambda x: ctypes.c_double(x).value,
  dtypes.uint8: lambda x: ctypes.c_uint8(x).value, dtypes.uint16: lambda x: ctypes.c_uint16(x).value,
@@ -324,7 +322,7 @@ truncate: dict[DType, Callable] = {dtypes.bool: bool,

 def _to_np_dtype(dtype:DType) -> type|None:
  import numpy as np
-  if dtype in { dtypes.bfloat16, *dtypes.fp8s }: return np.float32
+  if dtype == dtypes.bfloat16: return np.float32
  return np.dtype(dtype.fmt).type if dtype.fmt is not None else None
 def _from_np_dtype(npdtype:'np.dtype') -> DType: # type: ignore [name-defined] # noqa: F821
  import numpy as np
@@ -335,7 +333,6 @@ def _to_torch_dtype(dtype:DType) -> 'torch.dtype'|None:  # type: ignore [name-de
  import numpy as np, torch
  if dtype == dtypes.uint64: return torch.uint64
  if dtype == dtypes.bfloat16: return torch.bfloat16
-  if dtype in dtypes.fp8s: return torch.uint8
  # NOTE: torch doesn't expose this mapping with a stable API
  try: return torch.from_numpy(np.array([], dtype=_to_np_dtype(dtype))).dtype
  except TypeError: return None
@@ -140,13 +140,13 @@ class BufferXfer(BufferCopy):

 # **************** method cache ****************

-method_cache: dict[tuple[str, type, bytes, tuple[int, ...], bool], CompiledRunner] = {}
+method_cache: dict[tuple[str, bytes, tuple[int, ...], bool], CompiledRunner] = {}
 def get_runner(device:str, ast:UOp) -> CompiledRunner:
  # TODO: this should be all context relevant to rendering
  context = (BEAM.value, NOOPT.value, DEVECTORIZE.value)
-  ckey = (device, type(Device[device].compiler), ast.key, context, False)
+  ckey = (device, ast.key, context, False)
  if cret:=method_cache.get(ckey): return cret
-  bkey = (device.split(":")[0], type(Device[device].compiler), ast.key, context, True)
+  bkey = (device.split(":")[0], ast.key, context, True)
  if bret:=method_cache.get(bkey):
    method_cache[ckey] = ret = CompiledRunner(replace(bret.p, device=device), bret.lib)
  else:
@@ -135,7 +135,7 @@ USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1
 TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS", 0)
 FUSE_ARANGE, FUSE_CONV_BW = ContextVar("FUSE_ARANGE", 1), ContextVar("FUSE_CONV_BW", 0)
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1)
-PICKLE_BUFFERS, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("LRU", 1)
+PICKLE_BUFFERS, PROFILE, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("PROFILE", getenv("VIZ")), ContextVar("LRU", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
 DISABLE_COMPILER_CACHE, BLOCK_REORDER = ContextVar("DISABLE_COMPILER_CACHE", 0), ContextVar("BLOCK_REORDER", 1)
 DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES = ContextVar("DONT_REALIZE_EXPAND", 0), ContextVar("DONT_GROUP_REDUCES", 0)
@@ -146,8 +146,6 @@ RANGEIFY, FUSE_ATTENTION = ContextVar("RANGEIFY", 0), ContextVar("FUSE_ATTENTION
 EMULATE = ContextVar("EMULATE", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, (os.cpu_count() or 1) // (4 if ARCH_X86 else 2))) # take 1/2 of the cores, accounting HT
 CPU_LLVM, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("AMD_LLVM", 1)
-VIZ = PROFILE = ContextVar("VIZ", 0)
-SPEC = ContextVar("SPEC", 0)

@dataclass(frozen=True)
 class Metadata:
@@ -326,10 +324,7 @@ def cpu_objdump(lib, objdump_tool='objdump'):
    print(subprocess.check_output([objdump_tool, '-d', f.name]).decode('utf-8'))

 def capstone_flatdump(lib: bytes):
-  try: import capstone
-  except ImportError:
-    print("Disassembler Error: Capstone not installed.")
-    return
+  import capstone
  match platform.machine():
    case 'x86_64' | 'AMD64': cs = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)
    case 'aarch64' | 'arm64': cs = capstone.Cs(capstone.CS_ARCH_ARM64, capstone.CS_MODE_ARM)
@@ -3,7 +3,6 @@ import math, struct, sys
 from tinygrad.codegen.opt import tc
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import AMDRenderer
-from tinygrad.uop.decompositions import xexp2, xlog2
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, sint_to_uop
 from tinygrad.dtype import dtypes, DType, PtrDType, truncate
 from tinygrad.helpers import prod, AMX
@@ -107,8 +106,7 @@ base_rewrite = PatternMatcher([
   f"  {ctx[x]} = phi {ldt(x.dtype)} [ 0, %loop_entry_{x.arg[0]} ], [ {ctx[x]}phi, %loop_latch_{x.arg[0]} ]"),
  (UPat(Ops.ENDRANGE, name="x"), lambda ctx,x:
   f"  br label %loop_latch_{x.src[0].arg[0]}\nloop_latch_{x.src[0].arg[0]}:\n"
-   f"  {ctx[x.src[0]]}phi = add {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, 1\n"
-   f"  {ctx[x]} = icmp ult {ldt(x.src[0].dtype)} {ctx[x.src[0]]}phi, {ctx[x.src[0].src[0]]}\n"
+   f"  {ctx[x.src[0]]}phi = add i32 {ctx[x.src[0]]}, 1\n  {ctx[x]} = icmp ult i32 {ctx[x.src[0]]}phi, {ctx[x.src[0].src[0]]}\n"
   f"  br i1 {ctx[x]}, label %loop_body_{x.src[0].arg[0]}, label %loop_exit_{x.src[0].arg[0]}\nloop_exit_{x.src[0].arg[0]}:"),

  # if
@@ -199,7 +197,8 @@ barrier = 'fence syncscope("workgroup") release\ntail call void @llvm.amdgcn.s.b
 code_for_workitem = {"g": lambda x: f"tail call i32 @llvm.amdgcn.workgroup.id.{chr(120+int(x))}()",
                     "l": lambda x: f"tail call i32 @llvm.amdgcn.workitem.id.{chr(120+int(x))}()"}
 # https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#llvm-ir-intrinsics
-llvm_intrinsics = {Ops.SQRT: "sqrt", Ops.LOG2: "log2", Ops.EXP2: "exp2"}
+# llvm.log2/llvm.exp2 don't support double
+llvm_intrinsics = {Ops.SQRT: "sqrt"}
 class AMDLLVMRenderer(LLVMRenderer):
  device = "AMD"
  has_local = True
@@ -218,9 +217,6 @@ class AMDLLVMRenderer(LLVMRenderer):
      lambda x, y: UOp(Ops.VECTORIZE, dtypes.half.vec(16), tuple(y.gep(i // 2) if i % 2 == 0 else UOp.const(dtypes.half, 0.0) for i in range(16)))),
    (UPat(Ops.CAST, name="x", dtype=dtypes.half.vec(8), src=UPat.var("y", dtypes.half.vec(16))),
      lambda x, y: UOp(Ops.VECTORIZE, dtypes.half.vec(8), tuple(y.gep(i * 2) for i in range(8)))),
-    # amd llvm intrinsics llvm.log2/llvm.exp2 don't support double
-    (UPat(Ops.LOG2, dtype=dtypes.double, src=(UPat.var("d"),)), xlog2),
-    (UPat(Ops.EXP2, dtype=dtypes.double, src=(UPat.var("d"),)), xexp2),
  ])
  def _render_footer(self, uops: list[UOp]) -> str:
    # TODO: this is copied from cstyle
@@ -581,9 +581,6 @@ class KFDIface:
    if uncached: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED | kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT
    else: flags |= (kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR if host else kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM)

-    # Make mapped cpu address to be uncachable
-    if cpu_addr is not None: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED
-
    if cpu_access or host: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_PUBLIC

    if flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR:
@@ -48,8 +48,8 @@ class CLProgram:

  def __call__(self, *bufs:tuple[ctypes._CData, BufferSpec], global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]|None=None,
               vals:tuple[int, ...]=(), wait=False) -> float|None:
-    for i,(b,_) in enumerate(bufs): check(cl.clSetKernelArg(self.kernel, i, ctypes.sizeof(b), ctypes.byref(b)))
-    for i,v in enumerate(vals,start=len(bufs)): check(cl.clSetKernelArg(self.kernel, i, 4, ctypes.byref(ctypes.c_int32(v))))
+    for i,(b,_) in enumerate(bufs): cl.clSetKernelArg(self.kernel, i, ctypes.sizeof(b), ctypes.byref(b))
+    for i,v in enumerate(vals,start=len(bufs)): cl.clSetKernelArg(self.kernel, i, 4, ctypes.byref(ctypes.c_int32(v)))
    if local_size is not None: global_size = cast(tuple[int,int,int], tuple(int(g*l) for g,l in zip(global_size, local_size)))
    event = cl.cl_event() if wait else None
    check(cl.clEnqueueNDRangeKernel(self.dev.queue, self.kernel, len(global_size), None, (ctypes.c_size_t * len(global_size))(*global_size),
@@ -1,6 +1,6 @@
 from __future__ import annotations
 import platform, sys, ctypes, functools, time, mmap, threading, queue
-from tinygrad.helpers import from_mv, to_mv, OSX, WIN, mv_address, wait_cond, cpu_profile, suppress_finalizing, unwrap
+from tinygrad.helpers import from_mv, to_mv, OSX, WIN, mv_address, wait_cond, cpu_profile, suppress_finalizing
 from tinygrad.device import BufferSpec, DMACPURef
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
 from tinygrad.renderer.cstyle import ClangRenderer
@@ -56,9 +56,7 @@ class CPUComputeQueue(HWQueue):
 MAP_JIT = 0x0800

 class CPUProgram(HCQProgram):
-  rt_lib = None
-  try: rt_lib = ctypes.CDLL(ctypes.util.find_library('System' if OSX else 'kernel32') if OSX or WIN else 'libgcc_s.so.1')
-  except OSError: pass
+  rt_lib = ctypes.CDLL(ctypes.util.find_library('System' if OSX else 'kernel32') if OSX or WIN else 'libgcc_s.so.1')

  def __init__(self, dev, name:str, lib:bytes):
    if sys.platform == "win32": # mypy doesn't understand when WIN is used here
@@ -75,20 +73,15 @@ class CPUProgram(HCQProgram):
      # MAP_JIT allows us to easily flip pages from RW- to R-X and vice versa. It is a noop on intel cpus. (man pthread_jit_write_protect_np)
      self.mem = mmap.mmap(-1, len(lib), mmap.MAP_ANON|mmap.MAP_PRIVATE|(MAP_JIT if OSX else 0), mmap.PROT_READ|mmap.PROT_WRITE|mmap.PROT_EXEC)

-      if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(False)
+      if OSX: CPUProgram.rt_lib.pthread_jit_write_protect_np(False)
      self.mem.write(lib)
-      if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(True)
+      if OSX: CPUProgram.rt_lib.pthread_jit_write_protect_np(True)

      # __clear_cache isn't a normal libc function, but a compiler support routine found in libgcc_s for gcc and compiler-rt for clang.
      # libgcc_s comes as shared library but compiler-rt is only a bunch of static library archives which we can't directly load, but fortunately
      # it somehow found its way into libSystem on macos (likely because it used __builtin_clear_cache) and libgcc_s is ~always present on linux
      # Using ["name"] instead of .name because otherwise name is getting mangled: https://docs.python.org/3.12/reference/expressions.html#index-5
-      if CPUProgram.rt_lib is not None:
-        CPUProgram.rt_lib["__clear_cache"](ctypes.c_void_p(mv_address(self.mem)), ctypes.c_void_p(mv_address(self.mem) + len(lib)))
-      else:
-        # msync should be a universal POSIX way to do this
-        from tinygrad.runtime.autogen import libc
-        libc.msync(ctypes.c_void_p(mv_address(self.mem)), len(lib), libc.MS_SYNC | libc.MS_INVALIDATE)
+      CPUProgram.rt_lib["__clear_cache"](ctypes.c_void_p(mv_address(self.mem)), ctypes.c_void_p(mv_address(self.mem) + len(lib)))

      self.fxn = ctypes.CFUNCTYPE(None)(mv_address(self.mem))

@@ -39,11 +39,7 @@ class DiskDevice(Compiled):
  def _might_close(self):
    self.count -= 1
    if self.count == 0:
-      if self.fd is not None:
-        os.close(self.fd)
-      if hasattr(self, "mem"):
-        try: self.mem.close()
-        except BufferError: pass
+      if self.fd is not None: os.close(self.fd)
      self.size = None
  def _iouring_setup(self):
    DiskDevice._tried_io_uring_init = True
@@ -4,23 +4,21 @@
 # this is the (living) definition of uops
 from typing import Any, TYPE_CHECKING, cast
 import pickle, base64, itertools, time, struct, sys
-from tinygrad.dtype import DType, dtypes, ImageDType, PtrDType, truncate, float_to_bf16, float_to_fp8, fp8_to_float
+from tinygrad.dtype import DType, dtypes, ImageDType, PtrDType, truncate, float_to_bf16
 from tinygrad.helpers import all_same, getenv, flatten, get_single_element, EMULATE
 from tinygrad.device import Compiled, Compiler, Allocator
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
 from tinygrad.renderer import Renderer

-def storage_fmt_for_dtype(dtype: DType): return 'H' if dtype == dtypes.bfloat16 else 'B' if dtype in dtypes.fp8s else dtype.fmt
+def storage_fmt_for_dtype(dtype: DType): return 'H' if dtype == dtypes.bfloat16 else dtype.fmt

 def to_storage_scalar(x, dtype: DType):
  if dtype == dtypes.bfloat16: return (struct.unpack('I', struct.pack('f', float_to_bf16(x)))[0] >> 16) & 0xFFFF
-  if dtype in dtypes.fp8s: return float_to_fp8(float(x), dtype)
  return x

 def from_storage_scalar(x, dtype: DType):
  if dtype == dtypes.bfloat16: return struct.unpack('f', struct.pack('I', (x & 0xFFFF) << 16))[0]
-  if dtype in dtypes.fp8s: return fp8_to_float(int(x), dtype)
  return x

 def _load(m, i, dtype: DType):
@@ -84,7 +84,7 @@ class AMFirmware:
    self.descs += [self.desc(blob, hdr0.header.ucode_array_offset_bytes, hdr0.header.ucode_size_bytes, am.GFX_FW_TYPE_RLC_G)]

  def load_fw(self, fname:str, *headers, versioned_header:str|None=None):
-    fpath = fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/a9f26799247aa60fbaa3b64267a18f20b72b5235/amdgpu/{fname}", subdir="fw")
+    fpath = fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/45f59212aebd226c7630aff4b58598967c0c8c91/amdgpu/{fname}", subdir="fw")
    blob = memoryview(bytearray(fpath.read_bytes()))
    if AM_DEBUG >= 1: print(f"am {self.adev.devfmt}: loading firmware {fname}: {hashlib.sha256(blob).hexdigest()}")
    if versioned_header:
@@ -13,7 +13,7 @@ class ClangJITCompiler(Compiler):
    # x18 is a reserved platform register. It is clobbered on context switch in macos and is used to store TEB pointer in windows on arm, don't use it
    target = 'x86_64' if sys.platform == 'win32' else platform.machine()
    # on arm march means "runs on this arch and superset" instead of "optimize for this arch". x86 march == arm mcpu
-    arch = {'x86_64': '-march=native', 'AMD64': '-march=native', 'riscv64': '-march=rv64g'}.get(platform.machine(), "-mcpu=native")
+    arch = '-march=native' if platform.machine() in ('x86_64', 'AMD64') else '-mcpu=native'
    args = [arch, f'--target={target}-none-unknown-elf', '-O2', '-fPIC', '-ffreestanding', '-fno-math-errno', '-nostdlib', '-fno-ident']
    arch_args = ['-ffixed-x18'] if target == 'arm64' else []
    obj = subprocess.check_output([getenv("CC", 'clang'), '-c', '-x', 'c', *args, *arch_args, '-', '-o', '-'], input=src.encode('utf-8'))
@@ -29,7 +29,7 @@ def expect(x, err, ret=None):

 class LLVMCompiler(Compiler):
  jit = True
-  target_arch = {'arm64': 'AArch64', 'aarch64': 'AArch64', 'x86_64': 'X86', 'AMD64': 'X86', 'riscv64': 'riscv64'}[platform.machine()]
+  target_arch = {'arm64': 'AArch64', 'aarch64': 'AArch64', 'x86_64': 'X86', 'AMD64': 'X86'}[platform.machine()]
  def __init__(self, processor:str, feats:str):
    for component in ['Target', 'TargetInfo', 'TargetMC', 'AsmParser', 'AsmPrinter']: getattr(llvm, f'LLVMInitialize{self.target_arch}{component}')()

@@ -4,7 +4,7 @@ from tinygrad.helpers import to_char_p_p, colored, init_c_var, getenv
 import tinygrad.runtime.autogen.nvrtc as nvrtc
 from tinygrad.device import Compiler, CompileError

-CUDA_PATH = getenv("CUDA_PATH", "")
+CUDA_PATH = getenv("CUDA_PATH", "")  # PTX shouldn't be here, in fact, it shouldn't exist

 def _get_bytes(arg, get_str, get_sz, check) -> bytes:
  sz = init_c_var(ctypes.c_size_t(), lambda x: check(get_sz(arg, ctypes.byref(x))))
--- a/Show More
+++ b/Show More