ugh

fix tests
reset ending ranges
2025-10-17 23:14:36 +08:00 · 2025-10-17 23:05:01 +08:00 · 2025-10-17 22:52:43 +08:00 · 2025-10-17 22:41:31 +08:00 · 2025-10-17 22:41:18 +08:00 · 2025-10-17 10:32:41 -04:00
132 changed files with 24068 additions and 3272 deletions
@@ -41,6 +41,10 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
+  mesa:
+    description: "Install mesa"
+    required: false
+    default: 'false'
 runs:
  using: "composite"
  steps:
@@ -289,3 +293,13 @@ runs:
      if: inputs.llvm == 'true' && runner.os == 'macOS'
      shell: bash
      run: brew install llvm@20
+
+    # **** mesa ****
+    - name: Install mesa (linux)
+      if: inputs.mesa == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: sudo curl -L https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
+    - name: Install mesa (macOS)
+      if: inputs.mesa == 'true' && runner.os == 'macOS'
+      shell: bash
+      run: brew install sirhcm/tinymesa/tinymesa
@@ -36,8 +36,9 @@ jobs:
        cuda: 'true'
        webgpu: 'true'
        llvm: 'true'
+        pydeps: 'pyyaml mako'
    - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
+      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev llvm-20-dev
    - name: Verify OpenCL autogen
      run: |
        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
@@ -89,3 +90,8 @@ jobs:
        cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
        ./autogen_stubs.sh llvm
        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
+    - name: Verify mesa autogen
+      run: |
+        cp tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
+        ./autogen_stubs.sh mesa
+        diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
@@ -51,17 +51,18 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: python3.11 test/external/process_replay/reset.py
+    - name: Print macOS version
+      run: sw_vers
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=900 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
      # TODO: very slow step time
-      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=10000 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
    # process replay can't capture this, the graph is too large
-    # TODO: too slow
-    # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run model inference benchmark
      run: METAL=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
@@ -632,7 +633,7 @@ jobs:
      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 Space Lab policy + vision
      run: |
-        PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
+        PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
        PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
    - name: benchmark MobileNetV2 on DSP
      run: |
@@ -641,7 +642,7 @@ jobs:
        ln -s /data/home/tiny/tinygrad/testsig-*.so .
        PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
        # benchmark on DSP with NOOPT=1, the devectorizer has issues
-        PYTHONPATH=. CC=clang-19 DSP=1 DONT_REALIZE_EXPAND=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
    - uses: actions/upload-artifact@v4
@@ -12,7 +12,7 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 360
+    timeout-minutes: 720

    steps:
    - name: Checkout Code
@@ -89,64 +89,65 @@ jobs:
        clang -O2 recognize.c -lm -o recognize
        cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock

-  torchbackend:
-    name: Torch Backend Tests
-    runs-on: ubuntu-latest
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_minimal
-        pydeps: "pillow torchvision expecttest"
-        llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
-    - name: Lint with ruff
-      run: |
-        pip3 install --upgrade --force-reinstall ruff==0.11.0
-        python3 -m ruff check extra/torch_backend/backend.py
-    - name: Test one op
-      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
-    - name: Test ResNet-18
-      run: DEBUG=2 python3 extra/torch_backend/example.py
-    - name: My (custom) tests
-      run: python3 extra/torch_backend/test.py
-    - name: Test one op in torch tests
-      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
-    - name: Test Ops with TINY_BACKEND
-      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
-    - name: Test in-place operations on views
-      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
-    - name: Test multi-gpu
-      run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+  # TODO: fix the torch backend and reenable
+  # torchbackend:
+  #   name: Torch Backend Tests
+  #   runs-on: ubuntu-latest
+  #   timeout-minutes: 15
+  #   steps:
+  #   - name: Checkout Code
+  #     uses: actions/checkout@v4
+  #   - name: Setup Environment
+  #     uses: ./.github/actions/setup-tinygrad
+  #     with:
+  #       key: torch-backend-pillow-torchvision-et-pt
+  #       deps: testing_minimal
+  #       pydeps: "pillow torchvision expecttest"
+  #       llvm: 'true'
+  #   - name: Install ninja
+  #     run: |
+  #       sudo apt update || true
+  #       sudo apt install -y --no-install-recommends ninja-build
+  #   - name: Lint with ruff
+  #     run: |
+  #       pip3 install --upgrade --force-reinstall ruff==0.11.0
+  #       python3 -m ruff check extra/torch_backend/backend.py
+  #   - name: Test one op
+  #     run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+  #   - name: Test ResNet-18
+  #     run: DEBUG=2 python3 extra/torch_backend/example.py
+  #   - name: My (custom) tests
+  #     run: python3 extra/torch_backend/test.py
+  #   - name: Test one op in torch tests
+  #     run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+  #   - name: Test Ops with TINY_BACKEND
+  #     run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+  #   - name: Test in-place operations on views
+  #     run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+  #   - name: Test multi-gpu
+  #     run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py

-  torchbackendmore:
-    name: Torch Backend Tests More
-    runs-on: ubuntu-latest
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_minimal
-        llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
-    - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: CPU=1 CPU_LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
-    - name: Test some torch tests (expect failure)
-      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+  # torchbackendmore:
+  #   name: Torch Backend Tests More
+  #   runs-on: ubuntu-latest
+  #   timeout-minutes: 15
+  #   steps:
+  #   - name: Checkout Code
+  #     uses: actions/checkout@v4
+  #   - name: Setup Environment
+  #     uses: ./.github/actions/setup-tinygrad
+  #     with:
+  #       key: torch-backend-pillow-torchvision-et-pt
+  #       deps: testing_minimal
+  #       llvm: 'true'
+  #   - name: Install ninja
+  #     run: |
+  #       sudo apt update || true
+  #       sudo apt install -y --no-install-recommends ninja-build
+  #   - name: Test beautiful_mnist in torch with TINY_BACKEND
+  #     run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+  #   - 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
@@ -238,8 +239,6 @@ jobs:
        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check .
        python3 -m ruff check examples/mlperf/ --ignore E501
-    - name: Lint tinygrad with pylint
-      run: python -m pylint tinygrad/
    - name: Run mypy
      run: |
        python -m mypy --strict-equality --lineprecision-report .
@@ -274,6 +273,8 @@ jobs:
    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
    - name: Run Clip tests for SD MLPerf on NULL backend
      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
+    - name: Run AMD emulated BERT training on NULL backend
+      run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    # 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
@@ -309,10 +310,6 @@ jobs:
      run: python test/external/fuzz_symbolic.py
    - name: Fuzz Test fast idiv
      run: python test/external/fuzz_fast_idiv.py
-    - name: Fuzz Test shapetracker
-      run: CNT=50 python test/external/fuzz_shapetracker.py
-    - name: Fuzz Test shapetracker math
-      run: CNT=200 python test/external/fuzz_shapetracker_math.py
    - name: Fuzz Test shape ops
      run: python test/external/fuzz_shape_ops.py

@@ -386,8 +383,6 @@ jobs:
      #   run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
      - name: Test openpilot LLVM compile
        run: CPU=1 CPU_LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
-      - name: Test openpilot compile4
-        run: NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -677,7 +672,7 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        backend: [llvm, cpu, opencl]
+        backend: [llvm, cpu, opencl, lvp]

    name: Linux (${{ matrix.backend }})
    runs-on: ubuntu-22.04
@@ -691,9 +686,10 @@ jobs:
          key: ${{ matrix.backend }}-minimal
          deps: testing_minimal
          opencl: ${{ matrix.backend == 'opencl' && 'true' }}
-          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
+          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
@@ -895,7 +891,7 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        backend: [metal, llvm, cpu]
+        backend: [metal, llvm, cpu, lvp]
    name: MacOS (${{ matrix.backend }})
    runs-on: macos-15
    timeout-minutes: 20
@@ -908,12 +904,13 @@ jobs:
          key: macos-${{ matrix.backend }}-minimal
          deps: testing_minimal
          pydeps: "capstone"
-          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
+          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (${{ matrix.backend }})
        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
@@ -38,6 +38,7 @@ extra/huggingface_onnx/models/*
 extra/huggingface_onnx/*.yaml
 extra/weights
 venv
+venv_sd_mlperf
 examples/**/net.*[js,json]
 examples/**/*.safetensors
 node_modules
@@ -20,21 +20,15 @@ repos:
        language: system
        always_run: true
        pass_filenames: false
-      - id: tests
-        name: subset of tests
-        entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
-        language: system
-        always_run: true
-        pass_filenames: false
      - id: example
-        name: multi device tests
+        name: test all devices
        entry: python3 test/external/external_test_example.py
        language: system
        always_run: true
        pass_filenames: false
-      - id: pylint
-        name: pylint
-        entry: python3 -m pylint tinygrad/
+      - id: tests
+        name: subset of tests
+        entry: env PYTHONPATH="." python3 -m pytest -n=8 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
        language: system
        always_run: true
        pass_filenames: false
@@ -461,6 +461,85 @@ generate_libusb() {
  python3 -c "import tinygrad.runtime.autogen.libusb"
 }

+generate_mesa() {
+  MESA_TAG="mesa-25.2.4"
+  MESA_SRC=/tmp/mesa-$MESA_TAG
+  TINYMESA_TAG=tinymesa-32dc66c
+  TINYMESA_DIR=/tmp/tinymesa-$MESA_TAG-$TINYMESA_TAG/
+  TINYMESA_SO=$TINYMESA_DIR/libtinymesa_cpu.so
+  if [ ! -d "$MESA_SRC" ]; then
+    git clone --depth 1 --branch $MESA_TAG https://gitlab.freedesktop.org/mesa/mesa.git $MESA_SRC
+    pushd .
+    cd $MESA_SRC
+    git reset --hard $MESA_COMMIT_HASH
+    # clang 14 doesn't support packed enums
+    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/headers/nv_device_info.h
+    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/compiler/nak.h
+    sed -i "s/nir_instr_type \(\w\+\);/uint8_t \1;/" $MESA_SRC/src/compiler/nir/nir.h
+    mkdir -p gen/util/format
+    python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h
+    python3 src/compiler/nir/nir_opcodes_h.py > gen/nir_opcodes.h
+    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
+    python3 src/compiler/nir/nir_intrinsics_indices_h.py --outdir gen
+    python3 src/compiler/nir/nir_builder_opcodes_h.py > gen/nir_builder_opcodes.h
+    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
+    python3 src/compiler/builtin_types_h.py gen/builtin_types.h
+    popd
+  fi
+
+  if [ ! -d "$TINYMESA_DIR" ]; then
+    mkdir $TINYMESA_DIR
+    curl -L https://github.com/sirhcm/tinymesa/releases/download/$TINYMESA_TAG/libtinymesa_cpu-$MESA_TAG-linux-amd64.so -o $TINYMESA_SO
+  fi
+
+  clang2py -k cdefstu \
+    $MESA_SRC/src/compiler/nir/nir.h \
+    $MESA_SRC/src/compiler/nir/nir_builder.h \
+    $MESA_SRC/src/compiler/nir/nir_shader_compiler_options.h \
+    $MESA_SRC/src/compiler/nir/nir_serialize.h \
+    $MESA_SRC/gen/nir_intrinsics.h \
+    $MESA_SRC/src/nouveau/headers/nv_device_info.h \
+    $MESA_SRC/src/nouveau/compiler/nak.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_passmgr.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_misc.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_type.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_init.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_nir.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_struct.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_jit_types.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_flow.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_const.h \
+    $MESA_SRC/src/compiler/glsl_types.h \
+    $MESA_SRC/src/util/blob.h \
+    $MESA_SRC/src/util/ralloc.h \
+    --clang-args="-DHAVE_ENDIAN_H -DHAVE_STRUCT_TIMESPEC -DHAVE_PTHREAD -I$MESA_SRC/src -I$MESA_SRC/include -I$MESA_SRC/gen -I$MESA_SRC/src/compiler/nir -I$MESA_SRC/src/gallium/auxiliary -I$MESA_SRC/src/gallium/include -I$(llvm-config-20 --includedir)" \
+    -l $TINYMESA_SO \
+    -o $BASE/mesa.py
+
+  LVP_NIR_OPTIONS=$(./extra/mesa/lvp_nir_options.sh $MESA_SRC)
+
+  fixup $BASE/mesa.py
+  patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "f'{brew_prefix()}/lib/libtinymesa_cpu.dylib'"
+  echo "lvp_nir_options = gzip.decompress(base64.b64decode('$LVP_NIR_OPTIONS'))" >> $BASE/mesa.py
+  cat <<EOF | sed -i "/import ctypes.*/r /dev/stdin" $BASE/mesa.py
+def brew_prefix():
+  try: return subprocess.check_output(['brew', '--prefix', 'tinymesa']).decode().strip()
+  except Exception: return ''
+EOF
+  sed -i "/in_dll/s/.*/try: &\nexcept AttributeError: pass/" $BASE/mesa.py
+  sed -i "s/import ctypes/import ctypes, ctypes.util, os, gzip, base64, subprocess, tinygrad.helpers as helpers/" $BASE/mesa.py
+  sed -i "s/ctypes.CDLL('.\+')/(dll := _try_dlopen_tinymesa_cpu())/" $BASE/mesa.py
+  echo "def __getattr__(nm): raise AttributeError() if dll else FileNotFoundError(f'libtinymesa not found (MESA_PATH={BASE}). See https://github.com/sirhcm/tinymesa ($TINYMESA_TAG, $MESA_TAG)')" >> $BASE/mesa.py
+  sed -i "s/ctypes.glsl_base_type/glsl_base_type/" $BASE/mesa.py
+  # bitfield bug in clang2py
+  sed -i "s/('fp_fast_math', ctypes.c_bool, 9)/('fp_fast_math', ctypes.c_uint32, 9)/" $BASE/mesa.py
+  sed -i "s/('\(\w\+\)', pipe_shader_type, 8)/('\1', ctypes.c_ubyte)/" $BASE/mesa.py
+  sed -i "s/\([0-9]\+\)()/\1/" $BASE/mesa.py
+  sed -i "s/\(struct_nir_builder._pack_\) = 1/\1 = 0/" $BASE/mesa.py
+  python3 -c "import tinygrad.runtime.autogen.mesa"
+}
+
 if [ "$1" == "opencl" ]; then generate_opencl
 elif [ "$1" == "hip" ]; then generate_hip
 elif [ "$1" == "comgr" ]; then generate_comgr
@@ -484,6 +563,7 @@ elif [ "$1" == "pci" ]; then generate_pci
 elif [ "$1" == "vfio" ]; then generate_vfio
 elif [ "$1" == "webgpu" ]; then generate_webgpu
 elif [ "$1" == "libusb" ]; then generate_libusb
-elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu
+elif [ "$1" == "mesa" ]; then generate_mesa
+elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu; generate_mesa
 else echo "usage: $0 <type>"
 fi
@@ -232,7 +232,7 @@ if __name__ == "__main__":
  gpt2 = GPT2.build_gguf(args.model_size) if args.model_size.startswith("gpt2_gguf_") else GPT2.build(args.model_size)

  if args.benchmark != -1:
-    gpt2.model(Tensor.rand(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
+    gpt2.model(Tensor.randint(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
  else:
    texts = gpt2.generate(args.prompt, args.count, args.temperature, timing=args.timing, batch_size=args.batch_size)
    if not args.noshow:
@@ -2,7 +2,7 @@

 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -2,7 +2,7 @@

 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -5,7 +5,7 @@ set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." NV=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -2,7 +2,7 @@

 export PYTHONPATH="." AMD=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -2,7 +2,7 @@

 export PYTHONPATH="." AMD=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -5,7 +5,7 @@ set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_red"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -1,47 +0,0 @@
-import sys
-from tinygrad import Tensor, fetch, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp
-from tinygrad.nn.onnx import OnnxRunner
-from tinygrad.schedule.rangeify import get_rangeify_map
-from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.engine.realize import run_schedule
-
-# NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
-
-OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
-OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
-
-if __name__ == "__main__":
-  onnx_file = fetch(OPENPILOT_MODEL)
-  run_onnx = OnnxRunner(onnx_file)
-
-  inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
-  out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
-  root = out.uop
-  targets = [x.uop for x in inputs.values()]
-  print(targets)
-
-  # TODO: abstract this from gradient?
-
-  # compute the target path (top down)
-  in_target_path: dict[UOp, bool] = {}
-  for u in root.toposort(): in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
-  independent_set = {}
-  for u in root.toposort():
-    if in_target_path[u]:
-      for s in u.src:
-        if not in_target_path[s]:
-          independent_set[s] = None
-  independent = UOp.sink(*independent_set.keys())
-  kernelized = get_rangeify_map(independent)
-  independent = independent.substitute(kernelized)
-  schedule, var_vals = create_schedule_with_vars(independent)
-  run_schedule(schedule)
-
-  print("**** real ****")
-  GlobalCounters.reset()
-  out.uop = root.substitute(kernelized)
-  out.kernelize()
-
-  # realize
-  out.realize()
@@ -269,7 +269,8 @@ if __name__ == "__main__":

  # load in weights
  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], verbose=False, strict=False, realize=False)
+    model_bin = fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt')
+    load_state_dict(model, torch_load(model_bin)['state_dict'], verbose=False, strict=False, realize=False)

    if args.fp16:
      for k,v in get_state_dict(model).items():
@@ -19,8 +19,8 @@ from tinygrad.helpers import fetch, getenv

 # QUANT=1 python3 examples/test_onnx_imagenet.py
 #   https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx
-# DONT_REALIZE_EXPAND=1 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
-# VIZ=1 DONT_REALIZE_EXPAND=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx
+# python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+# VIZ=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx

 def imagenet_dataloader(cnt=0):
  input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
@@ -0,0 +1,23 @@
+#!/bin/sh
+
+if [ "$#" -ne 1 ] || ! [ -d $1 ]; then
+  echo "usage: $0 MESA_PREFIX"
+  exit 1
+fi
+
+TMP=$(mktemp)
+trap 'rm -f "$TMP"' EXIT
+
+(
+  cat <<EOF
+#define HAVE_ENDIAN_H
+#define HAVE_STRUCT_TIMESPEC
+#define HAVE_PTHREAD
+#include <unistd.h>
+#include "nir_shader_compiler_options.h"
+#include "compiler/shader_enums.h"
+EOF
+  sed -n '/struct nir_shader_compiler_options/,/^}/{p;/^}/q}' $1/src/gallium/drivers/llvmpipe/lp_screen.c
+  echo "int main(void) { write(1, &gallivm_nir_options, sizeof(gallivm_nir_options)); }"
+) | cc -x c -o $TMP - -I$1/src/compiler/nir -I$1/src -I$1/include && $TMP | gzip | base64 -w0
+
@@ -3,8 +3,21 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.nn.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort
+ort_options = ort.SessionOptions()
+ort_options.log_severity_level = 3

 def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
+  """
+  Generate example input tensors based on the provided ONNX graph input specifications.
+
+  NOTE: This is not guaranteed to be reliable. It's a best-effort helper
+  that uses heuristics to guess input shapes and values.
+
+  Example:
+    from tinygrad.nn.onnx import OnnxRunner
+    from extra.onnx_helpers import get_example_inputs
+    inputs = get_example_inputs(OnnxRunner(model_path).graph_inputs)
+  """
  def _get_shape(onnx_shape: tuple[str|int]):
    shape = []
    for onnx_dim in onnx_shape:
@@ -44,11 +57,9 @@ def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
    ret.update({name:value})
  return ret

-def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
+def _get_tinygrad_and_ort_np_outputs(onnx_file, inputs):
  run_onnx = OnnxRunner(onnx_file)

-  ort_options = ort.SessionOptions()
-  ort_options.log_severity_level = 3
  ort_sess = ort.InferenceSession(onnx_file, ort_options, ["CPUExecutionProvider"])
  np_inputs = {k:v.numpy() if isinstance(v, Tensor) else v for k,v in inputs.items()}
  out_names = list(run_onnx.graph_outputs)
@@ -56,9 +67,101 @@ def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
  ort_out = dict(zip(out_names, out_values))

  tinygrad_out = run_onnx(inputs)
+  Tensor.realize(*(x for x in tinygrad_out.values() if x is not None))
+  tinygrad_out = {k:v.numpy() if v is not None else None for k,v in tinygrad_out.items()}
+  return tinygrad_out, ort_out
+
+def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
+  """
+  Compares the final output tensors of an onnx model run in tinygrad and onnxruntime.
+  """
+  tinygrad_out, ort_out = _get_tinygrad_and_ort_np_outputs(onnx_file, inputs)

  assert tinygrad_out.keys() == ort_out.keys()
  for k in tinygrad_out.keys():
    tiny_v, onnx_v = tinygrad_out[k], ort_out[k]
    if tiny_v is None: assert onnx_v is None, f"{k}: {tiny_v=}, {onnx_v=}"
-    else: np.testing.assert_allclose(tiny_v.numpy(), onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tinygrad_out.keys()}")
+    else: np.testing.assert_allclose(tiny_v, onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tinygrad_out.keys()}")
+
+def validate_all_intermediates(onnx_file, inputs, rtol=1e-5, atol=1e-5):
+  """
+  Compares all intermediate node output of an onnx model run in tinygrad and onnxruntime.
+  """
+  report = generate_node_output_report(onnx_file, inputs)
+  for i, node in enumerate(report):
+    node_name = node["node"]
+    op = node["op"]
+    outputs = node["outputs"]
+    for output in outputs:
+      output_name = output["name"]
+      tinygrad_out = output["tinygrad"]
+      ort_out = output["onnxruntime"]
+      try:
+        if tinygrad_out is None: assert ort_out is None, f"None outputs are not equal {tinygrad_out=} {ort_out=}"
+        else: np.testing.assert_allclose(tinygrad_out, ort_out, rtol=rtol, atol=atol)
+        print(f"Validated {i}: {op=} {node_name=} {output_name=}")
+      except AssertionError as e:
+        print(f"FAILED {i}: {op=} {node_name=} {output_name=}")
+        print(str(e).strip() + "\n")
+
+def generate_node_output_report(onnx_file, inputs):
+  """
+  Build a report of all ONNX node outputs from tinygrad and onnxruntime
+
+  Returns:
+    A list of dictionaries, where each entry corresponds to one
+    node in the ONNX graph. The structure is as follows:
+    [
+      {
+        "node": str,  # The name of the ONNX node.
+        "op": str,    # The operation type of the ONNX node.
+        "outputs": [
+          {
+            "name": str,                       # The name of the output tensor.
+            "tinygrad": np.ndarray | None,     # The output value from tinygrad.
+            "onnxruntime": np.ndarray | None,  # The output value from onnxruntime.
+          },
+          ...
+        ]
+      },
+      ...
+    ]
+  """
+  import onnx_graphsurgeon as gs
+  import onnx
+  import tempfile
+
+  # rewrite the model to output all the node outputs
+  # `infer_shapes` here tries to fill the shapes and dtypes of intermediate values which graphsurgeon requires when assigning them as outputs
+  inferred_model = onnx.shape_inference.infer_shapes(onnx.load(onnx_file))
+  model = gs.import_onnx(inferred_model)
+  model_nodes = model.nodes
+  node_outputs = [n.outputs for n in model.nodes]
+  model.outputs = [
+      each_output for outputs in node_outputs for each_output in outputs
+      if not (each_output.dtype is None and each_output.shape is None)  # output with None dtype and None shape is likely a `None` value
+  ]
+  rewritten_model = gs.export_onnx(model)
+
+  # TODO: remove this once ORT supports 1.18.0
+  if getattr(rewritten_model, "ir_version", 0) > 10:
+    rewritten_model.ir_version = 10
+
+  with tempfile.NamedTemporaryFile(suffix=".onnx") as f:
+    onnx.save(rewritten_model, f.name)
+    rewritten_model_path = f.name
+    tinygrad_out, ort_out = _get_tinygrad_and_ort_np_outputs(rewritten_model_path, inputs)
+
+  report = []
+  for node in model_nodes:
+    outputs = []
+    for each_output in node.outputs:
+      if each_output.dtype is None and each_output.shape is None:
+        continue
+      name = each_output.name
+      tinygrad_output = tinygrad_out[name]
+      ort_output = ort_out[name]
+      outputs.append({"name": name, "tinygrad": tinygrad_output, "onnxruntime": ort_output})
+    report.append({"node": node.name, "op": node.op, "outputs": outputs})
+
+  return report
@@ -0,0 +1,11 @@
+from tinygrad.tensor import Tensor
+import argparse
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("hash", type=str, required=True, help="file hash to fetch")
+  parser.add_argument("len", type=int, required=True, help="file length to fetch")
+  parser.add_argument("dest", type=str, required=True, help="destination path to save the file")
+  args = parser.parse_args()
+
+  Tensor(bytes.fromhex(args.hash), device="CPU").load(args.len).to(f"disk:{args.dest}").realize()
@@ -0,0 +1,39 @@
+import json, multiprocessing
+from pathlib import Path
+
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import tqdm, getenv
+
+raid_root = Path(getenv("RAID_ROOT", "/raid"))
+
+def fetch_file(item):
+  path, info = item
+  h, size = info["hash"], info["size"]
+
+  path = raid_root / Path(path)
+  path.parent.mkdir(parents=True, exist_ok=True)
+
+  try:
+    pt = Tensor(bytes.fromhex(h), device="CPU").load(size).to(f"disk:{path.as_posix()}").realize()
+  except Exception as e:
+    print(f"error fetching {path}, {h}, {size}: {e}")
+    raise
+
+  pt.uop.buffer.deallocate()
+
+def fetch_mapping():
+  mapping_tensor = Tensor(bytes.fromhex("d734f5e3be9f1e9d863bfaa4fc6c1ef2")).load(175866113).realize()
+  mapping = mapping_tensor.data().tobytes().decode()
+  mapping = json.loads(mapping)
+  mapped_files = mapping.items()
+  return list(mapped_files)
+
+if __name__ == "__main__":
+  with multiprocessing.Pool(processes=1) as pool:
+    mapped_files = pool.apply(fetch_mapping)
+
+  print(f"fetched mapping for {len(mapped_files)} files")
+
+  with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
+    for _ in tqdm(pool.imap_unordered(fetch_file, mapped_files), total=len(mapped_files)):
+      pass
@@ -0,0 +1,31 @@
+from pathlib import Path
+import multiprocessing, json
+
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import tqdm
+
+raid_root = Path("/raid")
+
+def upload_file(path: Path):
+  pt = Tensor(path).realize()
+  h = pt.store().realize()
+  pt.uop.realized.deallocate()
+  return h.data().hex(), path, pt.nbytes()
+
+if __name__ == "__main__":
+  raid_files = sorted([p for p in raid_root.rglob("*") if p.is_file()])
+  print(f"found {len(raid_files)} files in /raid")
+
+  mapping = {}
+  with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
+    for h, p, s in tqdm(pool.imap_unordered(upload_file, raid_files), total=len(raid_files)):
+      mapping[p.relative_to(raid_root).as_posix()] = {"hash": h, "size": s}
+
+  # sort the mapping by key
+  mapping = dict(sorted(mapping.items()))
+
+  mapping = json.dumps(mapping).encode()
+  mapping_tensor = Tensor(mapping, device="CPU")
+  h = mapping_tensor.store().realize()
+
+  print(f"final hash: {h.data().hex()}, size: {len(mapping)}")
@@ -155,16 +155,14 @@ def index_tensor(x, y):
 def zero_(x):
  if TORCH_DEBUG: print(f"zero_ {x.shape}")
  tt = unwrap(x)
-  # NOTE: unconditional contiguous covers if x is contiguous (match it) or if x is view (realize for inplace)
-  # TODO: consolidate
-  tt.assign(tt.zeros_like().contiguous())
+  tt.assign(tt.zeros_like())

@torch.library.impl("aten::fill_.Scalar", "privateuseone")
@inplace_fn("x")
 def fill_scalar(x, y):
  if TORCH_DEBUG: print(f"fill_.Scalar {x.shape} {y}")
  tt = unwrap(x)
-  tt.assign(tt.full_like(y).contiguous())
+  tt.assign(tt.full_like(y))

@torch.library.impl("aten::_local_scalar_dense", "privateuseone")
 def _local_scalar_dense(tensor): return unwrap(tensor).item()
@@ -644,10 +642,11 @@ def get_real_tinygrad_buffers():
 torch.nn.modules.module.register_module_buffer_registration_hook(register_torch_buffer)

 from torch.nn.modules import Module
-def backward_hook(model:Module, _grad_input, _grad_out):
-  grads_to_realize = [unwrap(p.grad) for p in model.parameters() if p.grad is not None]
-  if len(grads_to_realize): Tensor.realize(*grads_to_realize)
-def module_hook(module:Module, _name, _submodule): module.register_backward_hook(backward_hook)
+def param_hook(_grad):
+  if _grad is not None and _grad.is_tiny: Tensor.realize(unwrap(_grad))
+def module_hook(module:Module, _name, _submodule):
+  for param in _submodule.parameters(recurse=False):
+    if param.requires_grad: param.register_hook(param_hook)
 torch.nn.modules.module.register_module_module_registration_hook(module_hook)

 def realize_optimizer_step(optimizer: torch.optim.Optimizer, *args, **kwargs):
@@ -0,0 +1,13 @@
+xcuserdata/
+
+**/*.xcodeproj/project.xcworkspace/*
+!**/*.xcodeproj/project.xcworkspace/xcshareddata
+
+**/*.xcodeproj/project.xcworkspace/xcshareddata/*
+!**/*.xcodeproj/project.xcworkspace/xcshareddata/WorkspaceSettings.xcsettings
+
+**/*.playground/playground.xcworkspace/*
+!**/*.playground/playground.xcworkspace/xcshareddata
+
+**/*.playground/playground.xcworkspace/xcshareddata/*
+!**/*.playground/playground.xcworkspace/xcshareddata/WorkspaceSettings.xcsettings
@@ -0,0 +1,11 @@
+{
+  "colors" : [
+    {
+      "idiom" : "universal"
+    }
+  ],
+  "info" : {
+    "author" : "xcode",
+    "version" : 1
+  }
+}
@@ -0,0 +1,148 @@
+{
+  "images" : [
+    {
+      "idiom" : "iphone",
+      "scale" : "2x",
+      "size" : "20x20"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "3x",
+      "size" : "20x20"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "2x",
+      "size" : "29x29"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "3x",
+      "size" : "29x29"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "2x",
+      "size" : "40x40"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "3x",
+      "size" : "40x40"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "2x",
+      "size" : "60x60"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "3x",
+      "size" : "60x60"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "1x",
+      "size" : "20x20"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "20x20"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "1x",
+      "size" : "29x29"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "29x29"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "1x",
+      "size" : "40x40"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "40x40"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "1x",
+      "size" : "76x76"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "76x76"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "83.5x83.5"
+    },
+    {
+      "idiom" : "ios-marketing",
+      "scale" : "1x",
+      "size" : "1024x1024"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "16x16"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "16x16"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "32x32"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "32x32"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "128x128"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "128x128"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "256x256"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "256x256"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "512x512"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "512x512"
+    }
+  ],
+  "info" : {
+    "author" : "xcode",
+    "version" : 1
+  }
+}
@@ -0,0 +1,6 @@
+{
+  "info" : {
+    "author" : "xcode",
+    "version" : 1
+  }
+}
@@ -0,0 +1,19 @@
+import AppKit
+import SwiftUI
+
+final class AppDelegate: NSObject, NSApplicationDelegate {
+	func applicationShouldTerminateAfterLastWindowClosed(_ sender: NSApplication) -> Bool {
+		true
+	}
+}
+
+@main
+struct TinyGPUApp: App {
+	@NSApplicationDelegateAdaptor(AppDelegate.self) var appDelegate
+
+    var body: some Scene {
+        WindowGroup {
+			TinyGPUView()
+        }
+    }
+}
@@ -0,0 +1,33 @@
+import SwiftUI
+
+struct TinyGPUView: View {
+	@ObservedObject var viewModel = TinyGPUViewModel()
+
+	var body: some View {
+#if os(macOS)
+		VStack(alignment: .center) {
+			Text("TinyGPU Intsaller")
+				.padding()
+				.font(.title)
+			Text(self.viewModel.dextLoadingState)
+				.multilineTextAlignment(.center)
+			HStack {
+				Button(
+					action: {
+						self.viewModel.activateMyDext()
+					}, label: {
+						Text("Install extension")
+					}
+				)
+			}
+		}
+		.frame(width: 500, height: 200, alignment: .center)
+#endif
+	}
+}
+
+struct TinyGPUView_Previews: PreviewProvider {
+    static var previews: some View {
+		TinyGPUView()
+    }
+}
@@ -0,0 +1,149 @@
+import Foundation
+import os.log
+import SystemExtensions
+
+class TinyGPUDriverLoadingStateMachine {
+  enum State { case unloaded, activating, needsApproval, activated, activationError }
+}
+
+class TinyGPUViewModel: NSObject {
+
+	@Published private var state: TinyGPUDriverLoadingStateMachine.State = .unloaded
+	
+	override init() {
+		super.init()
+		refreshInitialDextState()
+	}
+	
+	private func refreshInitialDextState() {
+#if os(macOS)
+		Task.detached { [dextIdentifier] in
+		  let newState = Self.queryDextState(bundleID: dextIdentifier)
+		  await MainActor.run { self.state = newState }
+		}
+#endif
+	}
+
+#if os(macOS)
+	private static func queryDextState(bundleID: String) -> TinyGPUDriverLoadingStateMachine.State {
+		let tool = "/usr/bin/systemextensionsctl"
+		let p = Process()
+		p.executableURL = URL(fileURLWithPath: tool)
+		p.arguments = ["list"]
+
+		let pipe = Pipe()
+		p.standardOutput = pipe
+		p.standardError = Pipe()
+
+		do {
+			try p.run()
+			p.waitUntilExit()
+			let data = pipe.fileHandleForReading.readDataToEndOfFile()
+			guard let output = String(data: data, encoding: .utf8) else { return .unloaded }
+
+			// Look for our bundle id line
+			if let line = output.split(separator: "\n").first(where: { $0.contains(bundleID) }) {
+				if line.contains("[activated enabled]") { return .activated }
+				if line.contains("[activated waiting for user]") { return .needsApproval }
+				if line.contains("terminated waiting to uninstall") { return .unloaded }
+				return .activating
+			} else {
+				return .unloaded
+			}
+		} catch {
+			return .unloaded
+		}
+	}
+#endif
+
+	private let dextIdentifier: String = "org.tinygrad.tinygpu.edriver"
+
+	public var dextLoadingState: String {
+		switch state {
+		case .unloaded:
+			return "TinyGPUDriver isn't loaded."
+		case .activating:
+			return "Activating TinyGPUDriver, please wait."
+		case .needsApproval:
+			return "Please follow the prompt to approve TinyGPUDriver."
+		case .activated:
+			return "TinyGPUDriver has been activated and is ready to use. You can close the installer."
+		case .activationError:
+			return "TinyGPUDriver has experienced an error during activation.\nPlease check the logs to find the error."
+		}
+	}
+}
+
+extension TinyGPUViewModel: ObservableObject {
+
+#if os(macOS)
+	func activateMyDext() {
+		activateExtension(dextIdentifier)
+	}
+	
+	func deactivateMyDext() {
+		deactivateExtension(dextIdentifier)
+	}
+
+	func activateExtension(_ dextIdentifier: String) {
+
+		let request = OSSystemExtensionRequest
+			.activationRequest(forExtensionWithIdentifier: dextIdentifier,
+							   queue: .main)
+		request.delegate = self
+		OSSystemExtensionManager.shared.submitRequest(request)
+
+		self.state = .activating
+	}
+	
+	func deactivateExtension(_ dextIdentifier: String) {
+
+		let request = OSSystemExtensionRequest.deactivationRequest(forExtensionWithIdentifier: dextIdentifier, queue: .main)
+		request.delegate = self
+		OSSystemExtensionManager.shared.submitRequest(request)
+
+		self.state = .unloaded
+	}
+#endif
+}
+
+#if os(macOS)
+extension TinyGPUViewModel: OSSystemExtensionRequestDelegate {
+
+	func request(
+		_ request: OSSystemExtensionRequest,
+		actionForReplacingExtension existing: OSSystemExtensionProperties,
+		withExtension ext: OSSystemExtensionProperties) -> OSSystemExtensionRequest.ReplacementAction {
+
+		var replacementAction: OSSystemExtensionRequest.ReplacementAction
+
+		os_log("sysex actionForReplacingExtension: %@ %@", existing, ext)
+
+		// Add appropriate logic here to determine whether to replace the extension
+		// with the new extension. Common things to check for include
+		// testing whether the new extension's version number is newer than
+		// the current version number, or whether the bundleIdentifier is different.
+		// For simplicity, this sample always replaces the current extension
+		// with the new one.
+		replacementAction = .replace
+
+		self.state = .activating
+		return replacementAction
+	}
+
+	func requestNeedsUserApproval(_ request: OSSystemExtensionRequest) {
+		os_log("sysex requestNeedsUserApproval")
+		self.state = .needsApproval
+	}
+
+	func request(_ request: OSSystemExtensionRequest, didFinishWithResult result: OSSystemExtensionRequest.Result) {
+		os_log("sysex didFinishWithResult: %d", result.rawValue)
+		self.state = .activated
+	}
+
+	func request(_ request: OSSystemExtensionRequest, didFailWithError error: Error) {
+		os_log("sysex didFailWithError: %@", error.localizedDescription)
+		self.state = .activationError
+	}
+}
+#endif
@@ -0,0 +1,589 @@
+// !$*UTF8*$!
+{
+	archiveVersion = 1;
+	classes = {
+	};
+	objectVersion = 54;
+	objects = {
+
+/* Begin PBXBuildFile section */
+		0ACB55392E9CB880007029EF /* PCIDriverKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 0ACB55382E9CB880007029EF /* PCIDriverKit.framework */; };
+		54798269286A3512009785F6 /* CoreAudio.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 54798268286A3512009785F6 /* CoreAudio.framework */; };
+		549EB121286A1A37009D38AB /* TinyGPUViewModel.swift in Sources */ = {isa = PBXBuildFile; fileRef = 549EB11F286A1A37009D38AB /* TinyGPUViewModel.swift */; };
+		549EB123286A1D48009D38AB /* org.tinygrad.tinygpu.edriver.dext in Embed System Extensions */ = {isa = PBXBuildFile; fileRef = C5B7D9BC26128AC50089B4C3 /* org.tinygrad.tinygpu.edriver.dext */; settings = {ATTRIBUTES = (RemoveHeadersOnCopy, ); }; };
+		549EB131286A2B98009D38AB /* IOKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 549EB130286A2B98009D38AB /* IOKit.framework */; };
+		54E42BC8286A1697000E1E9A /* TinyGPUApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = 54E42BB8286A1696000E1E9A /* TinyGPUApp.swift */; };
+		54E42BCA286A1697000E1E9A /* TinyGPUView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 54E42BB9286A1696000E1E9A /* TinyGPUView.swift */; };
+		54E42BCC286A1697000E1E9A /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = 54E42BBA286A1697000E1E9A /* Assets.xcassets */; };
+		C5B7D9C326128AC50089B4C3 /* TinyGPUDriver.cpp in Sources */ = {isa = PBXBuildFile; fileRef = C5B7D9C226128AC50089B4C3 /* TinyGPUDriver.cpp */; };
+		C5B7D9C526128AC50089B4C3 /* TinyGPUDriver.iig in Sources */ = {isa = PBXBuildFile; fileRef = C5B7D9C426128AC50089B4C3 /* TinyGPUDriver.iig */; };
+		C5C3BBB32612ACDC003C7BFE /* AudioDriverKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = C5C3BBB12612ACD3003C7BFE /* AudioDriverKit.framework */; };
+		C5C3BBB52612ACEF003C7BFE /* DriverKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = C5C3BBB42612ACEF003C7BFE /* DriverKit.framework */; };
+		C5D787AC261667FC006047E5 /* TinyGPUDriverUserClient.iig in Sources */ = {isa = PBXBuildFile; fileRef = C5D787AB261667FC006047E5 /* TinyGPUDriverUserClient.iig */; };
+		C5D787AE26168E59006047E5 /* TinyGPUDriverUserClient.cpp in Sources */ = {isa = PBXBuildFile; fileRef = C5D787AD26168D1E006047E5 /* TinyGPUDriverUserClient.cpp */; };
+/* End PBXBuildFile section */
+
+/* Begin PBXContainerItemProxy section */
+		549EB126286A1D66009D38AB /* PBXContainerItemProxy */ = {
+			isa = PBXContainerItemProxy;
+			containerPortal = C5B7D9B326128AC50089B4C3 /* Project object */;
+			proxyType = 1;
+			remoteGlobalIDString = C5B7D9BB26128AC50089B4C3;
+			remoteInfo = SimpleAudioDriver;
+		};
+/* End PBXContainerItemProxy section */
+
+/* Begin PBXCopyFilesBuildPhase section */
+		549EB122286A1D3A009D38AB /* Embed System Extensions */ = {
+			isa = PBXCopyFilesBuildPhase;
+			buildActionMask = 2147483647;
+			dstPath = "$(SYSTEM_EXTENSIONS_FOLDER_PATH)";
+			dstSubfolderSpec = 16;
+			files = (
+				549EB123286A1D48009D38AB /* org.tinygrad.tinygpu.edriver.dext in Embed System Extensions */,
+			);
+			name = "Embed System Extensions";
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXCopyFilesBuildPhase section */
+
+/* Begin PBXFileReference section */
+		0ACB55382E9CB880007029EF /* PCIDriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = PCIDriverKit.framework; path = System/DriverKit/System/Library/Frameworks/PCIDriverKit.framework; sourceTree = SDKROOT; };
+		54798268286A3512009785F6 /* CoreAudio.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreAudio.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX13.0.sdk/System/Library/Frameworks/CoreAudio.framework; sourceTree = DEVELOPER_DIR; };
+		549EB11F286A1A37009D38AB /* TinyGPUViewModel.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TinyGPUViewModel.swift; sourceTree = "<group>"; usesTabs = 1; };
+		549EB130286A2B98009D38AB /* IOKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = IOKit.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX13.0.sdk/System/Library/Frameworks/IOKit.framework; sourceTree = DEVELOPER_DIR; };
+		549EB132286A2B9D009D38AB /* IOKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = IOKit.framework; path = Platforms/iPhoneOS.platform/Developer/SDKs/iPhoneOS16.0.sdk/System/Library/Frameworks/IOKit.framework; sourceTree = DEVELOPER_DIR; };
+		54E42BB8286A1696000E1E9A /* TinyGPUApp.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TinyGPUApp.swift; sourceTree = "<group>"; };
+		54E42BB9286A1696000E1E9A /* TinyGPUView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TinyGPUView.swift; sourceTree = "<group>"; };
+		54E42BBA286A1697000E1E9A /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = "<group>"; };
+		54E42BC4286A1697000E1E9A /* TinyGPU.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = TinyGPU.app; sourceTree = BUILT_PRODUCTS_DIR; };
+		54E42BC6286A1697000E1E9A /* macOS.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = macOS.entitlements; sourceTree = "<group>"; };
+		C5B7D9BC26128AC50089B4C3 /* org.tinygrad.tinygpu.edriver.dext */ = {isa = PBXFileReference; explicitFileType = "wrapper.driver-extension"; includeInIndex = 0; path = org.tinygrad.tinygpu.edriver.dext; sourceTree = BUILT_PRODUCTS_DIR; };
+		C5B7D9BF26128AC50089B4C3 /* DriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = DriverKit.framework; path = Library/Frameworks/DriverKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5B7D9C226128AC50089B4C3 /* TinyGPUDriver.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; path = TinyGPUDriver.cpp; sourceTree = "<group>"; usesTabs = 1; };
+		C5B7D9C426128AC50089B4C3 /* TinyGPUDriver.iig */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.iig; path = TinyGPUDriver.iig; sourceTree = "<group>"; };
+		C5B7D9C626128AC50089B4C3 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = "<group>"; };
+		C5B7D9CC26128ADA0089B4C3 /* AudioDriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = AudioDriverKit.framework; path = System/DriverKit/System/Library/Frameworks/AudioDriverKit.framework; sourceTree = SDKROOT; };
+		C5B7D9CE26128B150089B4C3 /* TinyGPUDriver.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = TinyGPUDriver.entitlements; sourceTree = "<group>"; };
+		C5C0063326178F98003345D8 /* AppKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = AppKit.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/AppKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5C006352617ACB8003345D8 /* CoreAudio.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreAudio.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/CoreAudio.framework; sourceTree = DEVELOPER_DIR; };
+		C5C3BBB12612ACD3003C7BFE /* AudioDriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = AudioDriverKit.framework; path = Platforms/DriverKit.platform/Developer/SDKs/DriverKit.MacOSX21.0.Internal.sdk/System/DriverKit/System/Library/Frameworks/AudioDriverKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5C3BBB42612ACEF003C7BFE /* DriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = DriverKit.framework; path = Platforms/DriverKit.platform/Developer/SDKs/DriverKit.MacOSX21.0.Internal.sdk/System/DriverKit/System/Library/Frameworks/DriverKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5D787AB261667FC006047E5 /* TinyGPUDriverUserClient.iig */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.iig; path = TinyGPUDriverUserClient.iig; sourceTree = "<group>"; };
+		C5D787AD26168D1E006047E5 /* TinyGPUDriverUserClient.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; path = TinyGPUDriverUserClient.cpp; sourceTree = "<group>"; };
+		C5D787B026169723006047E5 /* IOKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = IOKit.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/IOKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5D787B22616973F006047E5 /* SystemExtensions.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = SystemExtensions.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/SystemExtensions.framework; sourceTree = DEVELOPER_DIR; };
+		C5D787B426169747006047E5 /* Foundation.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = Foundation.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/Foundation.framework; sourceTree = DEVELOPER_DIR; };
+/* End PBXFileReference section */
+
+/* Begin PBXFrameworksBuildPhase section */
+		54E42BC1286A1697000E1E9A /* Frameworks */ = {
+			isa = PBXFrameworksBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				54798269286A3512009785F6 /* CoreAudio.framework in Frameworks */,
+				549EB131286A2B98009D38AB /* IOKit.framework in Frameworks */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+		C5B7D9B926128AC50089B4C3 /* Frameworks */ = {
+			isa = PBXFrameworksBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				C5C3BBB32612ACDC003C7BFE /* AudioDriverKit.framework in Frameworks */,
+				C5C3BBB52612ACEF003C7BFE /* DriverKit.framework in Frameworks */,
+				0ACB55392E9CB880007029EF /* PCIDriverKit.framework in Frameworks */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXFrameworksBuildPhase section */
+
+/* Begin PBXGroup section */
+		54E42BB7286A1696000E1E9A /* Shared */ = {
+			isa = PBXGroup;
+			children = (
+				54E42BB8286A1696000E1E9A /* TinyGPUApp.swift */,
+				54E42BB9286A1696000E1E9A /* TinyGPUView.swift */,
+				549EB11F286A1A37009D38AB /* TinyGPUViewModel.swift */,
+				54E42BBA286A1697000E1E9A /* Assets.xcassets */,
+			);
+			path = Shared;
+			sourceTree = "<group>";
+		};
+		54E42BC5286A1697000E1E9A /* macOS */ = {
+			isa = PBXGroup;
+			children = (
+				54E42BC6286A1697000E1E9A /* macOS.entitlements */,
+			);
+			path = macOS;
+			sourceTree = "<group>";
+		};
+		C5B7D9B226128AC50089B4C3 = {
+			isa = PBXGroup;
+			children = (
+				C5B7D9C126128AC50089B4C3 /* TinyGPUDriverExtension */,
+				54E42BB7286A1696000E1E9A /* Shared */,
+				54E42BC5286A1697000E1E9A /* macOS */,
+				C5B7D9BE26128AC50089B4C3 /* Frameworks */,
+				C5B7D9BD26128AC50089B4C3 /* Products */,
+			);
+			sourceTree = "<group>";
+			usesTabs = 1;
+		};
+		C5B7D9BD26128AC50089B4C3 /* Products */ = {
+			isa = PBXGroup;
+			children = (
+				C5B7D9BC26128AC50089B4C3 /* org.tinygrad.tinygpu.edriver.dext */,
+				54E42BC4286A1697000E1E9A /* TinyGPU.app */,
+			);
+			name = Products;
+			sourceTree = "<group>";
+		};
+		C5B7D9BE26128AC50089B4C3 /* Frameworks */ = {
+			isa = PBXGroup;
+			children = (
+				0ACB55382E9CB880007029EF /* PCIDriverKit.framework */,
+				54798268286A3512009785F6 /* CoreAudio.framework */,
+				549EB130286A2B98009D38AB /* IOKit.framework */,
+				549EB132286A2B9D009D38AB /* IOKit.framework */,
+				C5C006352617ACB8003345D8 /* CoreAudio.framework */,
+				C5C0063326178F98003345D8 /* AppKit.framework */,
+				C5D787B426169747006047E5 /* Foundation.framework */,
+				C5D787B22616973F006047E5 /* SystemExtensions.framework */,
+				C5D787B026169723006047E5 /* IOKit.framework */,
+				C5C3BBB42612ACEF003C7BFE /* DriverKit.framework */,
+				C5B7D9CC26128ADA0089B4C3 /* AudioDriverKit.framework */,
+				C5C3BBB12612ACD3003C7BFE /* AudioDriverKit.framework */,
+				C5B7D9BF26128AC50089B4C3 /* DriverKit.framework */,
+			);
+			name = Frameworks;
+			sourceTree = "<group>";
+		};
+		C5B7D9C126128AC50089B4C3 /* TinyGPUDriverExtension */ = {
+			isa = PBXGroup;
+			children = (
+				C5B7D9C226128AC50089B4C3 /* TinyGPUDriver.cpp */,
+				C5B7D9C426128AC50089B4C3 /* TinyGPUDriver.iig */,
+				C5D787AD26168D1E006047E5 /* TinyGPUDriverUserClient.cpp */,
+				C5D787AB261667FC006047E5 /* TinyGPUDriverUserClient.iig */,
+				C5B7D9C626128AC50089B4C3 /* Info.plist */,
+				C5B7D9CE26128B150089B4C3 /* TinyGPUDriver.entitlements */,
+			);
+			path = TinyGPUDriverExtension;
+			sourceTree = "<group>";
+		};
+/* End PBXGroup section */
+
+/* Begin PBXHeadersBuildPhase section */
+		C5B7D9B726128AC50089B4C3 /* Headers */ = {
+			isa = PBXHeadersBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXHeadersBuildPhase section */
+
+/* Begin PBXNativeTarget section */
+		54E42BC3286A1697000E1E9A /* TinyGPU Installer (macOS) */ = {
+			isa = PBXNativeTarget;
+			buildConfigurationList = 54E42BD2286A1697000E1E9A /* Build configuration list for PBXNativeTarget "TinyGPU Installer (macOS)" */;
+			buildPhases = (
+				54E42BC0286A1697000E1E9A /* Sources */,
+				54E42BC1286A1697000E1E9A /* Frameworks */,
+				54E42BC2286A1697000E1E9A /* Resources */,
+				549EB122286A1D3A009D38AB /* Embed System Extensions */,
+			);
+			buildRules = (
+			);
+			dependencies = (
+				549EB127286A1D66009D38AB /* PBXTargetDependency */,
+			);
+			name = "TinyGPU Installer (macOS)";
+			productName = "SimpleAudioDriverExtension2 (macOS)";
+			productReference = 54E42BC4286A1697000E1E9A /* TinyGPU.app */;
+			productType = "com.apple.product-type.application";
+		};
+		C5B7D9BB26128AC50089B4C3 /* TinyGPUDriver */ = {
+			isa = PBXNativeTarget;
+			buildConfigurationList = C5B7D9C926128AC50089B4C3 /* Build configuration list for PBXNativeTarget "TinyGPUDriver" */;
+			buildPhases = (
+				C5B7D9B726128AC50089B4C3 /* Headers */,
+				C5B7D9B826128AC50089B4C3 /* Sources */,
+				C5B7D9B926128AC50089B4C3 /* Frameworks */,
+				C5B7D9BA26128AC50089B4C3 /* Resources */,
+			);
+			buildRules = (
+			);
+			dependencies = (
+			);
+			name = TinyGPUDriver;
+			productName = SimpleAudioDriverExtension;
+			productReference = C5B7D9BC26128AC50089B4C3 /* org.tinygrad.tinygpu.edriver.dext */;
+			productType = "com.apple.product-type.driver-extension";
+		};
+/* End PBXNativeTarget section */
+
+/* Begin PBXProject section */
+		C5B7D9B326128AC50089B4C3 /* Project object */ = {
+			isa = PBXProject;
+			attributes = {
+				BuildIndependentTargetsInParallel = YES;
+				DefaultBuildSystemTypeForWorkspace = Latest;
+				LastSwiftUpdateCheck = 1400;
+				LastUpgradeCheck = 1600;
+				ORGANIZATIONNAME = Apple;
+				TargetAttributes = {
+					54E42BC3286A1697000E1E9A = {
+						CreatedOnToolsVersion = 14.0;
+						LastSwiftMigration = 1400;
+					};
+					C5B7D9BB26128AC50089B4C3 = {
+						CreatedOnToolsVersion = 13.0;
+					};
+				};
+			};
+			buildConfigurationList = C5B7D9B626128AC50089B4C3 /* Build configuration list for PBXProject "TinyGPUDriverExtension" */;
+			compatibilityVersion = "Xcode 9.3";
+			developmentRegion = en;
+			hasScannedForEncodings = 0;
+			knownRegions = (
+				en,
+				Base,
+			);
+			mainGroup = C5B7D9B226128AC50089B4C3;
+			productRefGroup = C5B7D9BD26128AC50089B4C3 /* Products */;
+			projectDirPath = "";
+			projectRoot = "";
+			targets = (
+				C5B7D9BB26128AC50089B4C3 /* TinyGPUDriver */,
+				54E42BC3286A1697000E1E9A /* TinyGPU Installer (macOS) */,
+			);
+		};
+/* End PBXProject section */
+
+/* Begin PBXResourcesBuildPhase section */
+		54E42BC2286A1697000E1E9A /* Resources */ = {
+			isa = PBXResourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				54E42BCC286A1697000E1E9A /* Assets.xcassets in Resources */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+		C5B7D9BA26128AC50089B4C3 /* Resources */ = {
+			isa = PBXResourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXResourcesBuildPhase section */
+
+/* Begin PBXSourcesBuildPhase section */
+		54E42BC0286A1697000E1E9A /* Sources */ = {
+			isa = PBXSourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				549EB121286A1A37009D38AB /* TinyGPUViewModel.swift in Sources */,
+				54E42BCA286A1697000E1E9A /* TinyGPUView.swift in Sources */,
+				54E42BC8286A1697000E1E9A /* TinyGPUApp.swift in Sources */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+		C5B7D9B826128AC50089B4C3 /* Sources */ = {
+			isa = PBXSourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				C5B7D9C526128AC50089B4C3 /* TinyGPUDriver.iig in Sources */,
+				C5D787AE26168E59006047E5 /* TinyGPUDriverUserClient.cpp in Sources */,
+				C5D787AC261667FC006047E5 /* TinyGPUDriverUserClient.iig in Sources */,
+				C5B7D9C326128AC50089B4C3 /* TinyGPUDriver.cpp in Sources */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXSourcesBuildPhase section */
+
+/* Begin PBXTargetDependency section */
+		549EB127286A1D66009D38AB /* PBXTargetDependency */ = {
+			isa = PBXTargetDependency;
+			target = C5B7D9BB26128AC50089B4C3 /* TinyGPUDriver */;
+			targetProxy = 549EB126286A1D66009D38AB /* PBXContainerItemProxy */;
+		};
+/* End PBXTargetDependency section */
+
+/* Begin XCBuildConfiguration section */
+		54E42BCF286A1697000E1E9A /* Debug */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
+				ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
+				CLANG_ENABLE_MODULES = YES;
+				CODE_SIGN_ENTITLEMENTS = macOS/macOS.entitlements;
+				CODE_SIGN_IDENTITY = "-";
+				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "Apple Development";
+				CODE_SIGN_STYLE = Automatic;
+				COMBINE_HIDPI_IMAGES = YES;
+				CURRENT_PROJECT_VERSION = 1;
+				DEAD_CODE_STRIPPING = YES;
+				DEVELOPMENT_TEAM = 9YG3G8543N;
+				ENABLE_HARDENED_RUNTIME = YES;
+				ENABLE_PREVIEWS = YES;
+				ENABLE_USER_SCRIPT_SANDBOXING = YES;
+				GENERATE_INFOPLIST_FILE = YES;
+				INFOPLIST_KEY_NSHumanReadableCopyright = "";
+				LD_RUNPATH_SEARCH_PATHS = (
+					"$(inherited)",
+					"@executable_path/../Frameworks",
+				);
+				MACOSX_DEPLOYMENT_TARGET = 12.1;
+				MARKETING_VERSION = 1.0;
+				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.installer;
+				PRODUCT_NAME = TinyGPU;
+				PROVISIONING_PROFILE_SPECIFIER = "";
+				SDKROOT = macosx;
+				SWIFT_ACTIVE_COMPILATION_CONDITIONS = DEBUG;
+				SWIFT_EMIT_LOC_STRINGS = YES;
+				SWIFT_OPTIMIZATION_LEVEL = "-Onone";
+				SWIFT_VERSION = 5.0;
+			};
+			name = Debug;
+		};
+		54E42BD0286A1697000E1E9A /* Release */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
+				ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
+				CLANG_ENABLE_MODULES = YES;
+				CODE_SIGN_ENTITLEMENTS = macOS/macOS.entitlements;
+				CODE_SIGN_IDENTITY = "-";
+				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "Apple Development";
+				CODE_SIGN_STYLE = Automatic;
+				COMBINE_HIDPI_IMAGES = YES;
+				CURRENT_PROJECT_VERSION = 1;
+				DEAD_CODE_STRIPPING = YES;
+				DEVELOPMENT_TEAM = 9YG3G8543N;
+				ENABLE_HARDENED_RUNTIME = YES;
+				ENABLE_PREVIEWS = YES;
+				ENABLE_USER_SCRIPT_SANDBOXING = YES;
+				GENERATE_INFOPLIST_FILE = YES;
+				INFOPLIST_KEY_NSHumanReadableCopyright = "";
+				LD_RUNPATH_SEARCH_PATHS = (
+					"$(inherited)",
+					"@executable_path/../Frameworks",
+				);
+				MACOSX_DEPLOYMENT_TARGET = 12.1;
+				MARKETING_VERSION = 1.0;
+				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.installer;
+				PRODUCT_NAME = TinyGPU;
+				PROVISIONING_PROFILE_SPECIFIER = "";
+				SDKROOT = macosx;
+				SWIFT_COMPILATION_MODE = wholemodule;
+				SWIFT_EMIT_LOC_STRINGS = YES;
+				SWIFT_OPTIMIZATION_LEVEL = "-O";
+				SWIFT_VERSION = 5.0;
+			};
+			name = Release;
+		};
+		C5B7D9C726128AC50089B4C3 /* Debug */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ALWAYS_SEARCH_USER_PATHS = NO;
+				CLANG_ANALYZER_NONNULL = YES;
+				CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
+				CLANG_CXX_LANGUAGE_STANDARD = "gnu++17";
+				CLANG_CXX_LIBRARY = "libc++";
+				CLANG_ENABLE_MODULES = YES;
+				CLANG_ENABLE_OBJC_ARC = YES;
+				CLANG_ENABLE_OBJC_WEAK = YES;
+				CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
+				CLANG_WARN_BOOL_CONVERSION = YES;
+				CLANG_WARN_COMMA = YES;
+				CLANG_WARN_CONSTANT_CONVERSION = YES;
+				CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
+				CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
+				CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
+				CLANG_WARN_EMPTY_BODY = YES;
+				CLANG_WARN_ENUM_CONVERSION = YES;
+				CLANG_WARN_INFINITE_RECURSION = YES;
+				CLANG_WARN_INT_CONVERSION = YES;
+				CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
+				CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
+				CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
+				CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
+				CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
+				CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
+				CLANG_WARN_STRICT_PROTOTYPES = YES;
+				CLANG_WARN_SUSPICIOUS_MOVE = YES;
+				CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
+				CLANG_WARN_UNREACHABLE_CODE = YES;
+				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
+				COPY_PHASE_STRIP = NO;
+				DEBUG_INFORMATION_FORMAT = dwarf;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				ENABLE_STRICT_OBJC_MSGSEND = YES;
+				ENABLE_TESTABILITY = YES;
+				GCC_C_LANGUAGE_STANDARD = gnu11;
+				GCC_DYNAMIC_NO_PIC = NO;
+				GCC_NO_COMMON_BLOCKS = YES;
+				GCC_OPTIMIZATION_LEVEL = 0;
+				GCC_PREPROCESSOR_DEFINITIONS = (
+					"DEBUG=1",
+					"$(inherited)",
+				);
+				GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
+				GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
+				GCC_WARN_UNDECLARED_SELECTOR = YES;
+				GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
+				GCC_WARN_UNUSED_FUNCTION = YES;
+				GCC_WARN_UNUSED_VARIABLE = YES;
+				MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;
+				MTL_FAST_MATH = YES;
+				ONLY_ACTIVE_ARCH = YES;
+				SDKROOT = driverkit;
+			};
+			name = Debug;
+		};
+		C5B7D9C826128AC50089B4C3 /* Release */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ALWAYS_SEARCH_USER_PATHS = NO;
+				CLANG_ANALYZER_NONNULL = YES;
+				CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
+				CLANG_CXX_LANGUAGE_STANDARD = "gnu++17";
+				CLANG_CXX_LIBRARY = "libc++";
+				CLANG_ENABLE_MODULES = YES;
+				CLANG_ENABLE_OBJC_ARC = YES;
+				CLANG_ENABLE_OBJC_WEAK = YES;
+				CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
+				CLANG_WARN_BOOL_CONVERSION = YES;
+				CLANG_WARN_COMMA = YES;
+				CLANG_WARN_CONSTANT_CONVERSION = YES;
+				CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
+				CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
+				CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
+				CLANG_WARN_EMPTY_BODY = YES;
+				CLANG_WARN_ENUM_CONVERSION = YES;
+				CLANG_WARN_INFINITE_RECURSION = YES;
+				CLANG_WARN_INT_CONVERSION = YES;
+				CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
+				CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
+				CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
+				CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
+				CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
+				CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
+				CLANG_WARN_STRICT_PROTOTYPES = YES;
+				CLANG_WARN_SUSPICIOUS_MOVE = YES;
+				CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
+				CLANG_WARN_UNREACHABLE_CODE = YES;
+				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
+				COPY_PHASE_STRIP = NO;
+				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				ENABLE_NS_ASSERTIONS = NO;
+				ENABLE_STRICT_OBJC_MSGSEND = YES;
+				GCC_C_LANGUAGE_STANDARD = gnu11;
+				GCC_NO_COMMON_BLOCKS = YES;
+				GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
+				GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
+				GCC_WARN_UNDECLARED_SELECTOR = YES;
+				GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
+				GCC_WARN_UNUSED_FUNCTION = YES;
+				GCC_WARN_UNUSED_VARIABLE = YES;
+				MTL_ENABLE_DEBUG_INFO = NO;
+				MTL_FAST_MATH = YES;
+				SDKROOT = driverkit;
+				SWIFT_COMPILATION_MODE = wholemodule;
+			};
+			name = Release;
+		};
+		C5B7D9CA26128AC50089B4C3 /* Debug */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				AD_HOC_CODE_SIGNING_ALLOWED = YES;
+				CODE_SIGN_ENTITLEMENTS = TinyGPUDriverExtension/TinyGPUDriver.entitlements;
+				CODE_SIGN_IDENTITY = "Apple Development";
+				CODE_SIGN_STYLE = Automatic;
+				CURRENT_PROJECT_VERSION = 1;
+				DEVELOPMENT_TEAM = 9YG3G8543N;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				ENABLE_USER_SCRIPT_SANDBOXING = YES;
+				FRAMEWORK_SEARCH_PATHS = (
+					"$(inherited)",
+					"$(SDKROOT)/System/DriverKit/System/Library/Frameworks",
+				);
+				GENERATE_INFOPLIST_FILE = YES;
+				INFOPLIST_FILE = TinyGPUDriverExtension/Info.plist;
+				INFOPLIST_KEY_OSBundleUsageDescription = "Sample Code Audio Driver Kit Extension";
+				MARKETING_VERSION = 1.0;
+				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.edriver;
+				PRODUCT_NAME = "$(inherited)";
+				PROVISIONING_PROFILE_SPECIFIER = "";
+				RUN_CLANG_STATIC_ANALYZER = YES;
+				SDKROOT = driverkit;
+				SKIP_INSTALL = YES;
+			};
+			name = Debug;
+		};
+		C5B7D9CB26128AC50089B4C3 /* Release */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				AD_HOC_CODE_SIGNING_ALLOWED = YES;
+				CODE_SIGN_ENTITLEMENTS = TinyGPUDriverExtension/TinyGPUDriver.entitlements;
+				CODE_SIGN_IDENTITY = "Apple Development";
+				CODE_SIGN_STYLE = Automatic;
+				CURRENT_PROJECT_VERSION = 1;
+				DEVELOPMENT_TEAM = 9YG3G8543N;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				ENABLE_USER_SCRIPT_SANDBOXING = YES;
+				FRAMEWORK_SEARCH_PATHS = (
+					"$(inherited)",
+					"$(SDKROOT)/System/DriverKit/System/Library/Frameworks",
+				);
+				GENERATE_INFOPLIST_FILE = YES;
+				INFOPLIST_FILE = TinyGPUDriverExtension/Info.plist;
+				INFOPLIST_KEY_OSBundleUsageDescription = "Sample Code Audio Driver Kit Extension";
+				MARKETING_VERSION = 1.0;
+				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.edriver;
+				PRODUCT_NAME = "$(inherited)";
+				PROVISIONING_PROFILE_SPECIFIER = "";
+				RUN_CLANG_STATIC_ANALYZER = YES;
+				SDKROOT = driverkit;
+				SKIP_INSTALL = YES;
+			};
+			name = Release;
+		};
+/* End XCBuildConfiguration section */
+
+/* Begin XCConfigurationList section */
+		54E42BD2286A1697000E1E9A /* Build configuration list for PBXNativeTarget "TinyGPU Installer (macOS)" */ = {
+			isa = XCConfigurationList;
+			buildConfigurations = (
+				54E42BCF286A1697000E1E9A /* Debug */,
+				54E42BD0286A1697000E1E9A /* Release */,
+			);
+			defaultConfigurationIsVisible = 0;
+			defaultConfigurationName = Release;
+		};
+		C5B7D9B626128AC50089B4C3 /* Build configuration list for PBXProject "TinyGPUDriverExtension" */ = {
+			isa = XCConfigurationList;
+			buildConfigurations = (
+				C5B7D9C726128AC50089B4C3 /* Debug */,
+				C5B7D9C826128AC50089B4C3 /* Release */,
+			);
+			defaultConfigurationIsVisible = 0;
+			defaultConfigurationName = Release;
+		};
+		C5B7D9C926128AC50089B4C3 /* Build configuration list for PBXNativeTarget "TinyGPUDriver" */ = {
+			isa = XCConfigurationList;
+			buildConfigurations = (
+				C5B7D9CA26128AC50089B4C3 /* Debug */,
+				C5B7D9CB26128AC50089B4C3 /* Release */,
+			);
+			defaultConfigurationIsVisible = 0;
+			defaultConfigurationName = Release;
+		};
+/* End XCConfigurationList section */
+	};
+	rootObject = C5B7D9B326128AC50089B4C3 /* Project object */;
+}
@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>BuildSystemType</key>
+	<string>Latest</string>
+	<key>DerivedDataLocationStyle</key>
+	<string>Default</string>
+</dict>
+</plist>
@@ -0,0 +1,37 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>IOKitPersonalities</key>
+	<dict>
+		<key>TinyGPUDriver</key>
+		<dict>
+			<key>CFBundleIdentifier</key>
+			<string>$(PRODUCT_BUNDLE_IDENTIFIER)</string>
+			<key>IOClass</key>
+			<string>IOUserService</string>
+			<key>IOMatchCategory</key>
+			<string>TinyGPUDriver</string>
+			<key>IOPCIClassMatch</key>
+			<string>0x03000000</string>
+			<key>IOPCITunnelCompatible</key>
+			<true/>
+			<key>IOProviderClass</key>
+			<string>IOPCIDevice</string>
+			<key>IOResourceMatch</key>
+			<string>IOKit</string>
+			<key>IOUserClass</key>
+			<string>TinyGPUDriver</string>
+			<key>IOUserServerName</key>
+			<string>org.tinygrad.tinygpu.Driver</string>
+			<key>TinyGPUDriverUserClientProperties</key>
+			<dict>
+				<key>IOClass</key>
+				<string>IOUserUserClient</string>
+				<key>IOUserClass</key>
+				<string>TinyGPUDriverUserClient</string>
+			</dict>
+		</dict>
+	</dict>
+</dict>
+</plist>
@@ -0,0 +1,259 @@
+#include "TinyGPUDriver.h"
+#include "TinyGPUDriverUserClient.h"
+#include <AudioDriverKit/AudioDriverKit.h>
+#include <DriverKit/IOUserServer.h>
+#include <DriverKit/IOLib.h>
+#include <DriverKit/OSString.h>
+#include <DriverKit/IOMemoryMap.h>
+#include <DriverKit/IODMACommand.h>
+#include <DriverKit/IODispatchQueue.h>
+#include <PCIDriverKit/PCIDriverKit.h>
+#include <DriverKit/OSAction.h>
+
+struct TinyGPUDriver_IVars
+{
+	IOPCIDevice *pci = nullptr;
+};
+
+bool TinyGPUDriver::init()
+{
+	os_log(OS_LOG_DEFAULT, "tinygpu: init");
+
+	auto answer = super::init();
+	if (!answer) {
+		return false;
+	}
+
+	ivars = new TinyGPUDriver_IVars();
+	if (ivars == nullptr) {
+		return false;
+	}
+
+	return true;
+}
+
+void TinyGPUDriver::free()
+{
+	if (ivars != nullptr) {
+		
+	}
+	IOSafeDeleteNULL(ivars, TinyGPUDriver_IVars, 1);
+	super::free();
+}
+
+kern_return_t TinyGPUDriver::Start_Impl(IOService* in_provider)
+{
+	IOServiceName service_name;
+	os_log(OS_LOG_DEFAULT, "tinygpu: on gpu detected");
+
+	kern_return_t err = Start(in_provider, SUPERDISPATCH);
+	if (err) return err;
+
+	ivars->pci = OSDynamicCast(IOPCIDevice, in_provider);
+	if (!ivars->pci) return kIOReturnNoDevice;
+
+	err = ivars->pci->Open(this, 0);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: Open() failed 0x%08x", err);
+		ivars->pci = nullptr;
+		return err;
+	}
+
+	uint16_t ven = 0, dev = 0;
+	ivars->pci->ConfigurationRead16(kIOPCIConfigurationOffsetVendorID, &ven);
+	ivars->pci->ConfigurationRead16(kIOPCIConfigurationOffsetDeviceID, &dev);
+	os_log(OS_LOG_DEFAULT, "tinygpu: opened device ven=0x%04x dev=0x%04x", ven, dev);
+
+#if 0
+	uint32_t off = 0x100;
+	while (off) {
+		uint32_t hdr = 0, next = 0, cap_id = 0;
+		ivars->pci->ConfigurationRead32(off, &hdr);
+		cap_id = hdr & 0xFFFFu;
+		next = (hdr >> 20) & 0xFFCu;
+		os_log(OS_LOG_DEFAULT, "tinygpu: cap: %u", cap_id);
+		if (cap_id == 0x15) {
+			uint32_t cap = 0, ctrl = 0;
+			ivars->pci->ConfigurationRead32(off+0x4, &cap);
+			ivars->pci->ConfigurationRead32(off+0x8, &ctrl);
+
+			uint32_t new_bar_size = 31 - __builtin_clz(cap >> 4);
+			uint32_t new_ctrl = (ctrl & ~0x1f00) | (new_bar_size << 8);
+			ivars->pci->ConfigurationWrite32(off+0x8, new_ctrl);
+
+			os_log(OS_LOG_DEFAULT, "tinygpu: rebar: cap=%u ctrl=%u new_bar_size=%u new_ctrl=%u", cap, ctrl, new_bar_size, new_ctrl);
+			ivars->pci->Reset(0);
+			break;
+		}
+		off = next;
+	}
+	ivars->pci->Reset(kIOPCIDeviceResetTypeHotReset);
+#endif
+
+	uint16_t commandRegister;
+	ivars->pci->ConfigurationRead16(kIOPCIConfigurationOffsetCommand, &commandRegister);
+	commandRegister |= (kIOPCICommandIOSpace | kIOPCICommandBusMaster | kIOPCICommandMemorySpace);
+	ivars->pci->ConfigurationWrite16(kIOPCIConfigurationOffsetCommand, commandRegister);
+
+	memcpy((void*)service_name, (void*)"tinygpu\0", 8);
+	SetName(service_name);
+
+	os_log(OS_LOG_DEFAULT, "tinygpu: will register service %s", service_name);
+	RegisterService();
+
+	os_log(OS_LOG_DEFAULT, "tinygpu: service started %s", service_name);
+	return 0;
+}
+
+kern_return_t TinyGPUDriver::Stop_Impl(IOService* in_provider)
+{
+	ivars->pci->Close(this, 0);
+	return 0;
+}
+
+kern_return_t TinyGPUDriver::NewUserClient_Impl(uint32_t in_type, IOUserClient** out_user_client)
+{
+	kern_return_t err = 0;
+
+	IOService* user_client_service = nullptr;
+	err = Create(this, "TinyGPUDriverUserClientProperties", &user_client_service);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to create NewUserClient");
+		goto error;
+	}
+	*out_user_client = OSDynamicCast(IOUserClient, user_client_service);
+	os_log(OS_LOG_DEFAULT, "tinygpu: NewUserClient created");
+
+error:
+	return err;
+}
+
+kern_return_t TinyGPUDriver::MapBar(uint32_t bar, IOMemoryDescriptor** memory)
+{
+	kern_return_t err = 0;
+	uint8_t barMemoryIndex, barMemoryType;
+	uint64_t barMemorySize;
+	err = ivars->pci->GetBARInfo(bar, &barMemoryIndex, &barMemorySize, &barMemoryType);
+	if (err) return err;
+
+	os_log(OS_LOG_DEFAULT, "tinygpu: requested bar mapping %d, %d", bar, (uint32_t)barMemoryIndex);
+	err = ivars->pci->_CopyDeviceMemoryWithIndex(barMemoryIndex, memory, this);
+	return err;
+}
+
+kern_return_t TinyGPUDriver::CreateDMA(size_t size, TinyGPUCreateDMAResp* dmaDesc)
+{
+	kern_return_t err = 0;
+	IOMemoryMap* memoryMap = nullptr;
+	IOBufferMemoryDescriptor* sharedBuf = nullptr;
+	IODMACommand* dmaCmd = nullptr;
+	uint64_t flags = kIOMemoryDirectionInOut;
+	uint32_t segCount = 32;
+	IOAddressSegment segments[32];
+	IODMACommandSpecification dmaSpec = {
+		.options = 0,
+		.maxAddressBits = 40,
+	};
+
+	err = IOBufferMemoryDescriptor::Create(kIOMemoryDirectionInOut, size, IOVMPageSize, &sharedBuf);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to alloc user buffer, err=%d", err);
+		goto error;
+	}
+	
+	err = IODMACommand::Create(ivars->pci, kIODMACommandCreateNoOptions, &dmaSpec, &dmaCmd);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to create dma command, err=%d", err);
+		goto error;
+	}
+
+	err = dmaCmd->PrepareForDMA(kIODMACommandPrepareForDMANoOptions, sharedBuf, 0, size,
+								&flags, &segCount, segments);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to prepare for dma, err=%d", err);
+		goto error;
+	}
+
+	// pass addresses to userland
+	{
+		// debug
+		for (int i = 0; i < segCount; i++) {
+			os_log(OS_LOG_DEFAULT, "tinygpu: new dma mapping (sz=0x%zx) %d 0x%llx 0x%llx", size, i, segments[i].address, segments[i].length);
+		}
+
+		err = sharedBuf->CreateMapping(0, 0, 0, IOVMPageSize, IOVMPageSize, &memoryMap); // one page should be fine
+		if (err) {
+			os_log(OS_LOG_DEFAULT, "tinygpu: failed to map memory, err=%d", err);
+			goto error;
+		}
+
+		// Send back gpu addresses
+		uint64_t* addr = (uint64_t*)memoryMap->GetAddress();
+		for (int i = 0; i < segCount; i++) {
+			addr[i * 2] = segments[i].address;
+			addr[i * 2 + 1] = segments[i].length;
+		}
+		addr[segCount * 2] = 0;
+		addr[segCount * 2 + 1] = 0;
+
+		// free memoryMap
+		memoryMap->release();
+		memoryMap = nullptr;
+	}
+
+	dmaDesc->sharedBuf = sharedBuf;
+	dmaDesc->dmaCmd = dmaCmd;
+	return 0;
+
+error:
+	if (memoryMap) {
+		memoryMap->release();
+		memoryMap = nullptr;
+	}
+	if (dmaCmd) {
+		dmaCmd->CompleteDMA(kIODMACommandCompleteDMANoOptions);
+		dmaCmd->release();
+		dmaCmd = nullptr;
+	}
+	if (sharedBuf) {
+		sharedBuf->release();
+		sharedBuf = nullptr;
+	}
+	return err;
+}
+
+kern_return_t TinyGPUDriver::CfgRead(uint32_t off, uint32_t size, uint32_t* outVal)
+{
+  if (!ivars->pci || !outVal) return kIOReturnNotReady;
+
+  if (size == 1) {
+	uint8_t v8 = 0;
+	ivars->pci->ConfigurationRead8(off, &v8);
+	*outVal = v8;
+  } else if (size == 2) {
+	uint16_t v16 = 0;
+	ivars->pci->ConfigurationRead16(off, &v16);
+	*outVal = v16;
+  } else if (size == 4) {
+	uint32_t v32 = 0;
+	ivars->pci->ConfigurationRead32(off, &v32);
+	*outVal = v32;
+  }
+  return 0;
+}
+
+kern_return_t TinyGPUDriver::CfgWrite(uint32_t off, uint32_t size, uint32_t val)
+{
+  if (!ivars->pci) return kIOReturnNotReady;
+  if (size == 1) ivars->pci->ConfigurationWrite8 (off, (uint8_t)val);
+  else if (size == 2) ivars->pci->ConfigurationWrite16(off, (uint16_t)val);
+  else if (size == 4) ivars->pci->ConfigurationWrite32(off, (uint32_t)val);
+  return 0;
+}
+
+kern_return_t TinyGPUDriver::ResetDevice()
+{
+	if (!ivars->pci) return kIOReturnNotReady;
+	ivars->pci->Reset(kIOPCIDeviceResetTypeFunctionReset);
+	return 0;
+}
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>com.apple.developer.driverkit.transport.pci</key>
+	<array>
+		<dict>
+			<key>IOPCIMatch</key>
+			<string>0x00001002&amp;0x0000FFFF</string>
+		</dict>
+		<dict>
+			<key>IOPCIMatch</key>
+			<string>0x000010de&amp;0x0000FFFF</string>
+		</dict>
+	</array>
+	<key>com.apple.developer.driverkit.allow-any-userclient-access</key>
+	<true/>
+	<key>com.apple.developer.driverkit</key>
+	<true/>
+</dict>
+</plist>
@@ -0,0 +1,38 @@
+#ifndef TinyGPUDriver_h
+#define TinyGPUDriver_h
+
+#include <Availability.h>
+#include <DriverKit/IOService.iig>
+#include <PCIDriverKit/IOPCIDevice.iig>
+#include <DriverKit/IOMemoryMap.iig>
+#include <DriverKit/IODMACommand.iig>
+
+struct TinyGPUCreateDMAResp
+{
+	IOBufferMemoryDescriptor* sharedBuf;
+	IODMACommand* dmaCmd;
+};
+
+class TinyGPUDriver: public IOService
+{
+public:
+	virtual bool init() override;
+
+	virtual void free() override;
+
+	virtual kern_return_t Start(IOService * provider) override;
+
+	virtual kern_return_t Stop(IOService * provider) override;
+
+	virtual kern_return_t NewUserClient(uint32_t in_type, IOUserClient** out_user_client) override;
+
+	kern_return_t MapBar(uint32_t bar, IOMemoryDescriptor** memory) LOCALONLY;
+	kern_return_t CreateDMA(size_t size, TinyGPUCreateDMAResp* dmaDesc) LOCALONLY;
+
+	kern_return_t CfgRead(uint32_t off, uint32_t size, uint32_t* val) LOCALONLY;
+	kern_return_t CfgWrite(uint32_t off, uint32_t size, uint32_t val) LOCALONLY;
+	kern_return_t ResetDevice() LOCALONLY;
+	kern_return_t BarInfo() LOCALONLY;
+};
+
+#endif /* TinyGPUDriver_h */
@@ -0,0 +1,127 @@
+#include "TinyGPUDriverUserClient.h"
+#include "TinyGPUDriver.h"
+#include <DriverKit/DriverKit.h>
+#include <DriverKit/OSSharedPtr.h>
+#include <PCIDriverKit/PCIDriverKit.h>
+
+struct TinyGPUDriverUserClient_IVars
+{
+	OSSharedPtr<TinyGPUDriver> provider = nullptr;
+};
+
+bool TinyGPUDriverUserClient::init()
+{
+	auto theAnswer = super::init();
+	if (!theAnswer) {
+		return false;
+	}
+
+	ivars = IONewZero(TinyGPUDriverUserClient_IVars, 1);
+	if (ivars == nullptr) {
+		return false;
+	}
+
+	return true;
+}
+
+void TinyGPUDriverUserClient::free()
+{
+	if (ivars != nullptr) {
+		ivars->provider.reset();
+	}
+
+	IOSafeDeleteNULL(ivars, TinyGPUDriverUserClient_IVars, 1);
+	super::free();
+}
+
+kern_return_t TinyGPUDriverUserClient::Start_Impl(IOService* in_provider)
+{
+	kern_return_t err = kIOReturnSuccess;
+	if (!in_provider) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: provider is null");
+		err = kIOReturnBadArgument;
+		goto error;
+	}
+
+	err = Start(in_provider, SUPERDISPATCH);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to start super (%d)", err);
+		goto error;
+	}
+
+	ivars->provider = OSSharedPtr(OSDynamicCast(TinyGPUDriver, in_provider), OSRetain);
+	return 0;
+
+error:
+	ivars->provider.reset();
+	return err;
+}
+
+kern_return_t TinyGPUDriverUserClient::Stop_Impl(IOService* in_provider)
+{
+	return Stop(in_provider, SUPERDISPATCH);
+}
+
+kern_return_t TinyGPUDriverUserClient::ExternalMethod(uint64_t selector, IOUserClientMethodArguments* args, const IOUserClientMethodDispatch* in_dispatch, OSObject* in_target, void* in_reference)
+{
+	kern_return_t err = 0;
+
+	os_log(OS_LOG_DEFAULT, "tinygpu: rpc (%llu) in:%d, out:%d", selector, args->scalarInputCount, args->scalarOutputCount);
+
+	if (selector == TinyGPURPC::ReadCfg) {
+		if (args->scalarInputCount != 2 or args->scalarOutputCount < 1) return kIOReturnBadArgument;
+
+		uint32_t off = uint32_t(args->scalarInput[0]);
+		uint32_t size = uint32_t(args->scalarInput[1]);
+
+		uint32_t val = 0;
+		err = ivars->provider->CfgRead(off, size, &val);
+		os_log(OS_LOG_DEFAULT, "tinygpu: read cfg off:%x sz:%d, val:%x", off, size, val);
+
+		if (!err) {
+			args->scalarOutput[0] = val;
+			args->scalarOutputCount = 1;
+		}
+		return err;
+	} else if (selector == TinyGPURPC::WriteCfg) {
+		if (args->scalarInputCount != 3) return kIOReturnBadArgument;
+
+		uint32_t off = uint32_t(args->scalarInput[0]);
+		uint32_t size = uint32_t(args->scalarInput[1]);
+		uint32_t val = uint32_t(args->scalarInput[2]);
+
+		os_log(OS_LOG_DEFAULT, "tinygpu: wr cfg off:%x sz:%d, val:%x", off, size, val);
+		return ivars->provider->CfgWrite(off, size, val);
+	} else if (selector == TinyGPURPC::Reset) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: reset");
+		return ivars->provider->ResetDevice();
+	}
+
+	return kIOReturnUnsupported;
+}
+
+kern_return_t IMPL(TinyGPUDriverUserClient, CopyClientMemoryForType)
+{
+	if (!memory) {
+		return kIOReturnBadArgument;
+	}
+
+	if (ivars->provider.get() == nullptr) {
+		return kIOReturnNotAttached;
+	}
+
+	if (type < 6) {
+		uint32_t bar = (uint32_t)type;
+		return ivars->provider->MapBar(bar, memory);
+	}
+
+	// dma page buffer
+	TinyGPUCreateDMAResp buf;
+	kern_return_t err = ivars->provider->CreateDMA(type, &buf);
+	if (err) {
+		return err;
+	}
+
+	*memory = buf.sharedBuf;
+	return 0;
+}
@@ -0,0 +1,28 @@
+#ifndef TinyGPUDriverUserClient_h
+#define TinyGPUDriverUserClient_h
+
+#include <DriverKit/IOUserClient.iig>
+
+enum TinyGPURPC
+{
+	ReadCfg,
+	WriteCfg,
+	Reset
+};
+
+class TinyGPUDriverUserClient : public IOUserClient
+{
+public:
+	virtual bool init() final;
+	virtual void free() final;
+
+	virtual kern_return_t Start(IOService* in_provider) final;
+	virtual kern_return_t Stop(IOService* in_provider) final;
+
+	virtual kern_return_t ExternalMethod(uint64_t in_selector, IOUserClientMethodArguments* in_arguments, const IOUserClientMethodDispatch* in_dispatch, OSObject* in_target, void* in_reference) final;
+
+	virtual kern_return_t CopyClientMemoryForType(
+		uint64_t type, uint64_t *options, IOMemoryDescriptor **memory) final;
+};
+
+#endif /* TinyGPUDriverUserClient_h */
@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>com.apple.security.app-sandbox</key>
+	<false/>
+	<key>com.apple.security.files.user-selected.read-only</key>
+	<true/>
+	<key>com.apple.developer.system-extension.install</key>
+	<true/>
+</dict>
+</plist>
@@ -0,0 +1,39 @@
+#include <CoreFoundation/CoreFoundation.h>
+#include <IOKit/IOKitLib.h>
+#include <mach/mach.h>
+#include <stdio.h>
+#include <inttypes.h>
+
+static io_connect_t open_uc_by_name(const char *svc_name) {
+  io_connect_t conn = IO_OBJECT_NULL;
+  io_service_t service = IOServiceGetMatchingService(kIOMasterPortDefault, IOServiceNameMatching(svc_name));
+  if (!service) { fprintf(stderr, "service not found: %s\n", svc_name); return IO_OBJECT_NULL; }
+  kern_return_t kr = IOServiceOpen(service, mach_task_self(), /*type*/0, &conn);
+  IOObjectRelease(service);
+  if (kr) { fprintf(stderr, "IOServiceOpen 0x%x\n", kr); return IO_OBJECT_NULL; }
+  return conn;
+}
+
+int main(int argc, char **argv) {
+  uint32_t bar = (argc > 1) ? (uint32_t)strtoul(argv[1], NULL, 0) : 0;  // pick BAR index
+  io_connect_t conn = open_uc_by_name("tinygpu");
+  if (!conn) return 2;
+
+  mach_vm_address_t addr = 0;
+  mach_vm_size_t    size = 0;
+  kern_return_t kr = IOConnectMapMemory64(conn, bar, mach_task_self(), &addr, &size, kIOMapAnywhere);
+  if (kr) { fprintf(stderr, "Map BAR%u failed 0x%x\n", bar, kr); IOServiceClose(conn); return 3; }
+
+  printf("BAR%u mapped at 0x%llx, size 0x%llx\n", bar, (unsigned long long)addr, (unsigned long long)size);
+
+  // example: read a 32-bit register at offset 0x0 (make sure it’s safe!)
+  volatile uint32_t *mmio = (volatile uint32_t*)(uintptr_t)addr;
+  uint32_t v = mmio[0];
+  printf("mmio[0]=0x%08x\n", v);
+
+  kr = IOConnectUnmapMemory64(conn, bar, mach_task_self(), addr);
+  if (kr) fprintf(stderr, "Unmap failed 0x%x\n", kr);
+
+  IOServiceClose(conn);
+  return 0;
+}
@@ -0,0 +1,82 @@
+import ctypes, ctypes.util, sys
+
+cf = ctypes.CDLL(ctypes.util.find_library("CoreFoundation"))
+iokit = ctypes.CDLL(ctypes.util.find_library("IOKit"))
+libsys = ctypes.CDLL(ctypes.util.find_library("System"))
+
+kern_return_t = ctypes.c_int
+mach_port_t = ctypes.c_uint
+io_object_t = mach_port_t
+io_service_t = io_object_t
+io_connect_t = mach_port_t
+CFMutableDictionaryRef = ctypes.c_void_p
+CFStringRef = ctypes.c_void_p
+
+kIOMasterPortDefault = mach_port_t(0)
+
+libsys.mach_task_self_.restype = mach_port_t
+
+iokit.IOServiceNameMatching.argtypes = [ctypes.c_char_p]
+iokit.IOServiceNameMatching.restype  = CFMutableDictionaryRef
+
+iokit.IOServiceGetMatchingService.argtypes = [mach_port_t, CFMutableDictionaryRef]
+iokit.IOServiceGetMatchingService.restype  = io_service_t
+
+iokit.IOObjectRelease.argtypes = [io_object_t]
+iokit.IOObjectRelease.restype  = kern_return_t
+
+iokit.IOServiceOpen.argtypes = [io_service_t, mach_port_t, ctypes.c_uint32, ctypes.POINTER(io_connect_t)]
+iokit.IOServiceOpen.restype  = kern_return_t
+
+iokit.IOConnectCallMethod.argtypes = [io_connect_t, ctypes.c_uint32, ctypes.POINTER(ctypes.c_uint64), ctypes.c_uint32, ctypes.c_void_p,
+  ctypes.c_size_t, ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint32), ctypes.c_void_p, ctypes.POINTER(ctypes.c_size_t)]
+iokit.IOConnectCallMethod.restype = kern_return_t
+
+def open_userclient_by_name(name: str, uc_type: int = 0) -> io_connect_t:
+  mdict = iokit.IOServiceNameMatching(name.encode("utf-8"))
+  if not mdict: raise RuntimeError("IOServiceNameMatching returned NULL")
+
+  # Grab the first matching service
+  service = iokit.IOServiceGetMatchingService(kIOMasterPortDefault, mdict)
+  if not service: raise RuntimeError(f'service "{name}" not found')
+
+  # print("lol", service)
+  # print(libsys.mach_task_self_)
+  # cast libsys.mach_task_self_ to uint and print
+  # print("lol", ctypes.cast(libsys.mach_task_self_, ctypes.POINTER(ctypes.c_uint)).contents.value)
+
+  try:
+    # Open user client (type -> passed to NewUserClient_Impl)
+    conn = io_connect_t(0)
+    # print("lol", libsys.mach_task_self_)
+    kr = iokit.IOServiceOpen(service, ctypes.cast(libsys.mach_task_self_, ctypes.POINTER(ctypes.c_uint)).contents.value,
+                             ctypes.c_uint32(uc_type), ctypes.byref(conn))
+    if kr != 0: raise OSError(kr, f"IOServiceOpen failed (0x{kr:08x})")
+    return conn
+  finally: iokit.IOObjectRelease(service)
+
+def external_method(conn: io_connect_t, selector: int = 0) -> int:
+  # no scalars in/out, no struct in/out — just ping selector 0
+  in_scalars = ctypes.POINTER(ctypes.c_uint64)()  # NULL
+  out_scalars = (ctypes.c_uint64 * 1)()       # space if driver returns something
+  out_scalars_cnt = ctypes.c_uint32(0)      # driver can set this
+
+  return iokit.IOConnectCallMethod(conn, ctypes.c_uint32(selector), in_scalars, ctypes.c_uint32(0), None, ctypes.c_size_t(0),
+    out_scalars, ctypes.byref(out_scalars_cnt), None, ctypes.byref(ctypes.c_size_t(0)))
+
+def close_userclient(conn: io_connect_t) -> None:
+  # IOServiceClose is a macro; exported symbol is IOServiceClose in IOKit
+  iokit.IOServiceClose.argtypes = [io_connect_t]
+  iokit.IOServiceClose.restype  = kern_return_t
+  iokit.IOServiceClose(conn)
+
+if __name__ == "__main__":
+  try:
+    conn = open_userclient_by_name("tinygpu", uc_type=0)
+    kr = external_method(conn, selector=0)
+    print(f"ExternalMethod(0) -> 0x{kr:08x}")
+  except Exception as e:
+    print(e)
+    sys.exit(1)
+  finally:
+    if 'conn' in locals() and conn.value: close_userclient(conn)
@@ -42,7 +42,6 @@ setup(name='tinygrad',
        'tinygrad.runtime.support.am',
        'tinygrad.runtime.support.nv',
        'tinygrad.schedule',
-        'tinygrad.shape',
        'tinygrad.uop',
        'tinygrad.viz',
      ],
@@ -1,4 +1,4 @@
-from tinygrad import Tensor, dtypes, Context, GlobalCounters
+from tinygrad import Tensor, dtypes, GlobalCounters
 dtypes.default_float = dtypes.float16
 from tinygrad.dtype import to_dtype
 from tinygrad.helpers import getenv
@@ -13,6 +13,5 @@ if __name__ == "__main__":

  # test single kernel softmax
  GlobalCounters.reset()
-  with Context(DONT_GROUP_REDUCES=1):
-    single_kernel_softmax(t, -1, acc_dtype).realize()
+  single_kernel_softmax(t, -1, acc_dtype).realize()

@@ -1,46 +0,0 @@
-# ruff: noqa: E501
-from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.dtype import dtypes
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.codegen.opt.search import bufs_from_lin
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-
-ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-  UOp(Ops.STORE, dtypes.void, arg=None, src=(
-    UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=0, src=()),
-    UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(2, 1, 1280, 8, 8, 1, 1, 1), strides=(81920, 0, 64, 8, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),
-    UOp(Ops.ADD, dtypes.half, arg=None, src=(
-      UOp(Ops.ADD, dtypes.half, arg=None, src=(
-        UOp(Ops.CAST, dtypes.half, arg=None, src=(
-          UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (5, 6, 7)), src=(
-            UOp(Ops.CAST, dtypes.float, arg=None, src=(
-              UOp(Ops.MUL, dtypes.half, arg=None, src=(
-                UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-                  UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=1, src=()),
-                  UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 2, 1, 2560, 4, 10, 4, 10), strides=(0, 163840, 0, 64, 0, 8, 0, 1), offset=-9, mask=((0, 1), (0, 2), (0, 1), (0, 2560), (0, 4), (1, 9), (0, 4), (1, 9)), contiguous=False), View(shape=(2, 1, 1280, 8, 8, 2560, 3, 3), strides=(4096000, 0, 0, 40, 1, 1600, 440, 11), offset=0, mask=None, contiguous=False))), src=()),)),
-                UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-                  UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=2, src=()),
-                  UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(2, 1, 1280, 8, 8, 2560, 3, 3), strides=(0, 0, 23040, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),)),
-        UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=3, src=()),
-          x17:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(2, 1, 1280, 8, 8, 1, 1, 1), strides=(0, 0, 1, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
-      UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=4, src=()),
-         x17,)),)),)),))
-opts = [Opt(op=OptOps.UPCAST, axis=3, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UNROLL, axis=2, arg=0), Opt(op=OptOps.UNROLL, axis=1, arg=0), Opt(op=OptOps.LOCAL, axis=1, arg=8), Opt(op=OptOps.LOCAL, axis=2, arg=8), Opt(op=OptOps.LOCAL, axis=2, arg=2)]
-
-k = Kernel(ast)
-k.apply_opts(opts)
-bufs = bufs_from_lin(k)
-
-prg = CompiledRunner(get_program(k.ast, k.opts, k.applied_opts))
-
-for i in range(10):
-  speed = prg(bufs, var_vals={}, wait=True)
-  print(f"kernel time: {speed*1e3:.2f} ms")
-
-# on M1 Max
-# 11ms before block 9b0859d71780fef5cf3831e317f74e53f2483229
-# 15ms after block cbcc1c20eb09a1342f6581cfbb99632bade982a8
@@ -0,0 +1,39 @@
+import subprocess, unittest, os, sys
+from tinygrad.device import Device
+
+class TestTinygradSlow(unittest.TestCase):
+  def test_env_overwrite_default_device(self):
+    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}\\""'],
+                    shell=True, check=True)
+    subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                    shell=True, check=True)
+
+    if Device.DEFAULT != "CPU":
+      # setting multiple devices fail
+      with self.assertRaises(subprocess.CalledProcessError):
+        subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                        shell=True, check=True)
+
+      # setting device via DEV
+      subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+      subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+      subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+
+      with self.assertRaises(subprocess.CalledProcessError):
+        subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                        shell=True, check=True)
+
+class TestRunAsModule(unittest.TestCase):
+  def test_module_runs(self):
+    p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+      env={**os.environ, "DEBUG": "1"}, timeout=40,)
+    out = (p.stdout + p.stderr).decode()
+    self.assertEqual(p.returncode, 0, msg=out)
+
+if __name__ == '__main__':
+  unittest.main()
@@ -58,8 +58,8 @@ class TestExample(unittest.TestCase):
      print(f"WARNING: {device} test isn't running")
      return

-    x = Tensor.eye(64, device=device, requires_grad=True)
-    y = Tensor.eye(64, device=device, requires_grad=True)
+    x = Tensor.eye(8, device=device, requires_grad=True)
+    y = Tensor.eye(8, device=device, requires_grad=True)
    z = y.matmul(x).sum()
    z.backward()

@@ -1,55 +0,0 @@
-# ruff: noqa: E501
-import unittest
-
-from tinygrad.uop.ops import UOp, Ops
-from .search import Opt, OptOps
-from tinygrad.dtype import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from tinygrad.codegen.opt.kernel import Kernel
-
-from test.external.fuzz_linearizer import run_linearizer
-
-class TestTrainGpt2Kernel(unittest.TestCase):
-  def test_1(self):
-    # kernel 244
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(206045184), arg=0, src=()),
-        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 1), strides=(51511296, 50304, 1, 0), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (3,)), src=(
-        UOp(Ops.MUL, dtypes.float, arg=None, src=(
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(3145728), arg=1, src=()),
-          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 768), strides=(786432, 768, 0, 1), offset=0, mask=None, contiguous=False),)), src=()),)),
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(38633472), arg=2, src=()),
-          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 768), strides=(0, 0, 768, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),))
-
-    opts = [Opt(op=OptOps.LOCAL, axis=0, arg=16), Opt(op=OptOps.UPCAST, axis=1, arg=3), Opt(op=OptOps.LOCAL, axis=0, arg=2)]
-    kernel = Kernel(ast)
-    kernel.apply_opts(opts)
-    run_linearizer(kernel)
-
-  def test_2(self):
-    # kernel 254
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(3145728), arg=0, src=()),
-        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 1, 768), strides=(786432, 768, 0, 1), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (2,)), src=(
-        UOp(Ops.MUL, dtypes.float, arg=None, src=(
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(38633472), arg=1, src=()),
-          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 768), strides=(0, 0, 768, 1), offset=0, mask=None, contiguous=False),)), src=()),)),
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(205852672), arg=2, src=()),
-          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 768), strides=(51463168, 50257, 1, 0), offset=0, mask=((0, 4), (0, 1024), (0, 50257), (0, 768)), contiguous=False),)), src=()),)),)),)),)),))
-
-    opts = [Opt(op=OptOps.LOCAL, axis=1, arg=16), Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=3, arg=4)]
-    kernel = Kernel(ast)
-    kernel.apply_opts(opts)
-    run_linearizer(kernel)
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,6 +1,5 @@
 import gc
 from tinygrad import Tensor, UOp, Device, nn
-from tinygrad.shape.shapetracker import views_to_valid_uop
 from tinygrad.engine.realize import method_cache, get_program
 from tinygrad.schedule.indexing import apply_movement_op
 from test.test_tiny import TestTiny
@@ -69,7 +68,6 @@ if __name__ == "__main__":

    # these caches will keep uops alive
    method_cache.clear()
-    views_to_valid_uop.cache_clear()
    apply_movement_op.cache_clear()
    Tensor._device_seeds.clear()
    Tensor._device_rng_counters.clear()
@@ -2,7 +2,6 @@ import unittest
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.engine.realize import get_program
 from tinygrad.helpers import AMX

@@ -149,33 +148,5 @@ class TestFloat4(unittest.TestCase):

    assert TestFloat4.count_float4(uops) == (1, 1)

-  @unittest.skip("Ops.VIEW no longer exists")
-  def test_half4_load_unrolled(self):
-    # from llama 7B shard 4 gpus
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.VIEW, dtypes.float.ptr(96000), arg=ShapeTracker(views=(View(shape=(1, 3, 32000, 1), strides=(0, 32000, 1, 0), offset=0, mask=None, contiguous=True),)), src=( # noqa: E501
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(96000), arg=0, src=()),)),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (3,)), src=(
-          UOp(Ops.CAST, dtypes.float, arg=None, src=(
-            UOp(Ops.MUL, dtypes.half, arg=None, src=(
-              UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-                UOp(Ops.VIEW, dtypes.half.ptr(9216), arg=ShapeTracker(views=(View(shape=(1, 3, 32000, 1024), strides=(0, 4096, 0, 1), offset=0, mask=None, contiguous=False),)), src=( # noqa: E501
-                  UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(9216), arg=1, src=()),)),)),
-              UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-                UOp(Ops.VIEW, dtypes.half.ptr(32768000), arg=ShapeTracker(views=(View(shape=(1, 3, 32000, 1024), strides=(0, 0, 1024, 1), offset=0, mask=None, contiguous=False),)), src=( # noqa: E501
-                  UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(32768000), arg=2, src=()),)),)),)),)),)),)),))
-
-    # TODO: fix this, expected might change but should be positive
-    for expected, opts in [
-      ((7, 0), [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=3), Opt(op=OptOps.UNROLL, axis=0, arg=4)]),
-      ((5, 0), [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UNROLL, axis=0, arg=4)]),
-      ((2, 0), [Opt(op=OptOps.UNROLL, axis=0, arg=4)]),
-    ]:
-      program = get_program(ast, Device[Device.DEFAULT].renderer, opts=opts)
-
-      count = TestFloat4.count_half4(program.uops)
-      assert count == expected, f"{count=}, {expected=}"
-
 if __name__ == '__main__':
  unittest.main()
@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Device, Tensor, dtypes
-from tinygrad.helpers import CI
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError

 # TODO: write a clean version of this
@@ -177,9 +176,7 @@ class TestKernelOpts(unittest.TestCase):
    ], apply_tc=True, atol=atol, rtol=rtol)

  def test_padto_matmul(self):
-    if (CI and Device.DEFAULT in ["AMD", "NV", "CUDA"]):
-      self.skipTest("super slow on CUDA and AMD because of the big grid dims")
-    N = 17 * 17
+    N = 17
    Tensor.manual_seed(289)
    a = Tensor.rand(N, N)
    b = Tensor.rand(N, N)
@@ -213,7 +210,7 @@ class TestKernelOpts(unittest.TestCase):
      helper_linearizer_opt(a@b, [[Opt(OptOps.UNROLL, 0, 0), Opt(OptOps.PADTO, 2, 8)]])

  def test_padto_sum_ok(self):
-    N = 18 * 18
+    N = 18
    # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
    a = Tensor.rand(N, N).realize().shrink(((0, 17), (0, 17))) * 100
    b = (Tensor.rand(N, N) < 0.5).realize().shrink(((0, 17), (0, 17)))
@@ -244,7 +241,7 @@ class TestKernelOpts(unittest.TestCase):
    helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])

  def test_padto_sum_not_ok(self):
-    N = 18 * 18
+    N = 18
    # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
    a = Tensor.rand(N, N).shrink(((0, 17), (0, 17))).exp()
    # exp is not safe to pad
@@ -261,7 +258,7 @@ class TestKernelOpts(unittest.TestCase):
      helper_linearizer_opt(b.sum(0), [[Opt(OptOps.PADTO, 1, 32)],])

  def test_padto_max(self):
-    N = 18 * 18
+    N = 18
    # NOTE: this setup prevents 17 * 17 contiguous merged into one axis
    a = -Tensor.rand(N, N).shrink(((0, 17), (0, 17))) * 100

@@ -282,7 +279,7 @@ class TestKernelOpts(unittest.TestCase):

  def test_padto_where(self):
    Tensor.manual_seed(0)
-    N = 17 * 17
+    N = 17
    a = (Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1).where(1, 0)
    helper_linearizer_opt(a.max(0), [
      [Opt(OptOps.PADTO, 0, 32)],
@@ -291,7 +288,7 @@ class TestKernelOpts(unittest.TestCase):

  def test_padto_where_multioutput(self):
    Tensor.manual_seed(0)
-    N = 17 * 17
+    N = 17
    r = Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1
    a0 = r.where(1, 0)
    a1 = r.where(2, 0)
@@ -131,8 +131,7 @@ class TestIndexing(unittest.TestCase):
    # llama3 is 128256
    vocab_size, embed_size = (10, 3) if CI else (32000, 4096)
    emb = nn.Embedding(vocab_size, embed_size)
-    # TODO: why is a new realize needed here
-    emb_w = emb.weight.realize().numpy()
+    emb_w = emb.weight.numpy()
    x = Tensor([1,2,3,4])
    with Context(NOOPT=noopt):
      GlobalCounters.reset()
@@ -129,6 +129,7 @@ class TestAssign(unittest.TestCase):
  @unittest.expectedFailure
  def test_assign_changes_realized_alt(self): return self.test_assign_changes_alt(realize=True)

+  @unittest.skip("assign to contiguous shouldn't change the base buffer")
  def test_assign_changes_buffer_alt(self):
    a, b = [Tensor(Tensor(0).contiguous().realize().uop.as_buf()) for _ in range(2)]
    Tensor.realize(a.contiguous().assign(1), b.contiguous().assign(2))
@@ -1,7 +1,7 @@
 import unittest, io
 from contextlib import redirect_stdout
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.helpers import OSX, CPU_LLVM
+from tinygrad.helpers import OSX, CPU_LLVM, CPU_LVP
 from tinygrad.engine.realize import lower_schedule
 from tinygrad.device import is_dtype_supported
 from tinygrad.engine.realize import get_program
@@ -19,7 +19,7 @@ class TestCompileFailures(unittest.TestCase):

 class TestDisassembly(unittest.TestCase):
  # TODO: fails on llvm. llvm.LLVMGetHostCPUName() returns "generic"
-  @unittest.skipUnless(Device.DEFAULT in ("CPU",) and not CPU_LLVM and OSX, "m series cpus support fp16 arithmetic")
+  @unittest.skipUnless(Device.DEFAULT in ("CPU",) and not (CPU_LLVM or CPU_LVP) and OSX, "m series cpus support fp16 arithmetic")
  def test_float16_alu(self):
    c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
    s = c.schedule()[-1]
@@ -6,6 +6,7 @@ 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.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
 from test.helpers import rand_for_dtype
@@ -102,7 +103,7 @@ class TestDType(unittest.TestCase):
    ))

  @unittest.skipIf(Device.DEFAULT == "PYTHON", "skip for now")
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "skip for now")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "skip for now")
  def test_uint_overflow(self):
    if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
    v = dtypes.max(self.DTYPE)
@@ -261,7 +262,7 @@ class TestFloatDType(TestDType):
 class TestDoubleDType(TestDType):
  DTYPE = dtypes.double
  @unittest.skipIf((CI and Device.DEFAULT in {"CUDA", "NV"}) or \
-                    isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "conversion not supported on CI CUDA and PTX")  # TODO: why not?
+   isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "conversion not supported on CI CUDA, PTX, and NIR")  # TODO: why not?
  def test_float64_increased_precision(self):
    for func in [
      lambda t: t.exp(),
@@ -6,6 +6,7 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
 from tinygrad.runtime.ops_python import from_storage_scalar
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
 import numpy as np
 import pytest
 from hypothesis import assume, given, strategies as strat, settings, HealthCheck
@@ -29,8 +30,8 @@ unary_operations = [(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.
 # TODO: enable this (this is a dtype issue)
 #binary_operations.append(operator.truediv)

-# TODO: CI CUDA segfaults on sin, WEBGPU sin is not precise enough for large numbers
-if (getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}) or Device.DEFAULT == "WEBGPU":
+# TODO: CI CUDA segfaults on sin, WEBGPU and NIR sines are not precise enough for large numbers
+if (getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}) or Device.DEFAULT == "WEBGPU" or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer):
  unary_operations.remove((Tensor.sin, np.sin))
  unary_operations.remove((Tensor.cos, np.cos))

@@ -184,8 +185,8 @@ class TestDTypeALU(unittest.TestCase):
  @given(ht.int32, ht.int32, ht.float32, strat.sampled_from(integer_binary_operations), strat.sampled_from(binary_operations))
  def test_int32_midcast_float(self, a, b, c, op1, op2): universal_test_midcast(a, b, c, op1, op2, dtypes.int32, dtypes.float32)

-  # Metal and CUDA and HIP behave differently than numpy in CI for overflows
-  skip_overflow = CI and Device.DEFAULT in {"AMD", "NV", "CUDA"}
+  # Metal and CUDA and HIP and NIR behave differently than numpy in CI for overflows
+  skip_overflow = (CI and Device.DEFAULT in {"AMD", "NV", "CUDA"}) or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer)
  @given(strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
         strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
         ht.int32, strat.sampled_from(binary_operations), strat.sampled_from(integer_binary_operations))
@@ -26,8 +26,9 @@ import unittest
 import numpy as np
 import torch
 from tinygrad import Tensor, dtypes, nn
-from tinygrad.device import is_dtype_supported
+from tinygrad.device import Device, is_dtype_supported
 from tinygrad.helpers import getenv
+from tinygrad.renderer.nir import NIRRenderer

 MOCKGPU = getenv("MOCKGPU")

@@ -206,7 +207,8 @@ class TestUOpValidationIssue(unittest.TestCase):
  # these fail with UOp verification error.
  # we want more of these with diverse errors!

-  @unittest.skipIf((not is_dtype_supported(dtypes.long)) or MOCKGPU, "hangs gpuocelot")
+  @unittest.skipIf((not is_dtype_supported(dtypes.long)) or MOCKGPU or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer),
+                   "hangs gpuocelot, NIR cannot render")
  def test_tensor_index_overflow(self):
    val = Tensor([1])
    big = val.expand(2**31 + 3)
@@ -4,12 +4,9 @@

 import unittest
 from tinygrad import Device, dtypes
-from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
-from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.codegen.opt.search import Opt, OptOps
 from tinygrad.engine.realize import get_program
-from tinygrad.renderer.ptx import PTXRenderer

 class TestLinearizerFailure(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
@@ -30,91 +27,5 @@ class TestLinearizerFailure(unittest.TestCase):
    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)

-class TestLinearizerDumb(unittest.TestCase):
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "need local")
-  @unittest.skip("Ops.VALID no longer exists")
-  def test_max_simplify_and_cancel(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1000), arg=0, src=())
-    c1 = c0.view(ShapeTracker(views=(View(shape=(1000, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)))
-    c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1000), arg=1, src=())
-    c3 = c2.view(ShapeTracker(views=(View(shape=(1000, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)))
-    c4 = c3.load()
-    c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1), arg=2, src=())
-    c6 = c5.view(ShapeTracker(views=(View(shape=(1000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)))
-    c7 = c6.load()
-    c8 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c9 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1001, 1999), strides=(0, 0), offset=0, mask=((0, 1001), (999, 1999)), contiguous=False), View(shape=(1000, 1000), strides=(1, 2000), offset=0, mask=None, contiguous=False))), src=())
-    c10 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1000, 1000), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c11 = c1.store((c4.alu(Ops.CMPNE, c7).alu(Ops.CMPNE, UOp.const(dtypes.bool, True, src=c8)).cast(dtypes.int)*(c9.f(Ops.VALID, dtype=dtypes.bool).where(UOp.const(dtypes.int, -1, src=c10), UOp.const(dtypes.int, 0, src=c10)).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (1,)))+UOp.const(dtypes.int, 1000, src=c8))))
-    ast = c11.sink()
-    #opts = [Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8)]
-    opts = [Opt(op=OptOps.LOCAL, axis=0, arg=8)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
-    print(prg.src)
-    assert prg.uops is not None and not any(uop.op is Ops.MAX for uop in prg.uops), "leftover MAX"
-
-  # this was a bug in embedding, someday we should fold this anyway
-  @unittest.skipUnless(is_dtype_supported(dtypes.half), f"half dtype not supported on {Device.DEFAULT}")
-  @unittest.skip("UOp.view is no longer supported")
-  def test_llama_embedding(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(4096), arg=0, src=())
-    c1 = c0.view(ShapeTracker(views=(View(shape=(4096, 1, 1), strides=(1, 0, 0), offset=0, mask=None, contiguous=True),)))
-    c2 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(32001, 63999), strides=(0, 0), offset=0, mask=((0, 32001), (31999, 63999)), contiguous=False), View(shape=(4096, 32000, 32000), strides=(0, 1, 64000), offset=0, mask=None, contiguous=False))), src=())
-    c3 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4096, 32000, 32000), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c4 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=1, src=())
-    c6 = c5.view(ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)))
-    c7 = c6.load()
-    c8 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(131072000), arg=2, src=())
-    c9 = c8.view(ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(1, 4096, 0), offset=0, mask=None, contiguous=False),)))
-    c10 = c9.load()
-    c11 = c1.store(((c2.f(Ops.VALID, dtype=dtypes.bool).where(UOp.const(dtypes.int, 1, src=c3), UOp.const(dtypes.int, 0, src=c3)).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (2,)))+UOp.const(dtypes.int, -1, src=c4)).alu(Ops.CMPNE, c7).alu(Ops.CMPNE, UOp.const(dtypes.bool, True, src=c4)).cast(dtypes.half)*c10).cast(dtypes.float).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (1,))).cast(dtypes.half))
-    ast = c11.sink()
-    prg = get_program(ast, Device[Device.DEFAULT].renderer)
-    print(prg.src)
-
-  @unittest.expectedFailure
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need float4")
-  def test_unrolled_float4_align(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1), arg=0, src=())
-    c1 = c0.view(ShapeTracker(views=(View(shape=(1, 1), strides=(0, 0), offset=0, mask=None, contiguous=True),)))
-    c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(18), arg=1, src=())
-    c3 = c2.view(ShapeTracker(views=(View(shape=(3, 6), strides=(6, 1), offset=0, mask=None, contiguous=True),)))
-    c4 = c3.load()
-    c5 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(3, 6), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c6 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(18), arg=2, src=())
-    c7 = c6.view(ShapeTracker(views=(View(shape=(3, 6), strides=(6, 1), offset=0, mask=None, contiguous=True),)))
-    c8 = c7.load()
-    c9 = c1.store(c4.alu(Ops.CMPNE, UOp.const(dtypes.long, -1, src=c5)).alu(Ops.CMPNE, UOp.const(dtypes.bool, True, src=c5)).where(UOp.const(dtypes.float, 0.0, src=c5), c8).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (0, 1))))
-    ast = c9.sink()
-    opts = [Opt(op=OptOps.UNROLL, axis=0, arg=0)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
-    print(prg.src)
-    load_idxs = [x.src[1] for x in prg.uops if x.op is Ops.LOAD and x.src[0].arg == 2]
-    assert load_idxs[0] < load_idxs[1], f"first loaded idx {load_idxs[0].arg} then {load_idxs[1].arg}!"
-
-  @unittest.expectedFailure
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need float4")
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "this is somehow correct in PTX")
-  def test_upcasted_stores_out_of_order(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(9360), arg=0, src=())
-    c1 = c0.view(ShapeTracker(views=(View(shape=(4, 5, 13, 1, 1, 1, 1, 1, 4, 3, 3), strides=(2340, 468, 36, 0, 0, 0, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=True),)))
-    c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(144), arg=1, src=())
-    c3 = c2.view(ShapeTracker(views=(View(shape=(4, 5, 13, 1, 1, 1, 4, 1, 4, 3, 3), strides=(0, 0, 0, 0, 0, 0, 1, 0, 4, 48, 16), offset=0, mask=None, contiguous=False),)))
-    c4 = c3.load()
-    c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1040), arg=2, src=())
-    c6 = c5.view(ShapeTracker(views=(View(shape=(4, 5, 13, 1, 1, 1, 4, 1, 4, 3, 3), strides=(260, 13, 1, 0, 0, 0, 65, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)))
-    c7 = c6.load()
-    c8 = c1.store((c4*c7).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (6,))))
-    ast = c8.sink()
-    opts = [Opt(op=OptOps.UPCAST, axis=3, arg=0), Opt(op=OptOps.UPCAST, axis=2, arg=0)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
-    print(prg.src)
-    store_idxs = [x.src[1] for x in prg.uops if x.op is Ops.STORE]
-    for i in range(len(store_idxs) - 1):
-      first_bounds = store_idxs[i].vmin+store_idxs[i].vmax
-      next_bounds = store_idxs[i+1].vmin+store_idxs[i+1].vmax
-      assert first_bounds < next_bounds, f"first stored (max) idx {first_bounds} then {next_bounds}!"
-
 if __name__ == '__main__':
  unittest.main()
@@ -658,7 +658,7 @@ class TestMultiTensor(unittest.TestCase):

  # it doesn't work like this anymore
  # NOTE: this never failed in assign_multi, it failed tensor spec because MULTI was never pushed in the graph
-  @unittest.expectedFailure
+  @unittest.skip("this test is broken")
  def test_mlb_assign_change_axis(self):
    t_none = Tensor.zeros((16, 16)).shard(devices_2).contiguous().realize()
    t_zero = Tensor.ones((16, 16)).shard(devices_2, axis=0)
@@ -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, CPU_LLVM, AMD_LLVM
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, CPU_LVP, AMD_LLVM
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -698,8 +698,8 @@ class TestOps(unittest.TestCase):

  def test_pow_zero_tensor(self):
    helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [0.0]])
-    # TODO: fix WEBGPU
-    if Device.DEFAULT != "WEBGPU":
+    # TODO: fix WEBGPU and LVP
+    if Device.DEFAULT != "WEBGPU" and not CPU_LVP:
      helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [0.3]])
      helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [-0.3]])
  def test_pow_zero_const(self):
@@ -830,6 +830,7 @@ class TestOps(unittest.TestCase):
    self.assertEqual(a, b)
    self.assertEqual(Tensor(-1).contiguous().idiv(4).item(), 0)  # NOTE this is trunc-div behaviour

+  @unittest.skipIf(getenv("NV_NAK"), "MUFU.SIN is not accurate enough")
  def test_sin(self):
    helper_test_op([(45,65)], lambda x: x.sin())
    helper_test_op([()], lambda x: x.sin())
@@ -839,6 +840,7 @@ class TestOps(unittest.TestCase):
      helper_test_op(None, lambda x: x.sin(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                    atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and platform.system() == "Windows", "Not accurate enough with DirectX backend")
+  @unittest.skipIf(getenv("NV_NAK"), "MUFU.SIN is not accurate enough")
  def test_cos(self):
    helper_test_op([(45,65)], lambda x: x.cos())
    helper_test_op([()], lambda x: x.cos())
@@ -847,6 +849,7 @@ class TestOps(unittest.TestCase):
      helper_test_op(None, lambda x: x.cos(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                    atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and platform.system() == "Windows", "Not accurate enough with DirectX backend")
+  @unittest.skipIf(getenv("NV_NAK"), "MUFU.SIN is not accurate enough")
  def test_tan(self):
    # NOTE: backward has much higher diff with input close to pi/2 and -pi/2
    helper_test_op([(45,65)], lambda x: x.tan(), low=-1.5, high=1.5)
@@ -3177,6 +3180,7 @@ class TestOps(unittest.TestCase):
  def test_bitcast(self):
    helper_test_op([(3, 3)], lambda x: x.view(torch.int32), lambda x: x.bitcast(dtypes.int32), forward_only=True)

+  @unittest.skip("we have test_linalg, no need to test here. TODO: should be in torch backend tests")
  def test_svd(self):
    # test for tiny backend. real svd tests are in test_linalg
    A = torch.randn(5, 5)
@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, Device
-from tinygrad.helpers import CPU_LLVM
+from tinygrad.helpers import CPU_LLVM, CPU_LVP
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import get_program

@@ -12,7 +12,7 @@ class TestOpts(unittest.TestCase):
    out = (a+b).contiguous(arg=opts)
    s = out.schedule()
    self.assertEqual(s[-1].ast.arg.opts_to_apply, opts)
-    if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM:
+    if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM and not CPU_LVP:
      prg = get_program(s[-1].ast)
      self.assertIn('float4', prg.src)

@@ -217,9 +217,9 @@ class TestProfiler(unittest.TestCase):
        Tensor.realize(a, b)
    profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
    exec_points = [e for e in profile if isinstance(e, ProfilePointEvent) and e.name == "exec"]
-    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent)]
+    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent) and not e.is_copy]
    self.assertEqual(len(exec_points), len(range_events), 2)
-    self.assertEqual(len(dedup(e.key for e in exec_points)), 1)
+    self.assertEqual(len(dedup(e.arg['name'] for e in exec_points)), 1)
    self.assertEqual(len(dedup(e.arg['metadata'] for e in exec_points)), 1)

 if __name__ == "__main__":
@@ -72,7 +72,7 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
    out_file = get_quantized_model(sz)
    run_onnx = OnnxRunner(out_file)
    inp = Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32))
-    with Context(DONT_REALIZE_EXPAND=1, QUANTIZE=1):
+    with Context(QUANTIZE=1):
      sched = run_onnx({"input":inp})["output"].schedule()
      ei = lower_schedule_item(sched[-2])
      daccs = [u for u in ei.prg.p.uops if u.op is Ops.DEFINE_REG]
@@ -86,8 +86,7 @@ class TestQuantizeOnnx(unittest.TestCase):
    # divide is ~1500-2000 without reduce_range, 750-900 with it
    out_file = get_quantized_model(sz)
    run_onnx_jit, _ = load_onnx_model(out_file)
-    with Context(DONT_REALIZE_EXPAND=1):
-      run_onnx_jit(input=Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32)))
+    run_onnx_jit(input=Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32)))

  def test_prequant_conv2d_1x1(self):
    X = Tensor(np.random.uniform(0, 255, size=(1, 32, 128, 128)).astype(np.uint8))
@@ -109,11 +108,10 @@ class TestQuantizeOnnx(unittest.TestCase):
    N = 512
    X = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(xi))
    W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(wi))
-    with Context(DONT_REALIZE_EXPAND=1):
-      # this divide is interesting and forces the accumulator to actually be an int
-      out = (X.cast("int").matmul(W.cast("int"))//1000).cast("int8")
-      opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
-      sexec(out, opts)
+    # this divide is interesting and forces the accumulator to actually be an int
+    out = (X.cast("int").matmul(W.cast("int"))//1000).cast("int8")
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
+    sexec(out, opts)

  def test_prequant_gemm_handcode(self):
    src = """typedef int int128 __attribute__((aligned(512),vector_size(512)));
@@ -203,14 +201,12 @@ class TestQuantizeOnnx(unittest.TestCase):
  def test_prequant_gemm_intacc(self, xi=np.uint8, wi=np.uint8, replace_src=None, N=512, clip=True, opts=None):
    X = Tensor(m1:=(np.random.uniform(0, 255, size=(N,N)).astype(xi))).realize()
    W = Tensor(m2:=(np.random.uniform(0, 255, size=(N,N)).astype(wi))).realize()
-    # ugh, it's so broken with those casts. need DONT_REALIZE_EXPAND=1 python3 test/test_quantize_onnx.py TestQuantizeOnnx.test_prequant
    tg_dtype = dtypes.int8 if xi == np.int8 else dtypes.uint8
-    with Context(DONT_REALIZE_EXPAND=1):
-      out = (X.int().matmul(W.int())//1000)
-      if clip: out = out.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
-      out = out.cast(tg_dtype)
-      opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
-      sexec(out, opts, replace_src, run_count=1)
+    out = (X.int().matmul(W.int())//1000)
+    if clip: out = out.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
+    out = out.cast(tg_dtype)
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
+    sexec(out, opts, replace_src, run_count=1)
    tout = out.numpy()
    mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000)
    if clip: mout = mout.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
@@ -225,7 +221,6 @@ class TestQuantizeOnnx(unittest.TestCase):

  def test_prequant_gemv(self):
    N = 2048
-    # ugh, it's so broken with those casts. need DONT_REALIZE_EXPAND=1 python3 test/test_quantize_onnx.py TestQuantizeOnnx.test_prequant
    X = Tensor(np.random.uniform(0, 255, size=(1,N)).astype(np.uint8)).realize()
    W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8)).realize()
    #out = X.cast(dtypes.int) @ W.cast(dtypes.int)
@@ -6,6 +6,7 @@ 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 tinygrad.renderer.nir import NIRRenderer
 from test.helpers import not_support_multi_device

 import numpy as np
@@ -100,7 +101,7 @@ class TestRandomness(unittest.TestCase):

    np.testing.assert_allclose(jr, r)

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "fails with PTX")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "PTX and NIR use pointer arithmetic")
  def test_threefry_doesnt_use_long(self):
    for (_,ei) in lower_schedule(Tensor.rand(20).schedule()):
      if isinstance(ei.prg, CompiledRunner):
@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, nn
-from tinygrad.helpers import Context, GlobalCounters
+from tinygrad.helpers import Context, GlobalCounters, CI, CPU_LVP, getenv
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops

 class TestRangeifyAssign(unittest.TestCase):
@@ -17,8 +17,89 @@ class TestRangeifyAssign(unittest.TestCase):
    self.assertListEqual(lst, lst3)
    self.assertListEqual(lst2, B.permute(1, 0).tolist())

+class TestRangeifyEdgeCase(unittest.TestCase):
+  def test_matmul_relu_cat(self):
+    a = Tensor.ones(100, 512).contiguous().realize()
+    c = Tensor.ones(1, 512).contiguous().realize()
+    cm = Tensor.ones(512, 512)
+    c = c @ cm
+    c = c.relu()
+
+    res = Tensor.cat(a, c, dim=0)
+    self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
+
+if getenv("BIG") > 2:
+  # llama 8B (8192)
+  BS, HEADS, SEQLEN, EMB = 4, 32, 8192, 128
+elif getenv("BIG") > 1:
+  # llama 8B
+  BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
+elif getenv("BIG") > 0:
+  # bigger
+  BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
+else:
+  BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
+
+@unittest.skipIf(CPU_LVP, "broken in LVP")
+class TestPcontig(unittest.TestCase):
+  def test_flash_attention_bw(self):
+    def fa_bw():
+      Tensor.manual_seed(1337)
+      with Context(DEBUG=0):
+        q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize().requires_grad_() for _ in range(3)]
+        attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
+        attn_output.weight.requires_grad_().realize()
+        target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
+
+      GlobalCounters.reset()
+      attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward()
+      attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1)
+      out = attn_output(attn)
+      loss = (out - target).square().mean()
+      loss.backward()
+      #ret = [out, Tensor.stack(q.grad, k.grad, v.grad)]
+      ret = [out, q.grad, k.grad, v.grad]
+      Tensor.realize(*ret)
+      return ret
+
+    with Context(PCONTIG=2, REAL_SUBSTITUTE=1, DEBUG=2):
+      grads = fa_bw()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+
+    with Context(DEBUG=2):
+      cmp_grads = fa_bw()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+
+    with Context(DEBUG=0):
+      mses = [((x-y)**2).sum().item() for x,y in zip(grads, cmp_grads)]
+    mse = sum(mses)
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-6)
+
+  def test_flash_attention(self):
+    def fa():
+      Tensor.manual_seed(1337)
+      with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
+      GlobalCounters.reset()
+      return q.scaled_dot_product_attention(k, v).realize()
+
+    with Context(PCONTIG=2, DEBUG=2):
+      ret = fa()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+    with Context(DEBUG=2):
+      cmp = fa()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+    with Context(DEBUG=0):
+      mse = ((cmp-ret)**2).sum().item()
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-6)
+
+
+# *** non CI rangeify tests below this line ***
+
 N = 256

+@unittest.skipIf(CI, "useless in CI, doesn't test anything")
 class TestRangeifyOpt(unittest.TestCase):
  def test_randperm(self):
    Tensor.randperm(10000).realize()
@@ -54,6 +135,7 @@ class TestRangeifyOpt(unittest.TestCase):
    A = Tensor.empty(8,8,8,8).permute(1,0,3,2).flatten()
    A.sum().realize()

+@unittest.skipIf(CI, "useless in CI, doesn't test anything")
 class TestRangeify(unittest.TestCase):
  def test_groupnorm(self):
    # ranges 1 and 3 are merging
@@ -200,33 +282,6 @@ class TestRangeify(unittest.TestCase):
    out = blk._feed_forward(x)
    out.realize()

-  @unittest.skip("RANGEIFY=0 does nothing")
-  def test_flash_attention(self):
-    BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
-
-    # bigger
-    #BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
-
-    # llama 8B
-    #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
-
-    def fa():
-      Tensor.manual_seed(1337)
-      with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
-      return q.scaled_dot_product_attention(k, v).realize()
-
-    with Context(DEBUG=4):
-      GlobalCounters.reset()
-      ret = fa()
-    with Context(RANGEIFY=0):
-      with Context(DEBUG=2):
-        GlobalCounters.reset()
-        cmp = fa()
-      with Context(DEBUG=0):
-        mse = ((cmp-ret)**2).sum().item()
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-6)
-
 # contiguous + reduce can support ranges?

@unittest.skip("pm_rangeify no longer exists. test this in a different way")
@@ -280,16 +335,5 @@ class TestRangeifyPM(unittest.TestCase):
    b = self.base.pad(((0,1),(0,0))).pad(((0,0),(0,1)))
    self.assert_same(a, b)

-class TestRangeifyEdgeCase(unittest.TestCase):
-  def test_matmul_relu_cat(self):
-    a = Tensor.ones(100, 512).contiguous().realize()
-    c = Tensor.ones(1, 512).contiguous().realize()
-    cm = Tensor.ones(512, 512)
-    c = c @ cm
-    c = c.relu()
-
-    res = Tensor.cat(a, c, dim=0)
-    self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
-
 if __name__ == '__main__':
  unittest.main()
@@ -2,16 +2,14 @@
 # schedule confirms the right things are capable of fusing
 # NOTE: this has overlap with external_test_opt.py

-import unittest
+import unittest, functools
 import numpy as np
-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.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
 from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
 from tinygrad.schedule.rangeify import get_rangeify_map, Kernel
@@ -32,7 +30,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:
-    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):
@@ -118,8 +115,7 @@ class TestSchedule(unittest.TestCase):
    c = a+b
    with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 2)

-  @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")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  def test_expand_buffer_before_cast(self):
    a = Tensor.randn(4, 2, 1).realize().permute((1, 0, 2))
    b = a.cast(dtypes.half).expand((2, 4, 4))+2
@@ -129,7 +125,7 @@ class TestSchedule(unittest.TestCase):
  def test_indexing_scalars_simple(self):
    X = Tensor.randn(2, 2).realize()
    xt = X[Tensor(1)][Tensor(0)]
-    run_schedule(check_schedule(xt, 2))
+    run_schedule(check_schedule(xt, 1))
    np.testing.assert_equal(xt.numpy(), X.numpy()[1][0])

  @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
@@ -149,31 +145,30 @@ class TestSchedule(unittest.TestCase):
    assume(a<x and b<y)
    X = Tensor.randn(x, y).realize()
    xt = X[Tensor(a)][Tensor(b)]
-    run_schedule(check_schedule(xt, 2))
+    run_schedule(check_schedule(xt, 1))
    np.testing.assert_equal(xt.numpy(), X.numpy()[a][b])

  def test_push_pads_elementwise(self):
    x = Tensor.full((4,4), 2.).contiguous().realize()
    y = Tensor.full((4,4), 4.).contiguous().realize()
    z = (x.reciprocal()*y).pad((None, (0,1),)).sum()
-    run_schedule(check_schedule(z, 2))
+    run_schedule(check_schedule(z, 1))
    self.assertEqual(z.item(), 32)

  def test_push_pads_contiguous(self):
    x = Tensor.full((4,1), 2.).contiguous()
    y = Tensor.full((4,4), 4.).contiguous()
    z = (x.reciprocal().expand(4,4)*y).pad((None, (0,1),)).sum()
-    run_schedule(check_schedule(z, 2, [x,y]))
+    run_schedule(check_schedule(z, 1, [x,y]))
    self.assertEqual(z.item(), 32)

  def test_rand(self):
    x = Tensor.rand(32)
-    check_schedule(x, 4, [Tensor._device_rng_counters[x.device]])
+    check_schedule(x, 1, [Tensor._device_rng_counters[x.device]])

  def test_rand_recompute_arange(self):
    x = Tensor.rand(32)
-    with Context(DONT_GROUP_REDUCES=1):
-      check_schedule(x, 3, [Tensor._device_rng_counters[x.device]])
+    check_schedule(x, 1, [Tensor._device_rng_counters[x.device]])

  def test_empty_is_not_realized(self):
    a = Tensor.empty(10)
@@ -190,10 +185,7 @@ 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)
+    check_schedule(a, 3)

  def test_basic_binop_fusion(self):
    a = Tensor.empty(10)
@@ -266,18 +258,17 @@ class TestSchedule(unittest.TestCase):
    c = a.sum(axis=0) + b
    check_schedule(c, 1)

-  # not pushing permutes through reduces
  def test_reduce_permute_binop_fusion(self):
    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)
+    check_schedule(c, 1)

  def test_allow_push_permutes(self):
    a = Tensor.randn(10,10,10).realize()
    b = Tensor.randn(10,10,1).realize()
    c = a.sum(axis=0, keepdim=True).permute(2,1,0) + b
-    with Context(DONT_GROUP_REDUCES=1): run_schedule(check_schedule(c, 1))
+    run_schedule(check_schedule(c, 1))
    np.testing.assert_allclose(c.numpy(), np.sum(a.numpy(), axis=0, keepdims=True).transpose(2,1,0)+b.numpy())

  def test_binop_early_reshape_reduce_fusion(self):
@@ -342,7 +333,7 @@ class TestSchedule(unittest.TestCase):
    r1 = (x - r0).sum(axis=0).div(2)
    out0 = r0 + y
    out1 = r1 + y
-    schedule = check_schedule([out0, out1], 2)
+    schedule = check_schedule([out0, out1], 3)
    reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
    self.assertEqual(len(reduceops), 2) # why is RANGEIFY different?

@@ -375,7 +366,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.full((4,), 2.).contiguous()
    first = a.assign(b)
    second = a.assign(b)
-    check_schedule([first, second], 1)
+    check_schedule([first, second], 2) # TODO: 1?

  # NOTE: this is causing "LAZYCACHE=1 incorrectly reuses contiguous const" #4562
  # should contiguous dedup?
@@ -455,7 +446,7 @@ class TestSchedule(unittest.TestCase):
  @unittest.skipUnless(is_dtype_supported(dtypes.ulong), "Needs ulong")
  def test_fold_conv_batchnorm_optim(self):
    # this is too high
-    for optim, cnt in [(nn.optim.Adam, 30), (nn.optim.SGD, 11)]:
+    for optim, cnt in [(nn.optim.Adam, 30), (nn.optim.SGD, 13)]:
      with self.subTest(optim=optim.__name__):
        with Tensor.train():
          img = Tensor.ones(1,3,4,4)
@@ -476,7 +467,7 @@ class TestSchedule(unittest.TestCase):
      fw = bn(x).contiguous_backward().relu().contiguous()
      fw.sum().backward()
      # TODO: this is too many
-      check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 10)
+      check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 9)

  def test_fold_conv_relu(self):
    c1 = nn.Conv2d(3,16,3)
@@ -519,9 +510,8 @@ class TestSchedule(unittest.TestCase):
    img = Tensor.empty(64,64)
    x = (img.sum(0) + img.sum(1))
    out = x.relu()
-    check_schedule(out, 2)
+    check_schedule(out, 1)

-  #@unittest.skip("failing in old lazy")
  def test_push_permute_through_reshape(self):
    a = Tensor.empty(16,16)
    b = Tensor.empty(16,16)
@@ -555,7 +545,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)
+    check_schedule(out, 1)

  def test_children_dont_push(self):
    a = Tensor.empty(10, 10, 1)
@@ -563,7 +553,7 @@ class TestSchedule(unittest.TestCase):
    d = (a+b).expand(10, 10, 10)
    e = (a+b).permute(2,1,0)
    f = d+e
-    check_schedule(f, 2)
+    check_schedule(f, 1)

  # failing in new lazy
  @unittest.skip("always fusing elementwise")
@@ -602,13 +592,13 @@ class TestSchedule(unittest.TestCase):
    e = c[0] * d
    check_schedule(e, 1)

-  def test_expand_nofuse(self):
+  def test_expand_fuse(self):
    a = Tensor.empty(1, 16)
    b = Tensor.empty(1, 16)
    c = a * b
    d = Tensor.empty(8192, 16)
    e = c * d
-    check_schedule(e, 2)
+    check_schedule(e, 1)

  # this is the failing case in openpilot...it's very simple like this
  def test_image_conv_fusion(self):
@@ -626,7 +616,7 @@ class TestSchedule(unittest.TestCase):

    # NOOP, 3 convs, contiguous
    #check_schedule(x, 5)
-    check_schedule(x, 8)
+    check_schedule(x, 7)

  def test_image_conv_fusion_minimal(self):
    b1 = Tensor.empty(16)
@@ -809,13 +799,13 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.empty(32, 32, 32)
    y = Tensor.empty(32, 32)
    out = x.sum(axis=2).T+y
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  def test_two_elus_sum(self):
    x = Tensor.empty(32, 32)
    y = Tensor.empty(32, 32)
    out = x.sum(1).relu().elu() + y.sum(1).relu().elu()
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  @unittest.skipUnless(SPLIT_REDUCEOP, "Testing split reducop requires SPLIT_REDUCEOP")
  def test_preserve_multistage_reduce(self):
@@ -828,7 +818,7 @@ class TestSchedule(unittest.TestCase):
  def test_multistage_reduce(self):
    x = Tensor.empty(32, 32, 32)
    out = x.sum(2).relu().sum(1)
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  def test_multistage_reduce_fork(self):
    x = Tensor.empty(32, 32, 32)
@@ -844,7 +834,7 @@ class TestSchedule(unittest.TestCase):
    z = y.matmul(x).sum()
    z.backward()
    out = x.grad.contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), np.ones((64,64)))

  def test_example_matmul_contig(self):
@@ -853,7 +843,7 @@ class TestSchedule(unittest.TestCase):
    z = y.matmul(x).sum()
    z.backward()
    out = x.grad.contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), np.ones((64,64)))

  def test_example_matmul_same(self):
@@ -861,7 +851,7 @@ class TestSchedule(unittest.TestCase):
    z = x.matmul(x).sum()
    z.backward()
    out = x.grad.contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    # NOTE: the gradient flows twice
    np.testing.assert_allclose(out.numpy(), 2*np.ones((64,64)))

@@ -884,8 +874,7 @@ class TestSchedule(unittest.TestCase):
    x = x.sum(1)
    x = x[:16]
    out = x + y
-    # NOTE: this could be 1 kernel if we mask the store?
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  def test_multireduce_shrink(self):
    Tensor.manual_seed(0)
@@ -897,8 +886,7 @@ class TestSchedule(unittest.TestCase):
    b_out = b.sum(1)
    b_out = b_out[:16]
    out = a_out + b_out + c
-    # run_schedule(check_schedule(out, 2))  # TODO: this should be 1 (can we make it 1 with the new linearizer?)
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), a.numpy().sum(axis=1)[:16] + b.numpy().sum(axis=1)[:16] + c.numpy(), atol=1e-4, rtol=1e-4)

  # broken due to const folding and two contiguous are different kernels
@@ -915,7 +903,7 @@ class TestSchedule(unittest.TestCase):
    out0 = a.sum() + 2
    out1 = a.sum() + 4
    out2 = out0 * out1
-    run_schedule(check_schedule([out0, out1, out2], 1))
+    run_schedule(check_schedule([out0, out1, out2], 3)) # TODO: 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)
@@ -926,7 +914,7 @@ class TestSchedule(unittest.TestCase):
    out0 = a.sum().exp2()
    # out1 has two paths to a.sum()
    out1 = a.sum() + out0
-    run_schedule(check_schedule([out0, out1], 1))
+    run_schedule(check_schedule([out0, out1], 2)) # TODO: 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)

@@ -939,7 +927,7 @@ class TestSchedule(unittest.TestCase):
    out2 = b.sum().exp2()
    out3 = b.sum() + out2
    # run_schedule(check_schedule([out0, out1, out2, out3], 1))
-    run_schedule(check_schedule([out0, out1, out2, out3], 6))
+    run_schedule(check_schedule([out0, out1, out2, out3], 4))
    np.testing.assert_allclose(out0.numpy(), np_out0:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out1.numpy(), np_out1:=a.numpy().sum()+np_out0, atol=1e-4, rtol=1e-4)
    np_b = (a.numpy() + np_out0 + np_out1)
@@ -954,7 +942,7 @@ class TestSchedule(unittest.TestCase):
    out0 = a.sum() + b.sum() + 2
    out1 = a.sum() + b.sum() + 4
    # run_schedule(check_schedule([out0, out1], 1))
-    run_schedule(check_schedule([out0, out1], 4))
+    run_schedule(check_schedule([out0, out1], 2))
    np.testing.assert_allclose(out0.numpy(), a.numpy().sum()+b.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+b.numpy().sum()+4, atol=1e-4, rtol=1e-4)

@@ -981,7 +969,7 @@ class TestSchedule(unittest.TestCase):
    out1 = b.max() + out0*2
    out2 = a.sum() + out1
    # run_schedule(check_schedule([out0, out1, out2], 1))
-    run_schedule(check_schedule([out0, out1, out2], 4))
+    run_schedule(check_schedule([out0, out1, out2], 3))
    np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6)
    np.testing.assert_allclose(out1.numpy(), out1_np:=b.numpy().max() + out0_np*2, atol=1e-4, rtol=1e-6)
    np.testing.assert_allclose(out2.numpy(), a.numpy().sum() + out1_np, atol=1e-4, rtol=1e-6)
@@ -1018,7 +1006,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.empty(10,)
    c = a.sum() + b[0]
    d = a.sum() + 2
-    check_schedule([c, d], 1)
+    check_schedule([c, d], 2) # TODO: 1?

  def test_reduce_multiple_paths_midshrink(self):
    a = Tensor.empty(4, 4)
@@ -1047,7 +1035,7 @@ class TestSchedule(unittest.TestCase):
    k = Tensor.randn(32,8,16,8).realize()
    v = Tensor.randn(32,8,16,8).realize()
    out = Tensor.scaled_dot_product_attention(q,k,v)
-    run_schedule(check_schedule(out, 5))
+    run_schedule(check_schedule(out, 4))
    if getenv("CHECK", 1):
      import torch
      compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
@@ -1055,7 +1043,7 @@ class TestSchedule(unittest.TestCase):

    with Context(FUSE_ATTENTION=1):
      out = Tensor.scaled_dot_product_attention(q,k,v)
-      run_schedule(check_schedule(out, 1))
+      run_schedule(check_schedule(out, 4)) # TODO: should be 1?
      if getenv("CHECK", 1):
        import torch
        compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
@@ -1068,7 +1056,7 @@ class TestSchedule(unittest.TestCase):
    c = Tensor.randn(4, 32).realize()
    out = (c * a.sum(-1, keepdim=True)).sum(-1) + (b * a.sum(-1, keepdim=True)).sum(-1) # a.sum has >1 children but should still fuse
    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), \
      (c.numpy()*a.numpy().sum(axis=-1,keepdims=True)).sum(-1) + (b.numpy()*a.numpy().sum(axis=-1,keepdims=True)).sum(-1), atol=1e-4, rtol=1e-4)

@@ -1113,8 +1101,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(4, 32).realize()
    y = Tensor.randn(4, 32).realize()
    out = y.sum(axis=-1) + x.sum(axis=-1)
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), y.numpy().sum(axis=-1) + x.numpy().sum(axis=-1), atol=1e-4, rtol=1e-4)

  def test_multireduce_fusion_sequential(self):
@@ -1131,7 +1118,7 @@ class TestSchedule(unittest.TestCase):
    y = Tensor.randn(4, 32).realize()
    out = x.std(-1) + y.std(-1)
    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 4))
+    run_schedule(check_schedule(out, 3))
    np.testing.assert_allclose(out.numpy(), x.numpy().std(axis=-1, ddof=1) + y.numpy().std(axis=-1, ddof=1), atol=1e-4, rtol=1e-4)

  def test_multireduce_diffops_sequential(self):
@@ -1147,8 +1134,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(4, 32).realize()
    y = Tensor.randn(4, 32).realize()
    out = x.sum(-1) + y.max(-1)
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), x.numpy().sum(axis=-1) + y.numpy().max(axis=-1), atol=1e-4, rtol=1e-4)

  def test_multireduce_fusion_sequential_and_parallel(self):
@@ -1160,7 +1146,7 @@ class TestSchedule(unittest.TestCase):
    np_mu = (x.numpy() - x.numpy().max(axis=-1, keepdims=True)).mean(axis=-1, keepdims=True) + \
      (y.numpy() - y.numpy().max(axis=-1, keepdims=True)).mean(axis=-1, keepdims=True)
    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 6))
+    run_schedule(check_schedule(out, 5))
    np.testing.assert_allclose(out[0].numpy(), np.sqrt(np.square(x.numpy() - np_mu).sum(-1)/x.shape[-1]), atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out[1].numpy(), np.sqrt(np.square(y.numpy() - np_mu).sum(-1)/y.shape[-1]), atol=1e-4, rtol=1e-4)

@@ -1169,8 +1155,7 @@ class TestSchedule(unittest.TestCase):
    a,b = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()
    c,d = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()
    out = a@b + c@d
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), a.numpy()@b.numpy() + c.numpy()@d.numpy(), atol=1e-4, rtol=1e-4)

  def test_softmax_fusion(self):
@@ -1181,17 +1166,15 @@ class TestSchedule(unittest.TestCase):
    expected = (x_exp:=np.exp(x.numpy()-x.numpy().max(-1, keepdims=True)))/x_exp.sum(-1, keepdims=True)
    np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)

-  # TODO: rangeify stores the output in float32
  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  @unittest.expectedFailure
  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)
-    self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
+    self.assertEqual(len(sched), 3)
+    self.assertEqual(sched[0].bufs[0].dtype, dtypes.float)

    # input float, softmax in float
    Tensor.manual_seed(0)
@@ -1223,12 +1206,12 @@ class TestSchedule(unittest.TestCase):
  def test_scaled_dot_product_attention_fusion(self):
    x, y, z, m = (Tensor.empty(32, 8, 16, 16) for _ in range(4))
    out = Tensor.scaled_dot_product_attention(x, y, z, attn_mask=m)
-    check_schedule(out, 5)
+    check_schedule(out, 4)

  def test_scaled_dot_product_attention_causal_fusion(self):
    x, y, z = (Tensor.empty(32, 8, 16, 16) for _ in range(3))
    out = Tensor.scaled_dot_product_attention(x, y, z, is_causal=True)
-    check_schedule(out, 5)
+    check_schedule(out, 4)

  def test_adam_step_fusion(self):
    with Tensor.train():
@@ -1258,7 +1241,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.Adam(nn.state.get_parameters([c1, c2]), lr=1e-4)
      opt.zero_grad()
      c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 20)
+      check_schedule(opt.schedule_step(), 18)

  def test_sgd_conv_fuse(self):
    with Tensor.train():
@@ -1268,7 +1251,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters(c1))
      opt.zero_grad()
      c1(img).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 3)
+      check_schedule(opt.schedule_step(), 5) # TODO: 3?

  def test_sgd_2convs_fuse(self):
    with Tensor.train():
@@ -1291,7 +1274,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2]), nesterov=True, momentum=0.9, weight_decay=0.1)
      opt.zero_grad()
      c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 13)
+      check_schedule(opt.schedule_step(), 15)

  def test_sgd_4convs_fuse(self):
    with Tensor.train():
@@ -1304,7 +1287,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4]))
      opt.zero_grad()
      c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 17)
+      check_schedule(opt.schedule_step(), 15)

  def test_sgd_4convs_fuse_conv_bw(self):
    with Tensor.train():
@@ -1317,50 +1300,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4]))
      opt.zero_grad()
      c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 14)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  @unittest.expectedFailure
-  def test_prefer_half_buffer(self):
-    x = Tensor.ones(4).contiguous().realize()
-    # y = Tensor.ones(4).contiguous().realize()
-    z = Tensor.ones(4, 4).contiguous().realize()
-
-    # should not create extra kernel if output will be realized anyways
-    dummy = x.sum().half().float()
-    check_schedule(dummy, 1)
-    dummy = x.sum().half().float().contiguous() + 1
-    check_schedule(dummy, 2)
-
-    # shared between two outputs
-    shared = x.sum().half().float()
-    a = shared * 2
-    b = shared * 3
-    sched = check_schedule([a, b], 3)
-    # store reduceop in half
-    self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
-    # fuse cast with the child kernel
-    self.assertEqual(sched[1].bufs[0].dtype, dtypes.float)
-    self.assertEqual(sched[2].bufs[0].dtype, dtypes.float)
-
-    # reduce
-    a = z.sum(axis=0).half().float().sum(axis=0)
-    sched = check_schedule(a, 2)
-    self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
-    self.assertEqual(sched[1].bufs[0].dtype, dtypes.float)
-
-    # expand
-    # expand will realize just after the .float(), so requires change to realize-before-expand
-    # normal = (x.sum().half().float().reshape(1) * y).sum()
-    # sched = check_schedule(normal, 2)
-    # for si in sched[:-1]: assert all(out.dtype == dtypes.half for out in si.outputs[:-1])
-
-    # parallel reduce
-    # a = x.sum().half().float() * y.sum().half().float()
-    # b = a + 1
-    # c = a + 2
-    # sched = check_schedule([b, c], 4)
-    # doesn't store either in half because it doesn't chase
+      check_schedule(opt.schedule_step(), 15)

  def test_reduce_simple_chase(self):
    a = Tensor.empty(4, 4, 4)
@@ -1409,7 +1349,7 @@ class TestSchedule(unittest.TestCase):
    c = Tensor.empty(16, )
    r = a.sum(1) + c
    d = r[:4] * b
-    check_schedule(d, 2)
+    check_schedule(d, 1)

  def test_multireduce_push_shrink_chase(self):
    Tensor.manual_seed(0)
@@ -1419,22 +1359,20 @@ class TestSchedule(unittest.TestCase):
    d = Tensor.randn(16, 16).realize()
    r = a.sum(1) + c
    out = r[:4] * b + d.sum(1)[:4]
-    # schedule = check_schedule(out, 2)
-    schedule = check_schedule(out, 3)
+    schedule = check_schedule(out, 1)
    run_schedule(schedule)
    np.testing.assert_allclose(out.numpy(), (a.numpy().sum(1) + c.numpy())[:4] * b.numpy() + d.numpy().sum(1)[:4], atol=1e-4, rtol=1e-4)

  def test_midreduce_nochase(self):
    a = Tensor.empty(16, 16)
    b = (a.sum(0) + a.max(1)) + 2
-    check_schedule(b, 2)
+    check_schedule(b, 1)

  def test_multireduce_midreduce_nochase(self):
    Tensor.manual_seed(0)
    a = Tensor.randn(16, 16).realize()
    b = (a.sum(0)+a.max(0) + a.max(1)+a.sum(1)) + 2
-    # schedule = check_schedule(b, 2)
-    schedule = check_schedule(b, 4)
+    schedule = check_schedule(b, 1)
    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)

@@ -1446,7 +1384,7 @@ class TestSchedule(unittest.TestCase):
    c = a.sum() + 2
    d = (a.sum() - b.sum()) * 4
    # run_schedule(check_schedule([c, d], 1))
-    run_schedule(check_schedule([c, d], 3))
+    run_schedule(check_schedule([c, d], 2))
    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)

@@ -1472,7 +1410,7 @@ 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))
+    run_schedule(check_schedule([c, d, e, f], 4))
    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)
@@ -1487,7 +1425,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], 4))
    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)
@@ -1506,8 +1444,7 @@ class TestSchedule(unittest.TestCase):
    a = Tensor.randn(3, 4, 5).realize()
    b = Tensor.randn(3, 4, 5).realize()
    out = (a.pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum(keepdim=True)+b.pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum()).contiguous()
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), np.pad(a.numpy(), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(keepdims=True) + \
                                                   np.pad(b.numpy(), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=1e-4, rtol=1e-4)

@@ -1515,7 +1452,7 @@ class TestSchedule(unittest.TestCase):
    Tensor.manual_seed(0)
    a = Tensor.rand(3, 4, 5).realize()
    out = a.log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=1e-5, rtol=1e-6)

  def test_multireduce_pad_reduce_unsafe(self):
@@ -1524,7 +1461,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.randn(3, 4, 5).abs().realize()
    out = (a.log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum()+b).abs().log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 4))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), np.pad(np.log2(np.abs(np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum() + \
                                                   b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-5)

@@ -1538,7 +1475,7 @@ class TestSchedule(unittest.TestCase):
  def test_shrink_pad_unsafe(self):
    a = Tensor.ones((3, )).contiguous().realize()
    out = a.exp2().shrink(((0, 1),)).pad(((0, 1),)).contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_equal(out.numpy(), [2, 0])

  def test_base_change_shrink_pad(self):
@@ -1546,7 +1483,7 @@ class TestSchedule(unittest.TestCase):
    b = a.exp2()
    c = b[:-1, :-1]
    d = c.pad(((0, 1), (0, 1))) * 2
-    run_schedule(check_schedule(d, 2))
+    run_schedule(check_schedule(d, 1))
    np.testing.assert_equal(d.numpy(), np.pad(np.exp2(a.numpy())[:-1, :-1], ((0, 1), (0, 1)))*2)

  def test_base_change_expand_pad(self):
@@ -1554,14 +1491,14 @@ class TestSchedule(unittest.TestCase):
    b = a.exp2()
    c = b[:, None, :]
    d = c.pad(((0, 0), (1, 1), (0, 0))) * 2
-    run_schedule(check_schedule(d, 2))
+    run_schedule(check_schedule(d, 1))
    np.testing.assert_equal(d.numpy(), np.pad(np.exp2(a.numpy())[:, None, :], ((0, 0), (1, 1), (0, 0)))*2)

  def test_fuse_arange_pad_replicate_mode(self):
    x = Tensor.empty(3,3,3,3, requires_grad=True)
    y = x.pad((-1,2,2,-1), mode="replicate")
    dx = y.sum().gradient(x)[0]
-    sched = check_schedule(dx, 3)
+    sched = check_schedule(dx, 1)
    run_schedule(sched)
    np.testing.assert_allclose(dx.numpy(), [[[[0.,3.,9.],[0,1.,3.],[0.,0.,0.]]]*3]*3)

@@ -1571,7 +1508,7 @@ class TestSchedule(unittest.TestCase):
    a = Tensor.ones(4, 4).contiguous().realize()
    b = a.cast(dtypes.half).expand(2, 4, 4)
    c = b.cast(dtypes.int).expand(2, 2, 4, 4)
-    run_schedule(check_schedule(c, 2))
+    run_schedule(check_schedule(c, 1))
    np.testing.assert_equal(c.numpy(), np.ones(((2, 2, 4, 4)), dtype=np.int32))

  def test_base_change_pad_expand(self):
@@ -1579,7 +1516,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.full((4, 4), 2.).contiguous().realize()
    c = (a + b).pad(((1, 1), (1, 1)))
    d = c.cast(dtypes.int).expand((2, 6, 6)) * 4
-    run_schedule(check_schedule(d, 2))
+    run_schedule(check_schedule(d, 1))
    c_np = np.pad((np.full((4, 4), 2., dtype=np.float32) + np.full((4, 4), 1., dtype=np.float32)), ((1, 1), (1, 1)), constant_values=0.0)
    np.testing.assert_equal(d.numpy(), np.broadcast_to(c_np.astype(np.half), (2, *c_np.shape)) * 4)

@@ -1678,7 +1615,7 @@ class TestSchedule(unittest.TestCase):
    self._test_fusion([(4, 4), (1, 4)], lambda a,b:a.sum(1).reshape(b.shape)+b, 1)

  def test_late_fusion_post_permute(self):
-    self._test_fusion([(4, 6, 4), (4, 4, 1)], lambda a,b:a.sum(1, keepdim=True).permute((2, 0, 1))+b, 2)
+    self._test_fusion([(4, 6, 4), (4, 4, 1)], lambda a,b:a.sum(1, keepdim=True).permute((2, 0, 1))+b, 1)

  def test_late_fusion_double_transpose(self):
    self._test_fusion([(32, 16, 1)],
@@ -1716,6 +1653,7 @@ 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))
+  @unittest.skip("kernel count depends on input")
  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))
@@ -1729,7 +1667,7 @@ class TestSchedule(unittest.TestCase):
    X = Tensor.randn(10, 10).realize()
    idxs = Tensor([0, 2]).realize()
    xt = X[idxs]
-    run_schedule(check_schedule(xt, 2))
+    run_schedule(check_schedule(xt, 1))
    np.testing.assert_equal(xt.numpy(), X.numpy()[idxs.numpy()])

  def test_simple_indexing_alt(self):
@@ -1747,7 +1685,7 @@ class TestSchedule(unittest.TestCase):
  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))
+    run_schedule(check_schedule(xt, 1))
    np.testing.assert_equal(xt.numpy(), 6)

  def test_advanced_simple_indexing_combined(self):
@@ -1795,7 +1733,7 @@ class TestSchedule(unittest.TestCase):
    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))
+    run_schedule(check_schedule(out, 1))
    self.assertListEqual(out.tolist(), [0, 12, 4, 0])

  def test_arange_transposed_descendants(self):
@@ -1828,7 +1766,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(5, 2).realize()
    a = Tensor.arange(10).contiguous()
    out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 2))
    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,7 +1782,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(5, 2).realize()
    a = (Tensor.arange(10)+1).contiguous()
    out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)

  @unittest.skip("BUFFER_VIEW no longer supported on non-disk devices")
@@ -1859,10 +1797,10 @@ class TestSchedule(unittest.TestCase):
    from extra.models.llama import precompute_freqs_cis
    args = {"dim":32 if CI else 128, "end":2048 if CI else 8192, "theta":10000}
    fused = precompute_freqs_cis(**args)
-    run_schedule(check_schedule(fused, 3))
+    run_schedule(check_schedule(fused, 1))
    if getenv("CHECK", 1):
      ref = precompute_freqs_cis(**args)
-      run_schedule(check_schedule(ref, 3))
+      run_schedule(check_schedule(ref, 1))
      np.testing.assert_equal(fused.numpy(), ref.numpy())

  def test_fuse_assign_contiguous(self):
@@ -1904,7 +1842,7 @@ class TestSchedule(unittest.TestCase):
    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))
+    run_schedule(check_schedule(loss, 3))
    np.testing.assert_allclose(loss.item(), 0.878309, atol=1e-5, rtol=1e-6)

  def test_const_folding_alt(self):
@@ -1925,7 +1863,7 @@ 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))
+      run_schedule(check_schedule(loss, 5))
      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)
@@ -1935,7 +1873,7 @@ class TestSchedule(unittest.TestCase):
    r = (X+Tensor.arange(16).reshape(4, 4)).sum()
    out0 = r+2
    out1 = r+3
-    run_schedule(check_schedule([out0, out1], 1))
+    run_schedule(check_schedule([out0, out1], 2)) # TODO: 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)
@@ -1977,8 +1915,7 @@ class TestSwizzle(unittest.TestCase):
      a = Tensor.randint(32, 32).realize()
    r = (a+a).sum(1).sum(0)
    # double reduce collapses to a single reduce
-    with Context(DONT_GROUP_REDUCES=1):
-      run_schedule(check_schedule(r, 1))
+    run_schedule(check_schedule(r, 1))
    self.assertEqual(r.numpy(), (a.numpy()+a.numpy()).sum(1).sum(0))

  def test_single_swizzle(self):
@@ -1998,33 +1935,29 @@ class TestSwizzle(unittest.TestCase):
      b = Tensor.randint(4,).realize()
    # parallel reduce!
    add = a.sum(0)+b.sum(0)
-    with Context(DONT_GROUP_REDUCES=1):
-      run_schedule(check_schedule(add, 1))
+    run_schedule(check_schedule(add, 1))
    self.assertEqual(add.numpy(), a.numpy().sum(0)+b.numpy().sum(0))

-  @unittest.skip("TODO: how do we express the norm")
  def test_softmax_one_kernel(self):
    Tensor.manual_seed(0)
    with Context(DEBUG=0, TRACK_MATCH_STATS=0):
      a = Tensor.randn(32, 32).realize()
    t = a.softmax()
-    with Context(DONT_GROUP_REDUCES=1, DONT_REALIZE_EXPAND=1):
-      check_schedule(t, 1)
+    check_schedule(t, 3) # TODO: 1?

  def test_argmax_one_kernel(self):
    Tensor.manual_seed(0)
    with Context(DEBUG=0, TRACK_MATCH_STATS=0):
      a = Tensor.randn(10, 20).realize()
    t = a.argmax(0)
-    with Context(DONT_GROUP_REDUCES=1, DONT_REALIZE_EXPAND=1): t.realize()
+    check_schedule(t, 2) # TODO: 1?

  def test_swizzle_reduceop(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(4,4).realize()
    y = Tensor.randn(4,4,4).realize()
    out = x.reshape(4,4,1).expand(4,4,4).sum(axis=(1,))+y
-    with Context(DONT_REALIZE_EXPAND=1, DONT_GROUP_REDUCES=1):
-      run_schedule(check_schedule(out, 1))
+    run_schedule(check_schedule(out, 2)) # TODO: 1?
    np.testing.assert_allclose(out.numpy(), np.tile(x.numpy().reshape(4,4,1), (1,1,4)).sum(axis=1)+y.numpy())

  def test_permute_rewrite(self):
@@ -2032,7 +1965,7 @@ class TestSwizzle(unittest.TestCase):
    y = Tensor.randn(4, 1, 16).realize()
    z = Tensor.randn(4, 4, 1).realize()
    t = (x*y).sum(axis=(0, 2)).reshape(1, 4, 1).permute(0, 2, 1)+z
-    with Context(DONT_GROUP_REDUCES=1, DONT_REALIZE_EXPAND=1): run_schedule(check_schedule(t, 1))
+    run_schedule(check_schedule(t, 2)) # TODO: 1?
    t_np = (x.numpy()*y.numpy()).sum(axis=(0, 2)).reshape(1, 4, 1).transpose(0, 2, 1)+z.numpy()
    np.testing.assert_allclose(t.numpy(), t_np, atol=1e-6, rtol=1e-3)

@@ -2043,14 +1976,14 @@ class TestSwizzle(unittest.TestCase):
    a_reduce = a.sum(axis=(2,), keepdim=True).sum(axis=(1,))
    b_reduce = b.sum(axis=(0,))
    t = a_reduce+b_reduce
-    with Context(DONT_GROUP_REDUCES=1, DONT_REALIZE_EXPAND=1): run_schedule(check_schedule(t, 1))
+    run_schedule(check_schedule(t, 1))

  def test_parallel_reduce_possible(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 2, 2).realize()
    y = Tensor.randn(4, 2, 2).realize()
    t = x.sum(axis=1)+y.sum(axis=1)
-    with Context(DONT_GROUP_REDUCES=1): run_schedule(check_schedule(t, 1))
+    run_schedule(check_schedule(t, 1))
    np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)

  # kernels can only have 1 or n in each dim
@@ -2059,7 +1992,7 @@ class TestSwizzle(unittest.TestCase):
    x = Tensor.randn(4, 2, 2).realize()
    y = Tensor.randn(4, 3, 2).realize()
    t = x.sum(axis=1)+y.sum(axis=1)
-    with Context(DONT_GROUP_REDUCES=1): run_schedule(check_schedule(t, 1))
+    run_schedule(check_schedule(t, 1))
    np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)

  def test_unsafe_pad(self):
@@ -2152,7 +2085,7 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.arange(3).realize()
    zeros = Tensor.zeros(3).realize()
    b = (a*zeros).to("CPU")
-    run_schedule(check_schedule(b, 0, filter_sink=False))
+    run_schedule(check_schedule(b, 2, filter_sink=False)) # TODO: 0?
    self.assertListEqual(b.tolist(), [0, 0, 0])
    self.assertEqual(b.device, "CPU")

@@ -2160,8 +2093,8 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.ones((4,)).to("CPU")
    b = Tensor.empty(4, device="CPU")
    add = a+b
-    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]}"
+    add.kernelize()

  def test_alu_before_copy(self):
    buf = Tensor.ones(1).contiguous().realize()
@@ -2172,12 +2105,12 @@ class TestCopyFolding(unittest.TestCase):
  def test_copy_to_same_device(self):
    a = Tensor.empty(4).uop
    b = a.copy_to_device(a.device)
-    check_schedule(b, 0, filter_sink=False)
+    check_schedule(b, 1, filter_sink=False) # TODO: 0?

  def test_copy_to_same_device_alt(self):
    a = Tensor.empty(4, 4).uop
    b = a.copy_to_device(a.device)
-    check_schedule(b, 0, filter_sink=False)
+    check_schedule(b, 1, filter_sink=False) # TODO: 0?

  def test_copy_to_same_device_sched(self):
    a = Tensor.ones(4).contiguous().realize().uop.as_buf()
@@ -2192,13 +2125,11 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.empty(4)
    check_schedule(a.clone(), 1, filter_sink=False)

-  # NOTE: moving copy before view might change this
  def test_shrink_copy(self):
    a = Tensor.arange(4)
    view = a.shrink(((0, 2),))
    b = view.clone()
-    # NOTE: this was sort of a bug making this 2
-    run_schedule(check_schedule(b, 2, filter_sink=False))
+    run_schedule(check_schedule(b, 1, filter_sink=False))
    self.assertEqual(b.uop.base.buffer.size, 2)
    self.assertEqual(b.uop.size, 2)
    self.assertListEqual(b.tolist(), [0, 1])
@@ -2207,7 +2138,7 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.arange(2)
    view = a.reshape(2, 1).expand(2, 2)
    b = view.clone()
-    run_schedule(check_schedule(b, 2, filter_sink=False))
+    run_schedule(check_schedule(b, 1, filter_sink=False))
    self.assertEqual(b.uop.base.buffer.size, 4)
    self.assertEqual(b.uop.size, 4)
    self.assertListEqual(b.tolist(), [[0, 0], [1, 1]])
@@ -2251,7 +2182,7 @@ class TestBufferUOp(unittest.TestCase):
  def test_buffer_has_buffer(self):
    buf = Tensor.empty(10)
    self.assertIsNotNone(buf.uop.buffer)
-    self.assertEqual(buf.uop.st, ShapeTracker.from_shape((10,)))
+    self.assertEqual(buf.uop.shape, (10,))
    # the device Buffer remains unallocated until it's we run the schedule
    self.assertFalse(buf.uop.buffer.is_allocated())
    add = buf+1
@@ -2330,7 +2261,7 @@ class TestContiguous(unittest.TestCase):
  def test_double_contiguous_realizes_once(self):
    a = Tensor.empty(4, 1)
    b = a.expand((4, 4)).contiguous().contiguous()
-    check_schedule(b, 1)
+    check_schedule(b, 2) # TODO: should be 1?

  def test_view_does_not_realize(self):
    a = Tensor.empty(4)
@@ -2466,10 +2397,6 @@ class TestUOpBecome(unittest.TestCase):
    c = (a.reshape(1, 1, 4, 4)+0).shrink(((0, 1), (0, 1), (0, 3), (0, 3)))+0
    check_schedule([b, c], 0)
    assert all_same([x.uop.base.realized for x in [a,b,c]])
-    # these movement ops result in the same ShapeTracker
-    assert b.uop.st == c.uop.st
-    assert b.uop is c.uop
-    assert UPat(Ops.VIEW, src=(UPat(Ops.BUFFER),)).match(c.uop, {})

  def test_setitem_becomes_subbuffer(self):
    a = Tensor.full((4,), 2.).contiguous().realize()
@@ -52,7 +52,6 @@ class TestSetitem(unittest.TestCase):

  def test_setitem_into_noncontiguous(self):
    t = Tensor.ones(4)
-    self.assertFalse(t.uop.st.contiguous)
    with self.assertRaises(RuntimeError): t[1] = 5

  @unittest.skip("TODO: flaky")
@@ -165,8 +165,7 @@ class TestSoftmaxFusion(unittest.TestCase):
      sout.realize()

    print("*** single kernel softmax ***")
-    # NOTE: DONT_GROUP_REDUCES is required here
-    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2), DONT_GROUP_REDUCES=1):
+    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
      out = single_kernel_softmax(self.test)
      out.realize()

@@ -186,7 +185,6 @@ class TestSoftmaxFusion(unittest.TestCase):

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

-  @unittest.skip("recursion error no longer raised")
  def test_softmax_bw(self):
    print("*** softmax bw ***")
    self.test.requires_grad_()
@@ -197,14 +195,11 @@ class TestSoftmaxFusion(unittest.TestCase):
    self.test.grad = None

    print("*** single kernel softmax bw ***")
-    # NOTE: DONT_GROUP_REDUCES is required here
-    # TODO: fix RecursionError with DONT_GROUP_REDUCES
-    with self.assertRaises(RecursionError):
-      with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2), DONT_GROUP_REDUCES=1):
-        single_kernel_softmax(self.test).sum().backward()
-        g = self.test.grad.realize()
+    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
+      single_kernel_softmax(self.test).sum().backward()
+      g = self.test.grad.realize()

-      np.testing.assert_allclose(sg.numpy(), g.numpy(), atol=1e-7)
+    np.testing.assert_allclose(sg.numpy(), g.numpy(), atol=1e-7)

 if __name__ == '__main__':
  unittest.main()
@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Tensor, Variable, GlobalCounters
-from tinygrad.shape.shapetracker import View
 from tinygrad.uop.ops import sym_infer
 from tinygrad.dtype import dtypes
 from tinygrad.device import is_dtype_supported
@@ -64,14 +63,6 @@ class TestSymbolicOps(unittest.TestCase):
    self.test_attention(imin=4, imax=5, use_symbolic=False)
    self.test_attention(imin=4, imax=5, use_symbolic=True)

-  # until this works, symbolic single kernel softmax won't
-  @unittest.expectedFailure
-  def test_attention_simple_view(self):
-    i = Variable("i", 2, 10)
-    v1 = View.create((2,4,1,i,i), ((i*4),i,0,0,1))
-    v2 = View.create((2,4,1,i,i,i), (((i*i)*4),(i*i),0,0,i,1))
-    self.assertIsNotNone(v1+v2)
-
  def test_attention_training(self):
    with Tensor.train():
      self.test_attention(dropout_p=0.0)
@@ -1,4 +1,3 @@
-import subprocess
 import numpy as np
 import torch
 import unittest, copy, mmap, random, math, array
@@ -10,6 +9,7 @@ from hypothesis import given, settings, strategies as strat
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
 from tinygrad.codegen import full_rewrite
 from tinygrad.dtype import DType

@@ -515,32 +515,6 @@ class TestTinygrad(unittest.TestCase):
    print(a)
    print(c)

-  def test_env_overwrite_default_device(self):
-    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}\\""'],
-                    shell=True, check=True)
-    subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                    shell=True, check=True)
-
-    if Device.DEFAULT != "CPU":
-      # setting multiple devices fail
-      with self.assertRaises(subprocess.CalledProcessError):
-        subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                        shell=True, check=True)
-
-      # setting device via DEV
-      subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                      shell=True, check=True)
-      subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                      shell=True, check=True)
-      subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                      shell=True, check=True)
-
-      with self.assertRaises(subprocess.CalledProcessError):
-        subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                        shell=True, check=True)
-
  def test_no_attributeerror_after_apply_uop_exception(self):
    try:
      Tensor.arange(4).reshape(3,2)
@@ -553,6 +527,7 @@ class TestTinygrad(unittest.TestCase):
    self.assertListEqual(t.shrink_to(16).tolist(), list(range(16)))
    t = t.reshape(4, 8).contiguous().realize()
    self.assertListEqual(t.shrink_to(2, 2).tolist(), [[0, 1], [8, 9]])
+    self.assertListEqual(t.shrink_to(None, 2).tolist(), t.shrink_to(4, 2).tolist())
    with self.assertRaises(ValueError): t.shrink_to(2)
    with self.assertRaises(ValueError): t.shrink_to(2, 2, 2)

@@ -595,22 +570,6 @@ class TestMoveTensor(unittest.TestCase):
    np.testing.assert_equal(x.grad.numpy(), [[2,2,2],[0,0,0],[-2,-2,-2]])

 class TestZeroShapeTensor(unittest.TestCase):
-  def test_shape_is_expanded(self):
-    t = Tensor.empty(3, 2, 0)
-    assert t.shape == (3, 2, 0)
-    # numpy has stride 0, 0, 0; torch has stride 2, 1, 1
-    assert t.uop.st.is_expanded() == (True, True, True)
-
-    t = Tensor.empty(3, 0, 2)
-    assert t.shape == (3, 0, 2)
-    # numpy has stride 0, 0, 0; torch has stride 2, 2, 1
-    assert t.uop.st.is_expanded() == (True, True, True)
-
-    t = Tensor.empty(0, 0, 0)
-    assert t.shape == (0, 0, 0)
-    # numpy has stride 0, 0, 0; torch has stride 1, 1, 1
-    assert t.uop.st.is_expanded() == (True, True, True)
-
  def test_rand(self):
    t = Tensor.rand(3, 2, 0)
    assert t.shape == (3, 2, 0)
@@ -662,8 +621,10 @@ class TestZeroShapeTensor(unittest.TestCase):

    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)
+    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(1, 3).numpy(), [[1, 2, 0]])
+    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(None, 3).numpy(), [[1, 2, 0]])
+    with self.assertRaises(ValueError): Tensor([1, 2]).pad_to(2, 3)
+    with self.assertRaises(ValueError): Tensor([[1, 2]]).pad_to(3)

  def test_shrink_into_zero(self):
    t = Tensor.rand(3, 4).realize()
@@ -849,6 +810,13 @@ class TestTensorMetadata(unittest.TestCase):
    self.assertEqual(len(si.metadata), 1)
    self.assertEqual(si.metadata[0].name, "relu")

+  def test_assign(self):
+    x = Tensor.empty(10, 10).realize()
+    x.assign(Tensor.ones(10, 10).contiguous())
+    si = x.schedule()[-1]
+    self.assertEqual(len(si.metadata), 1)
+    self.assertEqual(si.metadata[0].name, "assign")
+
  def test_complex(self):
    x = Tensor.rand(3, requires_grad=True)
    y = Tensor.rand(3, requires_grad=True)
@@ -860,7 +828,6 @@ class TestTensorMetadata(unittest.TestCase):
    self.assertEqual(len(si.metadata), 3)
    self.assertEqual(set(m.name for m in si.metadata), {"relu", "sigmoid", "__mul__"})

-  @unittest.skip("not accurate")
  def test_complex_backward(self):
    x = Tensor.rand(3, requires_grad=True).realize()
    y = Tensor.rand(3, requires_grad=True).realize()
@@ -872,11 +839,12 @@ 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]
-    self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}")
-    self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "relu"})
+    self.assertEqual(len(si.metadata), 4, f"failed with {si.metadata}")
+    self.assertSetEqual(set(m.name for m in si.metadata), {"__mul__", "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(bw), 2)
+    self.assertEqual(bw[0].name, "__mul__")
+    self.assertEqual(bw[1].name, "sigmoid")

 class TestIdxUpcast(unittest.TestCase):
  def _find_op(self, ast: UOp, op: Ops):
@@ -898,7 +866,8 @@ class TestIdxUpcast(unittest.TestCase):
    store = next(uop for uop in uops if uop.op is Ops.STORE)
    assert store.op is Ops.STORE
    idx = self._find_op(store, Ops.INDEX)
-    if idx is not None: # PTX turns Ops.INDEX into pointer arithmetic earlier than cstyle, plus it's already cast to int64
+    # PTX and NIR turn Ops.INDEX into pointer arithmetic earlier than cstyle, plus it's already cast to int64
+    if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)):
      assert idx.op is Ops.INDEX
      idx_val = idx.src[1]
      assert idx_val.dtype is dtype
@@ -922,7 +891,7 @@ class TestIdxUpcast(unittest.TestCase):
  def test_regular_sym(self):
    self.do_op_then_assert(dtypes.int, 2048, 2048, UOp.variable("dim3", 1, 64).bind(32))

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX always convert Ops.INDEX to int64")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "PTX and NIR always converts Ops.INDEX to int64")
  def test_symfold(self):
    # This would cause an overflow, but after sym fold it's within int32
    a = Tensor.arange(65535)
@@ -11,8 +11,7 @@ class TestTensorUOp(unittest.TestCase):
    def helper(a: np.ndarray):
      print(a.shape, a.strides, a.flags.c_contiguous)
      b = Tensor(a).uop
-      #assert b.st.contiguous == a.flags.c_contiguous
-      assert b.st.shape == a.shape
+      assert b.shape == a.shape
      np.testing.assert_equal(a, Tensor(b).numpy())

    for ndims in range(1, 4):
@@ -93,7 +93,7 @@ class TestTensorVariable(unittest.TestCase):
    vb = v.bind(3)
    t = Tensor.empty(3, vb)
    assert t.uop.base.buffer.size == 30
-    assert t.uop.st.shape == (3, vb)
+    assert t.uop.shape == (3, vb)


 if __name__ == '__main__':
@@ -134,8 +134,8 @@ class TestTiny(unittest.TestCase):
  def test_mnist_backward(self):
    # NOTE: we don't have the whole model here for speed
    layers = [
-      nn.Conv2d(1, 32, 5), Tensor.relu,
-      nn.Conv2d(32, 32, 5), Tensor.relu]
+      nn.Conv2d(1, 8, 5), Tensor.relu,
+      nn.Conv2d(8, 8, 5), Tensor.relu]

    # replace random weights with ones
    # TODO: there's a bug here where it's tying two of the biases together. we need UNIQUE const
@@ -144,7 +144,7 @@ class TestTiny(unittest.TestCase):

    # realize gradients
    for x in nn.state.get_parameters(layers): x.requires_grad_()
-    Tensor.empty(4, 1, 28, 28).sequential(layers).sum().backward()
+    Tensor.empty(4, 1, 14, 14).sequential(layers).sum().backward()
    Tensor.realize(*[x.grad for x in nn.state.get_parameters(layers) if x.grad is not None])

  # *** image ***
@@ -149,6 +149,7 @@ class TestTranscendentalVectorized(unittest.TestCase):
    for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.log2, np.log2, (0.001, 200), vec_size)

  @unittest.skipIf(getenv("DSP"), "requires int division")
+  @unittest.skipIf(getenv("NV_NAK"), "MUFU.SIN is not accurate enough")
  def test_sin_vectorized(self):
    for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.sin, np.sin, (-100, 100), vec_size)

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
-import unittest, os, subprocess, sys
+import unittest, os, subprocess
 from tinygrad import Tensor
-from tinygrad.device import Device, Compiler
+from tinygrad.device import Device, Compiler, enumerate_devices_str
 from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, WIN, CI

 class TestDevice(unittest.TestCase):
@@ -100,10 +100,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=40,)
-    out = (p.stdout + p.stderr).decode()
-    self.assertEqual(p.returncode, 0, msg=out)
+    out = '\n'.join(enumerate_devices_str())
    self.assertIn("CPU", out) # for sanity check

 if __name__ == "__main__":
@@ -418,5 +418,32 @@ class TestPathTensor(unittest.TestCase):
      Tensor(pathlib.Path(test_file)).tolist()
    os.chmod(test_file, 0o644)
    assert Tensor(pathlib.Path(test_file)).tolist(), list(range(10))
+
+class TestDiskTensorMovement(unittest.TestCase):
+  def setUp(self):
+    self.fn = pathlib.Path(temp("custom_disk_range"))
+    self.fn.unlink(missing_ok=True)
+    Tensor.arange(100, dtype=dtypes.uint8).to(f"disk:{str(self.fn)}").realize()
+
+  def test_simple_read(self):
+    t = Tensor(self.fn)
+    self.assertTrue(Tensor.all(t.to(None) == Tensor.arange(100, dtype=dtypes.uint8)).item())
+
+  def test_slice_read(self):
+    t = Tensor(self.fn)
+    self.assertListEqual(t[16:18].tolist(), [16,17])
+
+  # TODO: fix this! at least assert on it
+  @unittest.expectedFailure
+  def test_slice_read_cat(self):
+    t = Tensor(self.fn)
+    self.assertListEqual(Tensor.cat(t[16:18], t[20:22]).tolist(), [16,17,20,21])
+
+  # TODO: fix this! at least assert on it
+  @unittest.expectedFailure
+  def test_slice_sum(self):
+    t = Tensor(self.fn)
+    self.assertListEqual((t[16:18]+t[20:22]).tolist(), [16+20,17+21])
+
 if __name__ == "__main__":
  unittest.main()
@@ -24,6 +24,15 @@ class TestElfLoader(unittest.TestCase):
    '''
    with self.assertRaisesRegex(RuntimeError, 'evil_external_function'):
      ClangJITCompiler().compile(src)
+  def test_link(self):
+    src = '''
+      float powf(float, float); // from libm
+      float test(float x, float y) { return powf(x, y); }
+    '''
+    args = ('-x', 'c', '-c', '-target', f'{platform.machine()}-none-unknown-elf', '-march=native', '-fPIC', '-O2', '-ffreestanding', '-nostdlib')
+    obj = subprocess.check_output(('clang',) + args + ('-', '-o', '-'), input=src.encode())
+    with self.assertRaisesRegex(RuntimeError, 'powf'): elf_loader(obj)
+    elf_loader(obj, link_libs=['m'])

 if __name__ == '__main__':
  unittest.main()
@@ -334,5 +334,19 @@ class TestBidirectional(unittest.TestCase):
    graph_rewrite(c, pm, ctx=ctx_list, bpm=bpm)
    self.assertListEqual(ctx_list, [('+', True), (1, True), (1, False), (2, True), (2, False), ('+', False)])

+class TestStopEarly(unittest.TestCase):
+  def test_stop_early(self):
+    a = UOp.const(dtypes.int, 3)
+    b = UOp.const(dtypes.int, 4)
+    c = a+b
+    cn = UOp.const(dtypes.int, 7)
+    d = UOp.const(dtypes.int, 2)
+    def visit_const(c:UOp):
+      print(f"visit {c.arg}")
+      assert c.arg not in (3,4)
+    pm_cvisit = PatternMatcher([(UPat(Ops.CONST, name="c"), visit_const),])
+    ret = (c+d).substitute({c:cn}, extra_pm=pm_cvisit)
+    assert ret == cn+d
+
 if __name__ == '__main__':
  unittest.main()
@@ -5,8 +5,6 @@ import numpy as np

 from tinygrad import Tensor, dtypes, Device, TinyJit
 from tinygrad.device import is_dtype_supported
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
 from tinygrad.helpers import CI, all_same, prod

 random.seed(42)
@@ -22,11 +20,13 @@ def consec(shape, start=1):
 # creates strided tensor with base set to reference tensor's base, equivalent to torch.set_()
 def set_(reference: Tensor, shape, strides, offset):
  raise NotImplementedError("need to implement without calling uop.view")
+  """
  if reference.uop.base.realized is None: reference.realize()
  assert reference.uop.base.realized, "base has to be realized before setting it to strided's base"
  strided = Tensor(reference.uop.view(ShapeTracker((View.create(shape=shape, strides=strides, offset=offset),))))
  assert strided.uop.st.real_strides() == strides, "real_strides should equal strides for strided"
  return strided
+  """

 def clone(original:Tensor): return original.clone()
 def copy_(src:Tensor, other:Tensor) -> Tensor: return src.clone()
@@ -180,474 +180,6 @@ class TestIndexing(unittest.TestCase):
    # def delitem(): del reference[0]
    # self.assertRaises(TypeError, delitem)

-  # TODO: LLVM is quite fast, why are other compiled backends slow?
-  @unittest.skipIf(CI and Device.DEFAULT in ["CPU", "CL", "METAL", "NV", "AMD"], "slow")
-  def test_advancedindex(self):
-    # integer array indexing
-
-    # pick a random valid indexer type
-    def ri(indices):
-      choice = random.randint(0, 2)
-      if choice == 0: return Tensor(indices)
-      if choice == 1: return list(indices)
-      return tuple(indices)
-
-    def validate_indexing(x):
-      numpy_testing_assert_equal_helper(x[[0]], consec((1,)))
-      numpy_testing_assert_equal_helper(x[ri([0]),], consec((1,)))
-      numpy_testing_assert_equal_helper(x[ri([3]),], consec((1,), 4))
-      numpy_testing_assert_equal_helper(x[[2, 3, 4]], consec((3,), 3))
-      numpy_testing_assert_equal_helper(x[ri([2, 3, 4]),], consec((3,), 3))
-      numpy_testing_assert_equal_helper(x[ri([0, 2, 4]),], np.array([1, 3, 5]))
-
-    def validate_setting(x):
-      x[[0]] = -2
-      numpy_testing_assert_equal_helper(x[[0]], np.array([-2]))
-      x[[0]] = -1
-      numpy_testing_assert_equal_helper(x[ri([0]), ], np.array([-1]))
-      x[[2, 3, 4]] = 4
-      numpy_testing_assert_equal_helper(x[[2, 3, 4]], np.array([4, 4, 4]))
-      x[ri([2, 3, 4]), ] = 3
-      numpy_testing_assert_equal_helper(x[ri([2, 3, 4]), ], np.array([3, 3, 3]))
-      x[ri([0, 2, 4]), ] = Tensor([5, 4, 3])
-      numpy_testing_assert_equal_helper(x[ri([0, 2, 4]), ], np.array([5, 4, 3]))
-
-    # Case 1: Purely Integer Array Indexing
-    reference = consec((10,))
-    validate_indexing(reference)
-    # setting values
-    validate_setting(reference)
-
-    # Tensor with stride != 1
-    # strided is [1, 3, 5, 7]
-
-    # # TODO: set stride
-    # reference = consec((10,))
-    # strided = set_(reference, (4,), (2,), 0)
-
-    # numpy_testing_assert_equal_helper(strided[[0]], np.array([1]))
-    # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([1]))
-    # numpy_testing_assert_equal_helper(strided[ri([3]), ], np.array([7]))
-    # numpy_testing_assert_equal_helper(strided[[1, 2]], np.array([3, 5]))
-    # numpy_testing_assert_equal_helper(strided[ri([1, 2]), ], np.array([3, 5]))
-    # numpy_testing_assert_equal_helper(strided[ri([[2, 1], [0, 3]]), ],
-    #                   np.array([[5, 3], [1, 7]]))
-
-    # stride is [4, 8]
-
-    # strided = set_(reference, (2,), (4,), offset=4)
-
-    # numpy_testing_assert_equal_helper(strided[[0]], np.array([5]))
-    # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([5]))
-    # numpy_testing_assert_equal_helper(strided[ri([1]), ], np.array([9]))
-    # numpy_testing_assert_equal_helper(strided[[0, 1]], np.array([5, 9]))
-    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ], np.array([5, 9]))
-    # numpy_testing_assert_equal_helper(strided[ri([[0, 1], [1, 0]]), ],
-    #                   np.array([[5, 9], [9, 5]]))
-
-    # reference is 1 2
-    #              3 4
-    #              5 6
-    reference = consec((3, 2))
-    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], np.array([1, 3, 5]))
-    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([1])], np.array([2, 4, 6]))
-    numpy_testing_assert_equal_helper(reference[ri([0]), ri([0])], consec((1,)))
-    numpy_testing_assert_equal_helper(reference[ri([2]), ri([1])], consec((1,), 6))
-    numpy_testing_assert_equal_helper(reference[[ri([0, 0]), ri([0, 1])]], np.array([1, 2]))
-    numpy_testing_assert_equal_helper(reference[[ri([0, 1, 1, 0, 2]), ri([1])]], np.array([2, 4, 4, 2, 6]))
-    numpy_testing_assert_equal_helper(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 2, 3, 3]))
-
-    rows = ri([[0, 0],
-               [1, 2]])
-    columns = [0],
-    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[1, 1],
-                                                                          [3, 5]]))
-
-    rows = ri([[0, 0],
-               [1, 2]])
-    columns = ri([1, 0])
-    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[2, 1],
-                                                                          [4, 5]]))
-    rows = ri([[0, 0],
-               [1, 2]])
-    columns = ri([[0, 1],
-                  [1, 0]])
-    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[1, 2],
-                                                                          [4, 5]]))
-
-    # setting values
-    reference[ri([0]), ri([1])] = -1
-    numpy_testing_assert_equal_helper(reference[ri([0]), ri([1])], np.array([-1]))
-    reference[ri([0, 1, 2]), ri([0])] = Tensor([-1, 2, -4])
-    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])],
-                      np.array([-1, 2, -4]))
-    reference[rows, columns] = Tensor([[4, 6], [2, 3]])
-    numpy_testing_assert_equal_helper(reference[rows, columns],
-                      np.array([[4, 6], [2, 3]]))
-
-    # Verify still works with Transposed (i.e. non-contiguous) Tensors
-    reference = Tensor([[0, 1, 2, 3],
-                        [4, 5, 6, 7],
-                        [8, 9, 10, 11]]).T
-
-    # Transposed: [[0, 4, 8],
-    #              [1, 5, 9],
-    #              [2, 6, 10],
-    #              [3, 7, 11]]
-
-    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], np.array([0, 1, 2]))
-    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([1])], np.array([4, 5, 6]))
-    numpy_testing_assert_equal_helper(reference[ri([0]), ri([0])], np.array([0]))
-    numpy_testing_assert_equal_helper(reference[ri([2]), ri([1])], np.array([6]))
-    numpy_testing_assert_equal_helper(reference[[ri([0, 0]), ri([0, 1])]], np.array([0, 4]))
-    numpy_testing_assert_equal_helper(reference[[ri([0, 1, 1, 0, 3]), ri([1])]], np.array([4, 5, 5, 4, 7]))
-    numpy_testing_assert_equal_helper(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([0, 4, 1, 1]))
-
-    rows = ri([[0, 0],
-               [1, 2]])
-    columns = [0],
-    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[0, 0], [1, 2]]))
-
-    rows = ri([[0, 0],
-               [1, 2]])
-    columns = ri([1, 0])
-    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[4, 0], [5, 2]]))
-    rows = ri([[0, 0],
-               [1, 3]])
-    columns = ri([[0, 1],
-                  [1, 2]])
-    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[0, 4], [5, 11]]))
-
-    # TODO: non contiguous setitem
-    '''
-    # setting values
-    reference[ri([0]), ri([1])] = -1
-    numpy_testing_assert_equal_helper(reference[ri([0]), ri([1])],
-                      np.array([-1]))
-    reference[ri([0, 1, 2]), ri([0])] = np.array([-1, 2, -4])
-    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])],
-                      np.array([-1, 2, -4]))
-    reference[rows, columns] = np.array([[4, 6], [2, 3]])
-    numpy_testing_assert_equal_helper(reference[rows, columns],
-                      np.array([[4, 6], [2, 3]]))
-    '''
-
-    # stride != 1
-
-    # strided is [[1 3 5 7],
-    #             [9 11 13 15]]
-
-    # # TODO: set stride
-    # reference = Tensor.arange(0., 24).reshape(3, 8)
-    # strided = set_(reference, (2,4), (8,2), 1)
-
-    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([0])], np.array([1, 9]))
-    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1])], np.array([3, 11]))
-    # numpy_testing_assert_equal_helper(strided[ri([0]), ri([0])], np.array([1]))
-    # numpy_testing_assert_equal_helper(strided[ri([1]), ri([3])], np.array([15]))
-    # numpy_testing_assert_equal_helper(strided[[ri([0, 0]), ri([0, 3])]], np.array([1, 7]))
-    # numpy_testing_assert_equal_helper(strided[[ri([1]), ri([0, 1, 1, 0, 3])]], np.array([9, 11, 11, 9, 15]))
-    # numpy_testing_assert_equal_helper(strided[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 3, 9, 9]))
-
-    # rows = ri([[0, 0],
-    #            [1, 1]])
-    # columns = [0],
-    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 1], [9, 9]]))
-
-    # rows = ri([[0, 1],
-    #            [1, 0]])
-    # columns = ri([1, 2])
-    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[3, 13], [11, 5]]))
-    # rows = ri([[0, 0],
-    #            [1, 1]])
-    # columns = ri([[0, 1],
-    #               [1, 2]])
-    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 3], [11, 13]]))
-
-    # setting values
-
-    # strided is [[10, 11],
-    #             [17, 18]]
-
-    # # TODO: set stride
-    # reference = Tensor.arange(0., 24).reshape(3, 8)
-    # strided = set_(reference, (2,2), (7,1), 10)
-
-    # numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])], np.array([11]))
-
-    # TODO non contiguous setitem
-    '''
-    strided[ri([0]), ri([1])] = -1
-    numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])],
-                      Tensor([-1]))
-    '''
-    # # TODO: set stride
-    # reference = Tensor.arange(0., 24).reshape(3, 8)
-    # strided = set_(reference, (2,2), (7,1), 10)
-
-    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])], np.array([11, 17]))
-
-    # TODO non contiguous setitem
-    '''
-    strided[ri([0, 1]), ri([1, 0])] = Tensor([-1, 2])
-    numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])],
-                      Tensor([-1, 2]))
-    '''
-
-    # # TODO: set stride
-    # reference = Tensor.arange(0., 24).realize().reshape(3, 8)
-    # strided = set_(reference, (2,2), (7,1), 10)
-
-    # rows = ri([[0],
-    #             [1]])
-    # columns = ri([[0, 1],
-    #               [0, 1]])
-    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[10, 11], [17, 18]]))
-
-    # TODO non contiguous setitem
-    '''
-    strided[rows, columns] = Tensor([[4, 6], [2, 3]])
-    numpy_testing_assert_equal_helper(strided[rows, columns],
-                      Tensor([[4, 6], [2, 3]]))
-    '''
-
-    # Tests using less than the number of dims, and ellipsis
-
-    # reference is 1 2
-    #              3 4
-    #              5 6
-    reference = consec((3, 2))
-    numpy_testing_assert_equal_helper(reference[ri([0, 2]),], np.array([[1, 2], [5, 6]]))
-    numpy_testing_assert_equal_helper(reference[ri([1]), ...], np.array([[3, 4]]))
-    numpy_testing_assert_equal_helper(reference[..., ri([1])], np.array([[2], [4], [6]]))
-
-    # verify too many indices fails
-    with self.assertRaises(IndexError): reference[ri([1]), ri([0, 2]), ri([3])]
-
-    # test invalid index fails
-    reference = Tensor.empty(10)
-    for err_idx in (10, -11):
-      with self.assertRaises(IndexError):
-        reference[err_idx]
-      # NOTE cannot check for out of bounds with Tensor indexing
-      # see tensor.py: __getitem__ (Tiny Things)
-      '''
-      with self.assertRaises(IndexError):
-        reference[Tensor([err_idx], dtype=dtypes.int64)]
-      with self.assertRaises(IndexError):
-        reference[[err_idx]]
-      '''
-
-    def tensor_indices_to_np(tensor: Tensor, indices):
-      npt = tensor.numpy()
-      idxs = tuple(i.numpy().tolist() if isinstance(i, Tensor) and i.dtype == dtypes.int64 else
-                  i for i in indices)
-      return npt, idxs
-
-    def get_numpy(tensor, indices):
-      npt, idxs = tensor_indices_to_np(tensor, indices)
-      return Tensor(npt[idxs])
-
-    def set_numpy(tensor:Tensor, indices, value):
-      if not isinstance(value, int):
-        value = value.numpy()
-      npt, idxs = tensor_indices_to_np(tensor, indices)
-      npt[idxs] = value
-      return npt
-
-    def assert_get_eq(tensor, indexer):
-      numpy_testing_assert_equal_helper(tensor[indexer], get_numpy(tensor, indexer))
-
-    def assert_set_eq(tensor: Tensor, indexer, val):
-      pyt = clone(tensor)
-      numt = clone(tensor)
-      pyt[indexer] = val
-      numt = set_numpy(numt, indexer, val)
-      numpy_testing_assert_equal_helper(pyt, numt)
-
-    # NOTE: torch initiates the gradients using g0cpu (rand as gradients)
-    def assert_backward_eq(tensor: Tensor, indexer):
-      cpu = clone(tensor.float())
-      cpu.requires_grad = True
-      outcpu = cpu[indexer].sum()
-      outcpu.backward()
-      dev = cpu.detach()
-      dev.requires_grad = True
-      outdev = dev[indexer].sum()
-      outdev.backward()
-      numpy_testing_assert_equal_helper(cpu.grad, dev.grad)
-
-    def get_set_tensor(indexed: Tensor, indexer):
-      set_size = indexed[indexer].shape
-      set_count = indexed[indexer].numel()
-      set_tensor = Tensor.randint(set_count, high=set_count).reshape(set_size) #.cast(dtypes.float64)
-      return set_tensor
-
-    # Tensor is  0  1  2  3  4
-    #            5  6  7  8  9
-    #           10 11 12 13 14
-    #           15 16 17 18 19
-    reference = Tensor.arange(0., 20).reshape(4, 5)
-
-    indices_to_test = [
-      # grab the second, fourth columns
-      [slice(None), [1, 3]],
-
-      # first, third rows,
-      [[0, 2], slice(None)],
-
-      # weird shape
-      [slice(None), [[0, 1],
-                      [2, 3]]],
-      # negatives
-      [[-1], [0]],
-      [[0, 2], [-1]],
-      [slice(None), [-1]],
-    ]
-
-    # only test dupes on gets
-    get_indices_to_test = indices_to_test + [[slice(None), [0, 1, 1, 2, 2]]]
-
-    for indexer in get_indices_to_test:
-      assert_get_eq(reference, indexer)
-      assert_backward_eq(reference, indexer)
-
-    for indexer in indices_to_test:
-      assert_set_eq(reference, indexer, 44)
-      assert_set_eq(reference, indexer, get_set_tensor(reference, indexer))
-
-    reference = Tensor.arange(0., 160).reshape(4, 8, 5)
-
-    indices_to_test = [
-      [slice(None), slice(None), [0, 3, 4]],
-      [slice(None), [2, 4, 5, 7], slice(None)],
-      [[2, 3], slice(None), slice(None)],
-      [slice(None), [0, 2, 3], [1, 3, 4]],
-      [slice(None), [0], [1, 2, 4]],
-      [slice(None), [0, 1, 3], [4]],
-      [slice(None), [[0, 1], [1, 0]], [[2, 3]]],
-      [slice(None), [[0, 1], [2, 3]], [[0]]],
-      [slice(None), [[5, 6]], [[0, 3], [4, 4]]],
-      [[0, 2, 3], [1, 3, 4], slice(None)],
-      [[0], [1, 2, 4], slice(None)],
-      [[0, 1, 3], [4], slice(None)],
-      [[[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None)],
-      [[[0, 1], [1, 0]], [[2, 3]], slice(None)],
-      [[[0, 1], [2, 3]], [[0]], slice(None)],
-      [[[2, 1]], [[0, 3], [4, 4]], slice(None)],
-      [[[2]], [[0, 3], [4, 1]], slice(None)],
-      # non-contiguous indexing subspace
-      [[0, 2, 3], slice(None), [1, 3, 4]],
-
-      # less dim, ellipsis
-      [[0, 2], ],
-      [[0, 2], slice(None)],
-      [[0, 2], Ellipsis],
-      [[0, 2], slice(None), Ellipsis],
-      [[0, 2], Ellipsis, slice(None)],
-      [[0, 2], [1, 3]],
-      [[0, 2], [1, 3], Ellipsis],
-      [Ellipsis, [1, 3], [2, 3]],
-      [Ellipsis, [2, 3, 4]],
-      [Ellipsis, slice(None), [2, 3, 4]],
-      [slice(None), Ellipsis, [2, 3, 4]],
-
-      # ellipsis counts for nothing
-      [Ellipsis, slice(None), slice(None), [0, 3, 4]],
-      [slice(None), Ellipsis, slice(None), [0, 3, 4]],
-      [slice(None), slice(None), Ellipsis, [0, 3, 4]],
-      [slice(None), slice(None), [0, 3, 4], Ellipsis],
-      [Ellipsis, [[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None)],
-      [[[0, 1], [1, 0]], [[2, 1], [3, 5]], Ellipsis, slice(None)],
-      [[[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None), Ellipsis],
-    ]
-
-    for indexer in indices_to_test:
-      assert_get_eq(reference, indexer)
-
-      assert_set_eq(reference, indexer, 212)
-      assert_set_eq(reference, indexer, get_set_tensor(reference, indexer))
-      assert_backward_eq(reference, indexer)
-
-    reference = Tensor.arange(0., 1296).reshape(3, 9, 8, 6)
-
-    indices_to_test = [
-      [slice(None), slice(None), slice(None), [0, 3, 4]],
-      [slice(None), slice(None), [2, 4, 5, 7], slice(None)],
-      [slice(None), [2, 3], slice(None), slice(None)],
-      [[1, 2], slice(None), slice(None), slice(None)],
-      [slice(None), slice(None), [0, 2, 3], [1, 3, 4]],
-      [slice(None), slice(None), [0], [1, 2, 4]],
-      [slice(None), slice(None), [0, 1, 3], [4]],
-      [slice(None), slice(None), [[0, 1], [1, 0]], [[2, 3]]],
-      [slice(None), slice(None), [[0, 1], [2, 3]], [[0]]],
-      [slice(None), slice(None), [[5, 6]], [[0, 3], [4, 4]]],
-      [slice(None), [0, 2, 3], [1, 3, 4], slice(None)],
-      [slice(None), [0], [1, 2, 4], slice(None)],
-      [slice(None), [0, 1, 3], [4], slice(None)],
-      [slice(None), [[0, 1], [3, 4]], [[2, 3], [0, 1]], slice(None)],
-      [slice(None), [[0, 1], [3, 4]], [[2, 3]], slice(None)],
-      [slice(None), [[0, 1], [3, 2]], [[0]], slice(None)],
-      [slice(None), [[2, 1]], [[0, 3], [6, 4]], slice(None)],
-      [slice(None), [[2]], [[0, 3], [4, 2]], slice(None)],
-      [[0, 1, 2], [1, 3, 4], slice(None), slice(None)],
-      [[0], [1, 2, 4], slice(None), slice(None)],
-      [[0, 1, 2], [4], slice(None), slice(None)],
-      [[[0, 1], [0, 2]], [[2, 4], [1, 5]], slice(None), slice(None)],
-      [[[0, 1], [1, 2]], [[2, 0]], slice(None), slice(None)],
-      [[[2, 2]], [[0, 3], [4, 5]], slice(None), slice(None)],
-      [[[2]], [[0, 3], [4, 5]], slice(None), slice(None)],
-      [slice(None), [3, 4, 6], [0, 2, 3], [1, 3, 4]],
-      [slice(None), [2, 3, 4], [1, 3, 4], [4]],
-      [slice(None), [0, 1, 3], [4], [1, 3, 4]],
-      [slice(None), [6], [0, 2, 3], [1, 3, 4]],
-      [slice(None), [2, 3, 5], [3], [4]],
-      [slice(None), [0], [4], [1, 3, 4]],
-      [slice(None), [6], [0, 2, 3], [1]],
-      [slice(None), [[0, 3], [3, 6]], [[0, 1], [1, 3]], [[5, 3], [1, 2]]],
-      [[2, 2, 1], [0, 2, 3], [1, 3, 4], slice(None)],
-      [[2, 0, 1], [1, 2, 3], [4], slice(None)],
-      [[0, 1, 2], [4], [1, 3, 4], slice(None)],
-      [[0], [0, 2, 3], [1, 3, 4], slice(None)],
-      [[0, 2, 1], [3], [4], slice(None)],
-      [[0], [4], [1, 3, 4], slice(None)],
-      [[1], [0, 2, 3], [1], slice(None)],
-      [[[1, 2], [1, 2]], [[0, 1], [2, 3]], [[2, 3], [3, 5]], slice(None)],
-
-      # less dim, ellipsis
-      [Ellipsis, [0, 3, 4]],
-      [Ellipsis, slice(None), [0, 3, 4]],
-      [Ellipsis, slice(None), slice(None), [0, 3, 4]],
-      [slice(None), Ellipsis, [0, 3, 4]],
-      [slice(None), slice(None), Ellipsis, [0, 3, 4]],
-      [slice(None), [0, 2, 3], [1, 3, 4]],
-      [slice(None), [0, 2, 3], [1, 3, 4], Ellipsis],
-      [Ellipsis, [0, 2, 3], [1, 3, 4], slice(None)],
-      [[0], [1, 2, 4]],
-      [[0], [1, 2, 4], slice(None)],
-      [[0], [1, 2, 4], Ellipsis],
-      [[0], [1, 2, 4], Ellipsis, slice(None)],
-      [[1], ],
-      [[0, 2, 1], [3], [4]],
-      [[0, 2, 1], [3], [4], slice(None)],
-      [[0, 2, 1], [3], [4], Ellipsis],
-      [Ellipsis, [0, 2, 1], [3], [4]],
-    ]
-
-    for indexer in indices_to_test:
-      assert_get_eq(reference, indexer)
-      assert_set_eq(reference, indexer, 1333)
-      assert_set_eq(reference, indexer, get_set_tensor(reference, indexer))
-
-    indices_to_test += [
-      [slice(None), slice(None), [[0, 1], [1, 0]], [[2, 3], [3, 0]]],
-      [slice(None), slice(None), [[2]], [[0, 3], [4, 4]]],
-    ]
-    for indexer in indices_to_test:
-      assert_get_eq(reference, indexer)
-      assert_set_eq(reference, indexer, 1333)
-      assert_backward_eq(reference, indexer)
-
  # TODO setitem backward
  '''
  def test_set_item_to_scalar_tensor(self):
@@ -969,10 +501,6 @@ class TestIndexing(unittest.TestCase):
    y = x[:, :, :, 1]
    z = y[:, 1:1, :]
    numpy_testing_assert_equal_helper((2, 0, 4), z.shape)
-    # this isn't technically necessary, but matches NumPy stride calculations.
-    # NOTE: this is empty and shouldn't have strides
-    numpy_testing_assert_equal_helper((True, True, True), z.uop.st.is_expanded())
-    self.assertTrue(z.uop.st.contiguous)

  @unittest.skip("bool indexing not supported")
  def test_index_getitem_copy_bools_slices(self):
@@ -1568,5 +1096,474 @@ class TestNumpy(unittest.TestCase):
    numpy_testing_assert_equal_helper(kernel, kernel2)
  '''

+def tensor_indices_to_np(tensor: Tensor, indices):
+  npt = tensor.numpy()
+  idxs = tuple(i.numpy().tolist() if isinstance(i, Tensor) and i.dtype == dtypes.int64 else
+              i for i in indices)
+  return npt, idxs
+
+def get_numpy(tensor, indices):
+  npt, idxs = tensor_indices_to_np(tensor, indices)
+  return Tensor(npt[idxs])
+
+def set_numpy(tensor:Tensor, indices, value):
+  if not isinstance(value, int):
+    value = value.numpy()
+  npt, idxs = tensor_indices_to_np(tensor, indices)
+  npt[idxs] = value
+  return npt
+
+def assert_get_eq(tensor, indexer):
+  numpy_testing_assert_equal_helper(tensor[indexer], get_numpy(tensor, indexer))
+
+def assert_set_eq(tensor: Tensor, indexer, val):
+  pyt = clone(tensor)
+  numt = clone(tensor)
+  pyt[indexer] = val
+  numt = set_numpy(numt, indexer, val)
+  numpy_testing_assert_equal_helper(pyt, numt)
+
+# NOTE: torch initiates the gradients using g0cpu (rand as gradients)
+def assert_backward_eq(tensor: Tensor, indexer):
+  cpu = clone(tensor.float())
+  cpu.requires_grad = True
+  outcpu = cpu[indexer].sum()
+  outcpu.backward()
+  dev = cpu.detach()
+  dev.requires_grad = True
+  outdev = dev[indexer].sum()
+  outdev.backward()
+  numpy_testing_assert_equal_helper(cpu.grad, dev.grad)
+
+def get_set_tensor(indexed: Tensor, indexer):
+  set_size = indexed[indexer].shape
+  set_count = indexed[indexer].numel()
+  set_tensor = Tensor.randint(set_count, high=set_count).reshape(set_size) #.cast(dtypes.float64)
+  return set_tensor
+
+@unittest.skipIf(CI and Device.DEFAULT in ["CPU", "CL", "METAL", "NV", "AMD"], "slow")
+class TestAdvancedIndexing(unittest.TestCase):
+  def test_integer_array_indexing(self):
+    # pick a random valid indexer type
+    def ri(indices):
+      choice = random.randint(0, 2)
+      if choice == 0: return Tensor(indices)
+      if choice == 1: return list(indices)
+      return tuple(indices)
+
+    def validate_indexing(x):
+      numpy_testing_assert_equal_helper(x[[0]], consec((1,)))
+      numpy_testing_assert_equal_helper(x[ri([0]),], consec((1,)))
+      numpy_testing_assert_equal_helper(x[ri([3]),], consec((1,), 4))
+      numpy_testing_assert_equal_helper(x[[2, 3, 4]], consec((3,), 3))
+      numpy_testing_assert_equal_helper(x[ri([2, 3, 4]),], consec((3,), 3))
+      numpy_testing_assert_equal_helper(x[ri([0, 2, 4]),], np.array([1, 3, 5]))
+
+    def validate_setting(x):
+      x[[0]] = -2
+      numpy_testing_assert_equal_helper(x[[0]], np.array([-2]))
+      x[[0]] = -1
+      numpy_testing_assert_equal_helper(x[ri([0]), ], np.array([-1]))
+      x[[2, 3, 4]] = 4
+      numpy_testing_assert_equal_helper(x[[2, 3, 4]], np.array([4, 4, 4]))
+      x[ri([2, 3, 4]), ] = 3
+      numpy_testing_assert_equal_helper(x[ri([2, 3, 4]), ], np.array([3, 3, 3]))
+      x[ri([0, 2, 4]), ] = Tensor([5, 4, 3])
+      numpy_testing_assert_equal_helper(x[ri([0, 2, 4]), ], np.array([5, 4, 3]))
+
+    # Case 1: Purely Integer Array Indexing
+    reference = consec((10,))
+    validate_indexing(reference)
+    # setting values
+    validate_setting(reference)
+
+    # Tensor with stride != 1
+    # strided is [1, 3, 5, 7]
+
+    # # TODO: set stride
+    # reference = consec((10,))
+    # strided = set_(reference, (4,), (2,), 0)
+
+    # numpy_testing_assert_equal_helper(strided[[0]], np.array([1]))
+    # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([1]))
+    # numpy_testing_assert_equal_helper(strided[ri([3]), ], np.array([7]))
+    # numpy_testing_assert_equal_helper(strided[[1, 2]], np.array([3, 5]))
+    # numpy_testing_assert_equal_helper(strided[ri([1, 2]), ], np.array([3, 5]))
+    # numpy_testing_assert_equal_helper(strided[ri([[2, 1], [0, 3]]), ],
+    #                   np.array([[5, 3], [1, 7]]))
+
+    # stride is [4, 8]
+
+    # strided = set_(reference, (2,), (4,), offset=4)
+
+    # numpy_testing_assert_equal_helper(strided[[0]], np.array([5]))
+    # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([5]))
+    # numpy_testing_assert_equal_helper(strided[ri([1]), ], np.array([9]))
+    # numpy_testing_assert_equal_helper(strided[[0, 1]], np.array([5, 9]))
+    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ], np.array([5, 9]))
+    # numpy_testing_assert_equal_helper(strided[ri([[0, 1], [1, 0]]), ],
+    #                   np.array([[5, 9], [9, 5]]))
+
+    # reference is 1 2
+    #              3 4
+    #              5 6
+    reference = consec((3, 2))
+    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], np.array([1, 3, 5]))
+    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([1])], np.array([2, 4, 6]))
+    numpy_testing_assert_equal_helper(reference[ri([0]), ri([0])], consec((1,)))
+    numpy_testing_assert_equal_helper(reference[ri([2]), ri([1])], consec((1,), 6))
+    numpy_testing_assert_equal_helper(reference[[ri([0, 0]), ri([0, 1])]], np.array([1, 2]))
+    numpy_testing_assert_equal_helper(reference[[ri([0, 1, 1, 0, 2]), ri([1])]], np.array([2, 4, 4, 2, 6]))
+    numpy_testing_assert_equal_helper(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 2, 3, 3]))
+
+    rows = ri([[0, 0],
+               [1, 2]])
+    columns = [0],
+    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[1, 1],
+                                                                          [3, 5]]))
+
+    rows = ri([[0, 0],
+               [1, 2]])
+    columns = ri([1, 0])
+    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[2, 1],
+                                                                          [4, 5]]))
+    rows = ri([[0, 0],
+               [1, 2]])
+    columns = ri([[0, 1],
+                  [1, 0]])
+    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[1, 2],
+                                                                          [4, 5]]))
+
+    # setting values
+    reference[ri([0]), ri([1])] = -1
+    numpy_testing_assert_equal_helper(reference[ri([0]), ri([1])], np.array([-1]))
+    reference[ri([0, 1, 2]), ri([0])] = Tensor([-1, 2, -4])
+    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])],
+                      np.array([-1, 2, -4]))
+    reference[rows, columns] = Tensor([[4, 6], [2, 3]])
+    numpy_testing_assert_equal_helper(reference[rows, columns],
+                      np.array([[4, 6], [2, 3]]))
+
+    # Verify still works with Transposed (i.e. non-contiguous) Tensors
+    reference = Tensor([[0, 1, 2, 3],
+                        [4, 5, 6, 7],
+                        [8, 9, 10, 11]]).T
+
+    # Transposed: [[0, 4, 8],
+    #              [1, 5, 9],
+    #              [2, 6, 10],
+    #              [3, 7, 11]]
+
+    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], np.array([0, 1, 2]))
+    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([1])], np.array([4, 5, 6]))
+    numpy_testing_assert_equal_helper(reference[ri([0]), ri([0])], np.array([0]))
+    numpy_testing_assert_equal_helper(reference[ri([2]), ri([1])], np.array([6]))
+    numpy_testing_assert_equal_helper(reference[[ri([0, 0]), ri([0, 1])]], np.array([0, 4]))
+    numpy_testing_assert_equal_helper(reference[[ri([0, 1, 1, 0, 3]), ri([1])]], np.array([4, 5, 5, 4, 7]))
+    numpy_testing_assert_equal_helper(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([0, 4, 1, 1]))
+
+    rows = ri([[0, 0],
+               [1, 2]])
+    columns = [0],
+    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[0, 0], [1, 2]]))
+
+    rows = ri([[0, 0],
+               [1, 2]])
+    columns = ri([1, 0])
+    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[4, 0], [5, 2]]))
+    rows = ri([[0, 0],
+               [1, 3]])
+    columns = ri([[0, 1],
+                  [1, 2]])
+    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[0, 4], [5, 11]]))
+
+    # TODO: non contiguous setitem
+    '''
+    # setting values
+    reference[ri([0]), ri([1])] = -1
+    numpy_testing_assert_equal_helper(reference[ri([0]), ri([1])],
+                      np.array([-1]))
+    reference[ri([0, 1, 2]), ri([0])] = np.array([-1, 2, -4])
+    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])],
+                      np.array([-1, 2, -4]))
+    reference[rows, columns] = np.array([[4, 6], [2, 3]])
+    numpy_testing_assert_equal_helper(reference[rows, columns],
+                      np.array([[4, 6], [2, 3]]))
+    '''
+
+    # stride != 1
+
+    # strided is [[1 3 5 7],
+    #             [9 11 13 15]]
+
+    # # TODO: set stride
+    # reference = Tensor.arange(0., 24).reshape(3, 8)
+    # strided = set_(reference, (2,4), (8,2), 1)
+
+    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([0])], np.array([1, 9]))
+    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1])], np.array([3, 11]))
+    # numpy_testing_assert_equal_helper(strided[ri([0]), ri([0])], np.array([1]))
+    # numpy_testing_assert_equal_helper(strided[ri([1]), ri([3])], np.array([15]))
+    # numpy_testing_assert_equal_helper(strided[[ri([0, 0]), ri([0, 3])]], np.array([1, 7]))
+    # numpy_testing_assert_equal_helper(strided[[ri([1]), ri([0, 1, 1, 0, 3])]], np.array([9, 11, 11, 9, 15]))
+    # numpy_testing_assert_equal_helper(strided[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 3, 9, 9]))
+
+    # rows = ri([[0, 0],
+    #            [1, 1]])
+    # columns = [0],
+    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 1], [9, 9]]))
+
+    # rows = ri([[0, 1],
+    #            [1, 0]])
+    # columns = ri([1, 2])
+    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[3, 13], [11, 5]]))
+    # rows = ri([[0, 0],
+    #            [1, 1]])
+    # columns = ri([[0, 1],
+    #               [1, 2]])
+    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 3], [11, 13]]))
+
+    # setting values
+
+    # strided is [[10, 11],
+    #             [17, 18]]
+
+    # # TODO: set stride
+    # reference = Tensor.arange(0., 24).reshape(3, 8)
+    # strided = set_(reference, (2,2), (7,1), 10)
+
+    # numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])], np.array([11]))
+
+    # TODO non contiguous setitem
+    '''
+    strided[ri([0]), ri([1])] = -1
+    numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])],
+                      Tensor([-1]))
+    '''
+    # # TODO: set stride
+    # reference = Tensor.arange(0., 24).reshape(3, 8)
+    # strided = set_(reference, (2,2), (7,1), 10)
+
+    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])], np.array([11, 17]))
+
+    # TODO non contiguous setitem
+    '''
+    strided[ri([0, 1]), ri([1, 0])] = Tensor([-1, 2])
+    numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])],
+                      Tensor([-1, 2]))
+    '''
+
+    # # TODO: set stride
+    # reference = Tensor.arange(0., 24).realize().reshape(3, 8)
+    # strided = set_(reference, (2,2), (7,1), 10)
+
+    # rows = ri([[0],
+    #             [1]])
+    # columns = ri([[0, 1],
+    #               [0, 1]])
+    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[10, 11], [17, 18]]))
+
+    # TODO non contiguous setitem
+    '''
+    strided[rows, columns] = Tensor([[4, 6], [2, 3]])
+    numpy_testing_assert_equal_helper(strided[rows, columns],
+                      Tensor([[4, 6], [2, 3]]))
+    '''
+
+    # Tests using less than the number of dims, and ellipsis
+
+    # reference is 1 2
+    #              3 4
+    #              5 6
+    reference = consec((3, 2))
+    numpy_testing_assert_equal_helper(reference[ri([0, 2]),], np.array([[1, 2], [5, 6]]))
+    numpy_testing_assert_equal_helper(reference[ri([1]), ...], np.array([[3, 4]]))
+    numpy_testing_assert_equal_helper(reference[..., ri([1])], np.array([[2], [4], [6]]))
+
+    # verify too many indices fails
+    with self.assertRaises(IndexError): reference[ri([1]), ri([0, 2]), ri([3])]
+
+    # test invalid index fails
+    reference = Tensor.empty(10)
+    for err_idx in (10, -11):
+      with self.assertRaises(IndexError):
+        reference[err_idx]
+      # NOTE cannot check for out of bounds with Tensor indexing
+      # see tensor.py: __getitem__ (Tiny Things)
+      '''
+      with self.assertRaises(IndexError):
+        reference[Tensor([err_idx], dtype=dtypes.int64)]
+      with self.assertRaises(IndexError):
+        reference[[err_idx]]
+      '''
+
+  def test_numpy_parity_and_backward_2d(self):
+    # Tensor is  0  1  2  3  4
+    #            5  6  7  8  9
+    #           10 11 12 13 14
+    #           15 16 17 18 19
+    reference = Tensor.arange(0., 20).reshape(4, 5)
+
+    indices_to_test = [
+      # grab the second, fourth columns
+      [slice(None), [1, 3]],
+
+      # first, third rows,
+      [[0, 2], slice(None)],
+
+      # weird shape
+      [slice(None), [[0, 1],
+                      [2, 3]]],
+      # negatives
+      [[-1], [0]],
+      [[0, 2], [-1]],
+      [slice(None), [-1]],
+    ]
+
+    # only test dupes on gets
+    get_indices_to_test = indices_to_test + [[slice(None), [0, 1, 1, 2, 2]]]
+
+    for indexer in get_indices_to_test:
+      assert_get_eq(reference, indexer)
+      assert_backward_eq(reference, indexer)
+
+    for indexer in indices_to_test:
+      assert_set_eq(reference, indexer, 44)
+      assert_set_eq(reference, indexer, get_set_tensor(reference, indexer))
+
+  def test_numpy_parity_and_backward_3d(self):
+    reference = Tensor.arange(0., 160).reshape(4, 8, 5)
+
+    indices_to_test = [
+      [slice(None), slice(None), [0, 3, 4]],
+      [slice(None), [2, 4, 5, 7], slice(None)],
+      [[2, 3], slice(None), slice(None)],
+      [slice(None), [0, 2, 3], [1, 3, 4]],
+      [slice(None), [0], [1, 2, 4]],
+      [slice(None), [0, 1, 3], [4]],
+      [slice(None), [[0, 1], [1, 0]], [[2, 3]]],
+      [slice(None), [[0, 1], [2, 3]], [[0]]],
+      [slice(None), [[5, 6]], [[0, 3], [4, 4]]],
+      [[0, 2, 3], [1, 3, 4], slice(None)],
+      [[0], [1, 2, 4], slice(None)],
+      [[0, 1, 3], [4], slice(None)],
+      [[[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None)],
+      [[[0, 1], [1, 0]], [[2, 3]], slice(None)],
+      [[[0, 1], [2, 3]], [[0]], slice(None)],
+      [[[2, 1]], [[0, 3], [4, 4]], slice(None)],
+      [[[2]], [[0, 3], [4, 1]], slice(None)],
+      # non-contiguous indexing subspace
+      [[0, 2, 3], slice(None), [1, 3, 4]],
+
+      # less dim, ellipsis
+      [[0, 2], ],
+      [[0, 2], slice(None)],
+      [[0, 2], Ellipsis],
+      [[0, 2], slice(None), Ellipsis],
+      [[0, 2], Ellipsis, slice(None)],
+      [[0, 2], [1, 3]],
+      [[0, 2], [1, 3], Ellipsis],
+      [Ellipsis, [1, 3], [2, 3]],
+      [Ellipsis, [2, 3, 4]],
+      [Ellipsis, slice(None), [2, 3, 4]],
+      [slice(None), Ellipsis, [2, 3, 4]],
+
+      # ellipsis counts for nothing
+      [Ellipsis, slice(None), slice(None), [0, 3, 4]],
+      [slice(None), Ellipsis, slice(None), [0, 3, 4]],
+      [slice(None), slice(None), Ellipsis, [0, 3, 4]],
+      [slice(None), slice(None), [0, 3, 4], Ellipsis],
+      [Ellipsis, [[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None)],
+      [[[0, 1], [1, 0]], [[2, 1], [3, 5]], Ellipsis, slice(None)],
+      [[[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None), Ellipsis],
+    ]
+
+    for indexer in indices_to_test:
+      assert_get_eq(reference, indexer)
+
+      assert_set_eq(reference, indexer, 212)
+      assert_set_eq(reference, indexer, get_set_tensor(reference, indexer))
+      assert_backward_eq(reference, indexer)
+
+  def test_numpy_parity_and_backward_4d(self):
+    reference = Tensor.arange(0., 1296).reshape(3, 9, 8, 6)
+
+    indices_to_test = [
+      [slice(None), slice(None), slice(None), [0, 3, 4]],
+      [slice(None), slice(None), [2, 4, 5, 7], slice(None)],
+      [slice(None), [2, 3], slice(None), slice(None)],
+      [[1, 2], slice(None), slice(None), slice(None)],
+      [slice(None), slice(None), [0, 2, 3], [1, 3, 4]],
+      [slice(None), slice(None), [0], [1, 2, 4]],
+      [slice(None), slice(None), [0, 1, 3], [4]],
+      [slice(None), slice(None), [[0, 1], [1, 0]], [[2, 3]]],
+      [slice(None), slice(None), [[0, 1], [2, 3]], [[0]]],
+      [slice(None), slice(None), [[5, 6]], [[0, 3], [4, 4]]],
+      [slice(None), [0, 2, 3], [1, 3, 4], slice(None)],
+      [slice(None), [0], [1, 2, 4], slice(None)],
+      [slice(None), [0, 1, 3], [4], slice(None)],
+      [slice(None), [[0, 1], [3, 4]], [[2, 3], [0, 1]], slice(None)],
+      [slice(None), [[0, 1], [3, 4]], [[2, 3]], slice(None)],
+      [slice(None), [[0, 1], [3, 2]], [[0]], slice(None)],
+      [slice(None), [[2, 1]], [[0, 3], [6, 4]], slice(None)],
+      [slice(None), [[2]], [[0, 3], [4, 2]], slice(None)],
+      [[0, 1, 2], [1, 3, 4], slice(None), slice(None)],
+      [[0], [1, 2, 4], slice(None), slice(None)],
+      [[0, 1, 2], [4], slice(None), slice(None)],
+      [[[0, 1], [0, 2]], [[2, 4], [1, 5]], slice(None), slice(None)],
+      [[[0, 1], [1, 2]], [[2, 0]], slice(None), slice(None)],
+      [[[2, 2]], [[0, 3], [4, 5]], slice(None), slice(None)],
+      [[[2]], [[0, 3], [4, 5]], slice(None), slice(None)],
+      [slice(None), [3, 4, 6], [0, 2, 3], [1, 3, 4]],
+      [slice(None), [2, 3, 4], [1, 3, 4], [4]],
+      [slice(None), [0, 1, 3], [4], [1, 3, 4]],
+      [slice(None), [6], [0, 2, 3], [1, 3, 4]],
+      [slice(None), [2, 3, 5], [3], [4]],
+      [slice(None), [0], [4], [1, 3, 4]],
+      [slice(None), [6], [0, 2, 3], [1]],
+      [slice(None), [[0, 3], [3, 6]], [[0, 1], [1, 3]], [[5, 3], [1, 2]]],
+      [[2, 2, 1], [0, 2, 3], [1, 3, 4], slice(None)],
+      [[2, 0, 1], [1, 2, 3], [4], slice(None)],
+      [[0, 1, 2], [4], [1, 3, 4], slice(None)],
+      [[0], [0, 2, 3], [1, 3, 4], slice(None)],
+      [[0, 2, 1], [3], [4], slice(None)],
+      [[0], [4], [1, 3, 4], slice(None)],
+      [[1], [0, 2, 3], [1], slice(None)],
+      [[[1, 2], [1, 2]], [[0, 1], [2, 3]], [[2, 3], [3, 5]], slice(None)],
+
+      # less dim, ellipsis
+      [Ellipsis, [0, 3, 4]],
+      [Ellipsis, slice(None), [0, 3, 4]],
+      [Ellipsis, slice(None), slice(None), [0, 3, 4]],
+      [slice(None), Ellipsis, [0, 3, 4]],
+      [slice(None), slice(None), Ellipsis, [0, 3, 4]],
+      [slice(None), [0, 2, 3], [1, 3, 4]],
+      [slice(None), [0, 2, 3], [1, 3, 4], Ellipsis],
+      [Ellipsis, [0, 2, 3], [1, 3, 4], slice(None)],
+      [[0], [1, 2, 4]],
+      [[0], [1, 2, 4], slice(None)],
+      [[0], [1, 2, 4], Ellipsis],
+      [[0], [1, 2, 4], Ellipsis, slice(None)],
+      [[1], ],
+      [[0, 2, 1], [3], [4]],
+      [[0, 2, 1], [3], [4], slice(None)],
+      [[0, 2, 1], [3], [4], Ellipsis],
+      [Ellipsis, [0, 2, 1], [3], [4]],
+    ]
+
+    for indexer in indices_to_test:
+      assert_get_eq(reference, indexer)
+      assert_set_eq(reference, indexer, 1333)
+      assert_set_eq(reference, indexer, get_set_tensor(reference, indexer))
+
+    indices_to_test += [
+      [slice(None), slice(None), [[0, 1], [1, 0]], [[2, 3], [3, 0]]],
+      [slice(None), slice(None), [[2]], [[0, 3], [4, 4]]],
+    ]
+    for indexer in indices_to_test:
+      assert_get_eq(reference, indexer)
+      assert_set_eq(reference, indexer, 1333)
+      assert_backward_eq(reference, indexer)
+
 if __name__ == '__main__':
  unittest.main()
@@ -26,18 +26,22 @@ class TestLinAlg(unittest.TestCase):
      orthogonality_helper(V)
      reconstruction_helper([U,s_diag,V],a)

-  def test_svd_nonfull(self):
-    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)
-      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)))
-      #reduced U,V is only orthogonal along smaller dim
-      if (m < n): orthogonality_helper(U),orthogonality_helper(V)
-      else: orthogonality_helper(U.transpose(-2,-1)),orthogonality_helper(V.transpose(-2,-1))
-      reconstruction_helper([U,s_diag,V],a)
+  def _test_svd_nonfull(self, size):
+    a = Tensor.randn(size).realize()
+    U,S,V = a.svd(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)))
+    #reduced U,V is only orthogonal along smaller dim
+    if (m < n): orthogonality_helper(U),orthogonality_helper(V)
+    else: orthogonality_helper(U.transpose(-2,-1)),orthogonality_helper(V.transpose(-2,-1))
+    reconstruction_helper([U,s_diag,V],a)
+
+  # faster for parallel pytest
+  def test_svd_nonfull_2_2(self): self._test_svd_nonfull((2,2))
+  def test_svd_nonfull_5_3(self): self._test_svd_nonfull((5,3))
+  def test_svd_nonfull_3_5(self): self._test_svd_nonfull((3,5))
+  def test_svd_nonfull_2_2_2_2_3(self): self._test_svd_nonfull((2,2,2,2,3))

  @unittest.skip("very big. recommend wrapping with TinyJit around inner function")
  def test_svd_large(self):
@@ -1,774 +0,0 @@
-#!/usr/bin/env python
-import unittest
-import numpy as np
-from tinygrad.dtype import dtypes, Invalid
-from tinygrad.helpers import prod
-from tinygrad.shape.shapetracker import ShapeTracker, View, views_to_valid_uop
-from tinygrad import Variable
-from tinygrad.uop.ops import UOp, Ops, graph_rewrite
-from tinygrad.codegen.late.devectorizer import sym
-from itertools import product
-
-def shapetracker_getitem(st:ShapeTracker, val:int):
-  valid_idx = views_to_valid_uop(st.reshape((st.size,)).views, (UOp.const(dtypes.int, val),))
-  idx, valid = valid_idx.get_idx(), valid_idx.get_valid()
-  idx, valid = graph_rewrite(idx, sym), graph_rewrite(valid, sym)
-  assert idx.op is Ops.CONST and valid.op is Ops.CONST
-  return idx.arg, valid.arg
-
-class CheckingShapeTracker:
-  def __init__(self, shape):
-    self.st = ShapeTracker.from_shape(shape)
-    self.t = np.arange(prod(shape), dtype=np.int32).reshape(shape)
-
-  @property
-  def shape(self):
-    return self.t.shape
-
-  def simplify(self):
-    self.st = self.st.simplify()
-    return self
-
-  def reshape(self, new_shape):
-    self.st = self.st.reshape(new_shape)
-    self.t = self.t.reshape(new_shape)
-    return self
-
-  def permute(self, axis):
-    self.st = self.st.permute(axis)
-    self.t = np.transpose(self.t, axis)
-    return self
-
-  def expand(self, new_shape):
-    self.st = self.st.expand(new_shape)
-    self.t = np.broadcast_to(self.t, new_shape)
-    return self
-
-  def flip(self, arg):
-    self.st = self.st.flip(arg)
-    self.t = np.flip(self.t, tuple(i for i in range(len(arg)) if arg[i]))
-    return self
-
-  def shrink(self, arg):
-    self.st = self.st.shrink(arg)
-    self.t = self.t[tuple([slice(x[0], x[1]) for x in arg])]
-    return self
-
-  def pad(self, arg):
-    self.st = self.st.pad(arg)
-    self.t = np.pad(self.t, arg, constant_values=-1)
-    return self
-
-  def __getitem__(self, val):
-    return self.t.flatten()[val]
-
-  @property
-  def views(self): return self.st.views
-
-  @property
-  def contiguous(self): return self.st.contiguous
-
-  def assert_same(self):
-    x = [(v[0] if (v:=shapetracker_getitem(self.st, i))[1] and v[0] is not Invalid else -1) for i in range(prod(self.st.shape))]
-    y = [self[i] for i in range(prod(self.shape))]
-    assert self.st.shape == self.shape
-    assert x == y, f"mismatch shapetracker:{x} real:{y}"
-
-@unittest.skip("don't create shapetrackers with views")
-class TestRealIssues(unittest.TestCase):
-  def test_reshape_doesnt_multiview(self):
-    self.st = ShapeTracker((View.create((256, 256, 2, 2, 2, 2, 2, 256, 8, 2), (0, 8, 0, 4, 0, 0, 2, 16384, 2048, 1), 0, None),))
-    self.st.reshape((128, 2, 256, 2, 2, 2, 2, 2, 256, 8, 2))
-    assert len(self.st.views) == 1
-
-  def test_reshape_stable_diffusion(self):
-    # regression test for https://github.com/tinygrad/tinygrad/pull/2616
-    st = ShapeTracker((View((2, 1920, 32, 32), (1310720, 1024, 32, 1), 0, ((0, 2), (0, 1280), (0, 32), (0, 32)), False),))
-    st = st.reshape((2, 32, 240, 256))
-    assert len(st.views) == 2
-
-  def test_reshape_trailing_invalid_ones(self):
-    st = ShapeTracker((View(shape=(1, 1, 5), strides=(0, 0, 1), offset=-5, mask=((1, 1), (0, 1), (0, 5)), contiguous=False),))
-    st = st.reshape((5,))
-    assert len(st.views) == 1
-    assert st.views[0].mask == ((0,0),)
-
-class TestRealDoesntSimplify(unittest.TestCase):
-  def tearDown(self):
-    self.st = self.st.simplify()
-    assert len(self.st.views) != 1
-
-  def test_1(self):
-    self.st = ShapeTracker((
-      View.create((8, 3, 1, 2, 11, 1), (33, 11, 0, 0, 1, 0), 0, None),
-      View.create((8, 6, 11), (66, 11, 1), 0, None)))
-    self.assertEqual(self.st.is_expanded(), (False, False, False))
-
-  def test_2(self):
-    self.st = ShapeTracker((
-      View.create((2, 2, 4, 3, 3), (72, 9, 18, -3, -1), 8, None),
-      View.create((4, 4, 3, 3), (36, 9, 3, 1), 0, None)))
-    self.assertEqual(self.st.is_expanded(), (False, False, False, False))
-
-class TestRealStrides(unittest.TestCase):
-  def test_1(self):
-    st = ShapeTracker((
-      View.create((2048,), (1,), 0, ((0, 512),)),
-      View.create((16, 32, 4), (128, 4, 1), 0, None),
-    ))
-    self.assertEqual(st.is_expanded(), (False, False, False))
-
-  def test_2(self):
-    # test/test_ops.py::TestOps::test_simple_padding_conv1d
-    st = ShapeTracker((
-      View.create((6, 2, 5, 14), (90, 45, 1, 5), 0, ((0, 6), (0, 2), (0, 5), (0, 9))),
-      View.create((6, 2, 78), (140, 70, 1), 0, ((0, 6), (0, 2), (0, 70))),
-      View.create((6, 2, 13, 6), (156, 78, 1, 13), 0, None),
-    ))
-    self.assertEqual(st.is_expanded(), (False, False, False, False))
-
-  def test_3(self):
-    # test/test_ops.py::TestOps::test_simple_cumsum
-    st = ShapeTracker((
-      View.create((4, 256, 512), (256, 0, 1), 0, ((0, 4), (0, 256), (0, 256))),
-      View.create((4, 131327), (131072, 1), 0, ((0, 4), (0, 131072))),
-      View.create((4, 511, 257), (131327, 1, 511), 0, None),
-    ))
-    self.assertEqual(st.is_expanded(), (False, False, False))
-
-  def test_4(self):
-    # test/test_nn.py::TestNN::test_conv_transpose1d
-    st = ShapeTracker((
-      View.create((4, 16, 56, 2), (896, 56, 1, 0), 0, ((0, 4), (0, 16), (0, 56), (0, 1))),
-      View.create((1, 4, 1, 16, 8, 121), (0, 1792, 0, 112, 0, 1), -5, ((0, 1), (0, 4), (0, 1), (0, 16), (0, 8), (5, 116))),
-      View.create((4, 64, 115, 16, 7), (15488, 0, 1, 968, 122), 0, None),
-    ))
-    self.assertEqual(st.is_expanded(), (False, True, False, False, False))
-
-  def test_5(self):
-    # test/test_ops.py::TestOps::test_conv2d
-    st = ShapeTracker((
-      View.create((1, 3, 1, 12, 2, 8), (0, 132, 0, 12, 1, 2), 0, ((0, 1), (0, 3), (0, 1), (0, 11), (0, 2), (0, 6))),
-      View.create((1, 3, 22, 21), (0, 192, 16, 1), 0, ((0, 1), (0, 3), (0, 12), (0, 16))),
-      View.create((3, 11, 7, 2, 3), (462, 21, 1, 231, 7), 0, None),
-    ))
-    self.assertEqual(st.is_expanded(), (False, False, False, True, False))
-
-class TestIndexExpressions2d(unittest.TestCase):
-  def setUp(self):
-    shapes = [(30, 5), (15, 10), (15, 1), (5, 10), (5, 1)] # Make sure dim0 is a multiple of 5, one of the tests divides this dimension by 5
-    offsets = [0, 1, 15, 28, 10000]
-    self.sts = [ShapeTracker.from_shape((prod(base_shape)+offset,)).shrink(((offset, offset+prod(base_shape)),)).\
-                reshape(base_shape) for base_shape in shapes for offset in offsets]
-    self.offset = [offset for base_shape in shapes for offset in offsets]
-    self.shapes = [shape for shape in shapes for offset in offsets]
-    self.idxs_exprs = []
-
-  def tearDown(self):
-    for st, offset, shape, idxs_expr in zip(self.sts, self.offset, self.shapes, self.idxs_exprs):
-      numel = prod(shape)
-      self.check_bounds(idxs_expr(self.default_idxs(st.shape)), offset, numel)
-      idx0s = [(0,0), (0, min(1, st.shape[0]-1)), (0, st.shape[0]-1), (min(3, st.shape[0]-1), min(6, st.shape[0]-1)), (st.shape[0]-1, st.shape[0]-1)]
-      idx1s = [(0,0), (0, min(1, st.shape[1]-1)), (0, st.shape[1]-1), (min(3, st.shape[1]-1), min(6, st.shape[1]-1)), (st.shape[1]-1, st.shape[1]-1)]
-      idx2s = [(0,0), (0, min(1, st.shape[2]-1)), (0, st.shape[2]-1), (min(3, st.shape[2]-1), min(6, st.shape[2]-1)),
-               (st.shape[2]-1, st.shape[2]-1)] if len(st.shape) == 3 else [None for _ in idx0s]
-      for idx0, idx1, idx2 in product(idx0s, idx1s, idx2s):
-        idxs = [Variable(f"idx{i}", idx[0], idx[1]) for i, idx in enumerate((idx0, idx1, idx2)) if idx is not None]
-        self.check_bounds(idxs_expr(idxs), offset, numel)
-
-  def default_idx(self, shape):
-    return Variable("idx", 0, prod(shape)-1)
-
-  def default_idxs(self, shape):
-    return [Variable(f"idx{i}", 0, d-1) for i,d in enumerate(shape)]
-
-  def check_bounds(self, expr, offset, numel):
-    assert expr.vmin >= offset
-    assert expr.vmax <= offset + numel - 1
-
-  def test_noop(self):
-    for st, base_shape, offset in zip(self.sts, self.shapes, self.offset):
-      self.idxs_exprs.append(lambda idxs, base_shape=base_shape, offset=offset: idxs[0]*base_shape[1] + idxs[1] + offset)
-
-  def test_permute(self):
-    new_st = []
-    for st, base_shape, offset in zip(self.sts, self.shapes, self.offset):
-      st = st.permute((1, 0))
-      self.idxs_exprs.append(lambda idxs, base_shape=base_shape, offset=offset: idxs[0] + idxs[1]*base_shape[1] + offset)
-      new_st.append(st)
-    self.sts = new_st
-
-  def test_reshape(self):
-    new_st = []
-    for st, base_shape, offset in zip(self.sts, self.shapes, self.offset):
-      st = st.reshape((base_shape[0], 1, base_shape[1]))
-      self.idxs_exprs.append(lambda idxs, base_shape=base_shape, offset=offset: idxs[0]*base_shape[1] + idxs[2] + offset)
-      new_st.append(st)
-    self.sts = new_st
-
-  def test_reshape_expand(self):
-    new_st = []
-    for st, base_shape, offset in zip(self.sts, self.shapes, self.offset):
-      st = st.reshape((base_shape[0], 1, base_shape[1]))
-      st = st.expand((base_shape[0], base_shape[1], base_shape[1]))
-      self.idxs_exprs.append(lambda idxs, base_shape=base_shape, offset=offset: idxs[0]*base_shape[1] + idxs[2] + offset)
-      new_st.append(st)
-    self.sts = new_st
-
-  def test_permute_reshape_1(self): # This tests multiple views
-    new_st = []
-    for st, base_shape, offset in zip(self.sts, self.shapes, self.offset):
-      st = st.permute((1, 0))
-      st = st.reshape((base_shape[0]//5, 1, base_shape[1]*5))
-      self.idxs_exprs.append(lambda idxs, base_shape=base_shape, offset=offset: (idxs[0]*(base_shape[1]*5)+idxs[2])%base_shape[0]*base_shape[1] + \
-                             (idxs[0]*(base_shape[1]*5)+idxs[2])//base_shape[0] + offset)
-      new_st.append(st)
-    self.sts = new_st
-
-  def test_permute_reshape_2(self):
-    new_st = []
-    for st, base_shape, offset in zip(self.sts, self.shapes, self.offset):
-      st = st.permute((1, 0))
-      st = st.reshape((1, base_shape[0]//5, base_shape[1]*5))
-      self.idxs_exprs.append(lambda idxs, base_shape=base_shape, offset=offset: (idxs[1]*(base_shape[1]*5)+idxs[2])%base_shape[0]*base_shape[1] + \
-                             (idxs[1]*(base_shape[1]*5)+idxs[2])//base_shape[0] + offset)
-      new_st.append(st)
-    self.sts = new_st
-
-  def test_reshaping_splitting(self):
-    self.st = CheckingShapeTracker((5,10,5,10))
-    self.st.permute((1, 0, 3, 2))
-    self.st.pad(((0,0), (0,5), (0,0), (0,5)))
-    self.st.reshape((10,2,5,10,2,5))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_reshape_splitting_1(self):
-    self.st = CheckingShapeTracker((1,10,1))
-    self.st.pad(((0,4),(0,0),(1,0)))
-    self.st.reshape((5,5,2,2))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_reshape_combining_1(self):
-    self.st = CheckingShapeTracker((2,1,10))
-    self.st.pad(((2,6), (0,0), (0,0)))
-    self.st.reshape((100,))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_reshape_combining_2(self):
-    self.st = CheckingShapeTracker((1,1,5))
-    self.st.pad(((3,6), (0,0), (0,5)))
-    self.st.reshape((100,))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_reshape_combining_3(self):
-    self.st = CheckingShapeTracker((1,1,4))
-    self.st.pad(((3,6), (0,0), (1,5)))
-    self.st.reshape((100,))
-    assert len(self.st.views) == 1
-    assert self.st.views[0].mask[0] == (31, 35)
-    self.st.assert_same()
-
-  def test_reshape_combining_4(self):
-    # interestingly this one is quite slow
-    self.st = CheckingShapeTracker((1,1,5,5,1,1,5))
-    self.st.pad(((2,1), (0,0), (0,2), (0,0), (2,1), (0,0), (0,2)))
-    self.st.reshape((28,5,28))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_reshape_splitting_combining(self):
-    self.st = CheckingShapeTracker((1,5,5))
-    self.st.pad(((0,4), (0,5), (0,0)))
-    self.st.reshape((10,25))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_reshape_only_1s(self):
-    self.st = CheckingShapeTracker((1, 1, 1, 4, 1, 3, 5, 1))
-    self.st.pad(((0,4), (0,0), (0,0), (1,1), (0,0), (0,0), (0,0), (0,0)))
-    self.st.reshape((5, 6, 3, 5))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-    self.st.reshape((1, 1, 5, 6, 3, 5, 1, 1))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-    self.st.reshape((1, 5, 6, 1, 3, 1, 5, 1))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_zero_mask_1(self):
-    self.st = CheckingShapeTracker((1, 3, 2))
-    self.st.pad(((0,0), (0,3), (0,0)))
-    self.st.shrink(((0,1), (3,6), (0,2)))
-    self.st.reshape((3,2))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-    self.st.reshape((1, 3, 1, 2, 1))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_zero_mask_2(self):
-    self.st = CheckingShapeTracker((1, 3, 2))
-    self.st.pad(((0,2), (0,3), (0,0)))
-    self.st.shrink(((2,3), (3,6), (0,2)))
-    self.st.reshape((3,2))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-    self.st.reshape((1, 3, 1, 2, 1))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_expanded_reshaped(self):
-    self.st = CheckingShapeTracker((1, 3, 2, 1))
-    self.st.expand((5, 3, 2, 2))
-    self.st.pad(((0,0), (0,3), (0,0), (0, 0)))
-    self.st.reshape((5, 2, 3, 2, 2))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-  def test_splitting_big(self):
-    self.st = CheckingShapeTracker((1, 5, 1, 15, 1))
-    self.st.pad(((0,0), (0,5), (0,0), (0,15), (0,0)))
-    self.st.reshape((10, 1, 30))
-    self.st.permute((2,1,0))
-    self.st.reshape((2,3,5,2,5))
-    assert len(self.st.views) == 1
-    v = self.st.views[-1]
-    assert v.strides == (0, 5, 1, 0, 15) and v.mask == ((0, 1), (0, 3), (0, 5), (0, 1), (0, 5))
-    self.st.assert_same()
-
-  def test_combining_big(self):
-    self.st = CheckingShapeTracker((1,3,1,5,3,1))
-    self.st.pad(((0,0),(2,2),(0,0),(0,0),(0,0),(0,0)))
-    self.st.reshape((1,1,1,105,1,1))
-    assert len(self.st.views) == 1
-    v = self.st.views[-1]
-    assert v.strides == (0, 0, 0, 1, 0, 0) and v.mask == ((0, 1), (0, 1), (0, 1), (30, 75), (0, 1), (0, 1)) and v.offset == -30
-    self.st.assert_same()
-
-  def test_pad_reshape(self):
-    self.st = CheckingShapeTracker((4,))
-    self.st.pad(((2,2),))
-    self.st.reshape((4,2))
-    assert len(self.st.views) == 1
-    self.st.assert_same()
-
-class TestSimplifyingShapeTracker(unittest.TestCase):
-  def setUp(self):
-    self.st = CheckingShapeTracker((1, 10))
-
-  def tearDown(self):
-    self.st.assert_same()
-
-  # multiview simplify
-  def test_expand_contract_simple(self):
-    self.st = self.st.expand((10, 10))
-    self.st = self.st.reshape((100,))
-    print(self.st.views)
-    assert (len(self.st.views) == 2)
-    self.st = self.st.reshape((10, 10))
-    print(self.st.views)
-
-    self.st = self.st.simplify()
-    print(self.st.views)
-    assert (len(self.st.views) == 1)
-
-  # multiview simplify
-  def test_expand_contract_different_shape(self):
-    self.st.expand((10, 10))
-    self.st.reshape((100,))
-    print(self.st.views)
-    assert (len(self.st.views) == 2)
-    self.st.reshape((2, 5, 2, 5))
-    print(self.st.views)
-
-    self.st = self.st.simplify()
-    print(self.st.views)
-    assert (len(self.st.views) == 1)
-
-  # multiview simplify
-  def test_expand_contract_still_complex(self):
-    self.st.expand((10, 10))
-    self.st.reshape((100,))
-    print(self.st.views)
-    assert (len(self.st.views) == 2)
-    self.st.reshape((5, 20))
-
-    self.st = self.st.simplify()
-    print(self.st.views)
-    assert (len(self.st.views) == 2)
-
-# Tensor.zeros(2, 4).permute(1,0).reshape(2, 4)
-# (d1*4 + d0%4), d1=x//4, d0=x%4 = ((x//4)*4) + (x%4)%4
-
-class TestComplexShapeTracker(unittest.TestCase):
-  def test_add_1s(self):
-    self.st = CheckingShapeTracker((4, 4))
-    self.st.permute((1,0))
-    self.st.reshape((1,4,1,4,1))
-    assert not self.st.contiguous
-    self.st.permute((0,3,2,1,4))
-    assert self.st.contiguous
-
-  def test_permute_1s_simple(self):
-    self.st = CheckingShapeTracker((1, 16, 9,9))
-    self.st.permute((1,0,2,3))
-    assert self.st.contiguous
-    self.st = CheckingShapeTracker((2, 16, 9,9))
-    self.st.permute((1,0,2,3))
-    assert not self.st.contiguous
-
-  def test_remove_1s_simple(self):
-    self.st = CheckingShapeTracker((1, 16, 1, 1))
-    self.st.reshape((16,))
-    assert self.st.contiguous
-
-  def test_remove_1s(self):
-    self.st = CheckingShapeTracker((1, 4, 1, 4, 1))
-    self.st.permute((0,3,2,1,4))
-    self.st.reshape((4,4))
-    assert not self.st.contiguous
-    self.st.permute((1,0))
-    assert self.st.contiguous
-
-  def test_permute_reshape(self):
-    self.st = CheckingShapeTracker((4, 4))
-    self.st.permute((1,0))
-    self.st.reshape((2, 2, 2, 2))
-    # TODO: should also be tested by test_super_complex
-    assert len(self.st.views) == 1
-
-  def test_factorize_split(self):
-    self.st = CheckingShapeTracker((4, 4))
-    self.st.permute((1,0))
-    self.st.reshape((2, 2, 2, 2))
-    self.st.permute((2,3,0,1))
-    assert self.st.contiguous
-
-  def test_factorize_combine(self):
-    self.st = CheckingShapeTracker((4, 4, 4))
-    self.st.permute((2, 0, 1))
-    self.st.reshape((4, 16))
-    self.st.permute((1, 0))
-    assert self.st.contiguous
-
-  def test_factorize_combine_add_ones(self):
-    self.st = CheckingShapeTracker((4, 4, 4))
-    self.st.permute((2, 0, 1))
-    self.st.reshape((4, 16, 1, 1))
-    self.st.permute((1, 0, 2, 3))
-    assert self.st.contiguous
-
-  def test_fancy_factorize(self):
-    self.st = CheckingShapeTracker((32, 3, 3, 1))
-    self.st.reshape((8, 4, 3, 3))
-    assert len(self.st.views) == 1
-
-  def test_super_complex_2_fail(self):
-    self.st = CheckingShapeTracker((4, 4, 4))
-    self.st.permute((2, 0, 1))
-    self.st.reshape((16, 4))
-    assert len(self.st.views) != 1
-
-  def test_work(self):
-    self.st = CheckingShapeTracker((64, 1024, 4))
-    self.st.reshape((1, 64, 128, 32))
-    self.st.permute((0, 3, 1, 2))
-    self.st.reshape((1, 32, 1, 64, 128))
-    self.st.permute((0, 3, 4, 1, 2))
-    assert self.st.contiguous
-
-  def test_work2(self):
-    self.st = CheckingShapeTracker((64, 1024, 4))
-    self.st.reshape((1, 64, 128, 32))
-    self.st.permute((0, 3, 1, 2))
-    self.st.reshape((1, 1, 32, 64, 128))
-    self.st.permute((0, 3, 4, 1, 2))
-    self.st.reshape((64, 1024, 4))
-    print(self.st.views)
-    assert self.st.contiguous
-
-class TestShapeTrackerEquality(unittest.TestCase):
-  def test_simple_equals(self):
-    self.assertEqual(ShapeTracker.from_shape((10,10)), ShapeTracker.from_shape((10,10)))
-  def test_other_equals(self):
-    st1 = ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True)))
-    st2 = ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True)))
-    self.assertEqual(st1, st2)
-
-class TestSingleShapeTracker(unittest.TestCase):
-  def setUp(self):
-    self.st = CheckingShapeTracker((7,4))
-
-  def tearDown(self):
-    self.st.assert_same()
-
-  def test_reshape(self):
-    self.st.reshape((7,1,4))
-    assert self.st.contiguous
-
-  def test_permute(self):
-    self.st.permute((1,0))
-    assert not self.st.contiguous
-
-  def test_shrink(self):
-    self.st.shrink(((1,2), (0,4)))
-    assert not self.st.contiguous
-
-  def test_double_permute(self):
-    self.st.permute((1,0))
-    self.st.permute((1,0))
-    assert self.st.contiguous
-
-  def test_reshape_permute(self):
-    self.st.reshape((7,1,4))
-    self.st.permute((0,1,2))
-    assert self.st.contiguous
-
-  def test_reshape_permute_yes(self):
-    self.st.reshape((7,1,4))
-    self.st.permute((0,2,1))
-    assert self.st.contiguous
-
-  def test_reshape_permute_no(self):
-    self.st.reshape((4,7))
-    self.st.permute((1,0))
-    assert not self.st.contiguous
-
-class TestShapeTrackerFuzzFailures(unittest.TestCase):
-  def setUp(self):
-    self.st = CheckingShapeTracker((3,3,3))
-  def tearDown(self):
-    self.st.assert_same()
-  def test_case_1(self):
-    self.st.shrink(((1, 2), (1, 3), (1, 3)))
-    self.st.reshape((1, 4))
-    self.st.shrink(((0, 1), (1, 3)))
-    self.st = self.st.simplify()
-  def test_case_2(self):
-    self.st.flip( (True, False, True) )
-    self.st.reshape( (3, 9) )
-    self.st.shrink( ((1, 2), (1, 5)) )
-    self.st.flip( (True, True) )
-  def test_case_3(self):
-    self.st.shrink( ((0, 2), (0, 2), (0, 1)) )
-    self.st.permute( (1, 0, 2) )
-    self.st.reshape( (4,) )
-    self.st.shrink( ((0, 3),) )
-    self.st.flip( (True, False) )
-  def test_case_4(self):
-    self.st.reshape( (3, 3, 3, 1) )
-    self.st.pad( ((0, 0), (0, 0), (0, 0), (1, 1)) )
-    self.st.shrink( ((0, 2), (1, 2), (0, 2), (0, 1)) )
-    self.st.expand( (2, 1, 2, 3) )
-
-class TestMaskedShapeTracker(unittest.TestCase):
-  def test_pad_1x1(self):
-    self.st = CheckingShapeTracker((1,1))
-    self.st.pad(((1,1), (1,1)))
-    self.st.assert_same()
-
-  def test_pad_2x2(self):
-    self.st = CheckingShapeTracker((2,2))
-    self.st.pad(((1,1), (1,1)))
-    self.st.assert_same()
-
-  def test_pad_reshape(self):
-    st1 = CheckingShapeTracker((1, 2))
-    st1.pad(((1, 0), (0, 1)))
-    st1.reshape((3, 2))
-    st1.assert_same()
-
-    st2 = CheckingShapeTracker((1, 2))
-    st2.pad(((1, 1), (0, 2)))
-    st2.reshape((4, 3))
-    st2.assert_same()
-
-    st3 = CheckingShapeTracker((1, 1, 1, 2))
-    st3.pad(((0, 2), (1, 2), (2, 2), (0, 4)))
-    st3.reshape((4, 3, 6, 5))
-    st3.assert_same()
-
-class TestShapeTracker(unittest.TestCase):
-  def setUp(self):
-    self.st = CheckingShapeTracker((7,4))
-    self.apply = lambda fxn: [fxn(x) for x in [self.st]]
-
-  def tearDown(self):
-    self.st.assert_same()
-
-  def test_noop(self):
-    pass
-
-  def test_simple_split(self):
-    self.test_permute()
-    self.apply(lambda x: x.reshape((prod(self.st.shape), )))
-
-  def test_simple_pad(self):
-    self.st.pad(((1,1), (1,1)))
-
-  def test_pad_shrink(self):
-    self.st.pad(((1,1), (1,1)))
-    self.st.shrink(((0,4), (0,4)))
-
-  def test_pad_one_sided(self):
-    self.st.pad(((0,1), (0,0)))
-
-  def test_pad_reshape(self):
-    self.st.pad(((0,1), (0,0)))
-    self.st.reshape((8*4,))
-
-  def test_pad_pad(self):
-    self.st.pad(((1,1), (1,1)))
-    self.st.pad(((1,1), (1,1)))
-
-  def test_pad_permute(self):
-    self.st.pad(((1,1), (2,2)))
-    self.st.permute((1,0))
-
-  def test_pad_expand(self):
-    self.st.reshape((7,4,1))
-    self.st.pad(((1,1), (1,1), (0,0)))
-    self.st.expand((9,6,4))
-
-  def test_pad_expand_alt(self):
-    self.st.pad(((1,1), (1,1)))
-    self.st.reshape((9,6,1))
-    self.st.expand((9,6,4))
-
-  def test_pad_flip(self):
-    self.st.pad(((1,4), (1,3)))
-    self.st.flip((True, False))
-
-  def test_pad_flip_int(self):
-    self.st.pad(((1,4), (1,3)))
-    self.st.flip((0, 1))
-
-  def test_reshape(self):
-    new_shape = self.st.shape[::-1]
-    self.apply(lambda x: x.reshape(new_shape))
-
-  def test_permute(self):
-    if len(self.st.shape) == 2: self.apply(lambda x: x.permute((1,0)))
-    elif len(self.st.shape) == 3: self.apply(lambda x: x.permute((2,0,1)))
-
-  def test_reshape_with_1(self):
-    new_shape = (self.st.shape[0], 1, self.st.shape[1])
-    self.apply(lambda x: x.reshape(new_shape))
-
-  def test_expand(self):
-    self.test_reshape_with_1()
-    new_shape = list(self.st.shape)
-    new_shape[1] = 2
-    self.apply(lambda x: x.expand(tuple(new_shape)))
-
-  def test_flip_0(self):
-    self.apply(lambda x: x.flip((True, False)))
-
-  def test_flip_1(self):
-    self.apply(lambda x: x.flip((False, True)))
-
-  def test_flip_01(self):
-    self.apply(lambda x: x.flip((True, True)))
-
-  def test_slice_0(self):
-    self.apply(lambda x: x.shrink(((1, x.shape[0]), (0, x.shape[1]))))
-
-  def test_slice_1(self):
-    self.apply(lambda x: x.shrink(((0, x.shape[0]), (1, x.shape[1]))))
-
-  def test_slice_1c1(self):
-    self.apply(lambda x: x.shrink(((0, 1), (0, 1))))
-
-  def test_slice_1c2(self):
-    self.apply(lambda x: x.shrink(((1, 2), (1, 2))))
-
-  def test_double_permute(self):
-    self.apply(lambda x: x.permute((1, 0)))
-    self.apply(lambda x: x.permute((1, 0)))
-
-  def test_slice_permute(self):
-    self.apply(lambda x: x.shrink(((0, 2), (2, 4))))
-    self.apply(lambda x: x.permute((1, 0)))
-
-  def test_slice_expand(self):
-    self.apply(lambda x: x.shrink(((0, 2), (3, 4))))
-    self.apply(lambda x: x.expand((2, 10)))
-
-  def test_double_flip(self):
-    self.apply(lambda x: x.flip((True, False)))
-    self.apply(lambda x: x.flip((True, False)))
-
-  def test_flip(self): self.apply(lambda x: x.flip((True, False)))
-  def test_flip2(self): self.apply(lambda x: x.flip((False, True)))
-  def test_flip3(self): self.apply(lambda x: x.flip((True, True)))
-
-  def test_reshape_then_permute(self):
-    self.test_reshape()
-    self.test_permute()
-
-  def test_reshape_then_expand(self):
-    self.test_reshape()
-    self.test_expand()
-
-  def test_permute_then_reshape(self):
-    self.test_permute()
-    self.test_reshape()
-
-  def test_expand_then_reshape(self):
-    self.test_expand()
-    self.test_reshape()
-
-  def test_combo(self):
-    self.test_permute()
-    self.test_reshape()
-    self.test_slice_1()
-    self.test_expand()
-    self.test_permute()
-
-class TestVariableShrink(unittest.TestCase):
-  def test_shrink(self):
-    st = ShapeTracker.from_shape((10,))
-    st = st.shrink(((0, Variable("i", 1, 10)),))
-    assert len(st.views) == 1
-
-  def test_shrink_bound(self):
-    st = ShapeTracker.from_shape((10,))
-    st = st.shrink(((0, Variable("i", 1, 10).bind(3)),))
-    assert len(st.views) == 1
-
-class TestVariableMerge(unittest.TestCase):
-  def test_add_reshape(self):
-    vi = Variable("i", 1, 10)
-    st1 = ShapeTracker.from_shape((vi,))
-    st2 = ShapeTracker.from_shape((1, vi,))
-    st = st1+st2
-    assert len(st.views) == 1
-
-  def test_add_stride_0(self):
-    st1 = ShapeTracker.from_shape((3,), (0,))
-    st2 = ShapeTracker.from_shape((Variable("i", 1, 10).bind(3),), (0,))
-    st = st1+st2
-    assert len(st.views) == 1, f"multiview {st}"
-
-  def test_add_reshape_bound(self):
-    vi = Variable("i", 1, 10).bind(3)
-    st1 = ShapeTracker.from_shape((vi,))
-    st2 = ShapeTracker.from_shape((1, vi,))
-    st = st1+st2
-    assert len(st.views) == 1
-
-  def test_simplify(self):
-    vi = Variable("i", 1, 10).bind(3)
-    st1 = ShapeTracker.from_shape((vi,))
-    st2 = ShapeTracker.from_shape((1, vi,))
-    st = ShapeTracker((st1.views[0], st2.views[0]))
-    st = st.simplify()
-    assert len(st.views) == 1
-
-if __name__ == '__main__':
-  unittest.main()
@@ -1,108 +0,0 @@
-import unittest
-from tinygrad.helpers import prod
-from tinygrad.shape.view import View
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad import Variable
-from test.unit.test_shapetracker import shapetracker_getitem
-
-class MultiShapeTracker:
-  def __init__(self, sts:list[ShapeTracker]): self.sts = sts
-  @property
-  def shape(self): return self.sts[0].shape
-  def reshape(self, arg): self.sts = [x.reshape(arg) for x in self.sts]
-  def permute(self, arg): self.sts = [x.permute(arg) for x in self.sts]
-  def expand(self, arg): self.sts = [x.expand(arg) for x in self.sts]
-  def shrink(self, arg): self.sts = [x.shrink(arg) for x in self.sts]
-  def flip(self, arg): self.sts = [x.flip(arg) for x in self.sts]
-  def pad(self, arg): self.sts = [x.pad(arg) for x in self.sts]
-
-def st_equal(st1:ShapeTracker, st2:ShapeTracker) -> bool:
-  if st1.shape != st2.shape: return False
-  if st1 == st2: return True
-  for i in range(0, prod(st1.shape)):
-    st1_off, st1_v = shapetracker_getitem(st1, i)
-    st2_off, st2_v = shapetracker_getitem(st2, i)
-    if st1_v != st2_v or (st1_off != st2_off and st1_v):
-      print(f"ST MISMATCH @ {i}, {st1_v=} != {st2_v=}, {st1_off=} != {st2_off=}")
-      print(st1)
-      print(st2)
-      return False
-  return True
-
-class TestShapeTrackerBasics(unittest.TestCase):
-  def test_pad_shrink_removes_mask(self):
-    a = ShapeTracker.from_shape((10, 10))
-    a = a.pad(((0,2), (0,2)))
-    a = a.shrink(((0,10), (0,10)))
-    assert len(a.views) == 1 and a.views[-1].mask is None
-
-  def test_pad_shrink_leaves_mask(self):
-    a = ShapeTracker.from_shape((10, 10))
-    a = a.pad(((0,2), (0,2)))
-    a = a.shrink(((0,10), (0,11)))
-    assert len(a.views) == 1 and a.views[-1].mask is not None
-
-  def test_reshape_makes_same(self):
-    a = ShapeTracker.from_shape((2, 5))
-    x = a.pad( ((2, 0), (0, 0)) )
-    x = x.reshape( (2, 2, 5) )
-    x1 = x.reshape( (4, 5) )
-    x1 = x1.reshape( (2, 2, 5) )
-    assert x == x1.simplify()
-
-  def test_simplify_is_correct(self):
-    multiv = ShapeTracker(views=(View(shape=(15, 3), strides=(9, 1), offset=6, mask=None, contiguous=False),
-                                 View(shape=(4, 3), strides=(12, 4), offset=0, mask=None, contiguous=False)))
-    assert st_equal(multiv, multiv.simplify())
-
-class TestShapeTrackerAdd(unittest.TestCase):
-  def test_simple_add_reshape(self):
-    a = ShapeTracker.from_shape((10, 10))
-    a = a.reshape((100,))
-    b = ShapeTracker.from_shape((100,))
-    assert a+b == b
-
-  @unittest.skip("no longer simplifies")
-  def test_simple_add_permute(self):
-    a = ShapeTracker.from_shape((10, 10))
-    a = a.permute((1,0))
-    b = ShapeTracker.from_shape((10, 10))
-    b = b.permute((1,0))
-    assert a+b == ShapeTracker.from_shape((10, 10))
-
-  def test_plus_real1(self):
-    st = MultiShapeTracker([ShapeTracker.from_shape((15, 9))])
-    st.shrink( ((0, 15), (6, 9)) )
-    backup = st.sts[0]
-    st.sts.append(ShapeTracker.from_shape(backup.shape))
-    st.reshape( (45,) )
-    st.flip( (True,) )
-    st.reshape( (15, 3) )
-    assert st_equal(backup + st.sts[1], st.sts[0])
-
-  def test_off_by_one(self):
-    st1 = ShapeTracker(views=(View(shape=(5,), strides=(1,), offset=0, mask=None, contiguous=True),
-                              View(shape=(5,), strides=(1,), offset=0, mask=None, contiguous=True)))
-    st2 = ShapeTracker(views=(View(shape=(4,), strides=(1,), offset=0, mask=None, contiguous=True),
-                              View(shape=(5,), strides=(1,), offset=0, mask=None, contiguous=True)))
-    assert not (st_equal(st1, st2))
-
-class TestShapeTrackerAddVariable(unittest.TestCase):
-  def test_merge_symbolic_views(self):
-    var_i = Variable('i', 1, 10)
-    var_j = Variable('i', 1, 10)
-    vm1 = View(shape=(var_i, var_j, 3), strides=(3, 0, 1), offset=0, mask=None, contiguous=False)
-    vm2 = View(shape=(var_i, var_j, 3), strides=(var_j*3, 3, 1), offset=0, mask=None, contiguous=True)
-    ShapeTracker((vm1,)) + ShapeTracker((vm2,))
-
-  def test_merge_symbolic_views_2(self):
-    var_i = Variable('i', 1, 10)
-    var_j = Variable('j', 1, 10)
-    vm1 = View(shape=(var_i, var_j), strides=(0, 0), offset=0, mask=None, contiguous=False)
-    vm2 = View(shape=(var_i, var_j), strides=(var_j, 1), offset=0, mask=None, contiguous=True)
-    ret = (ShapeTracker((vm1,)) + ShapeTracker((vm2,))).reshape((var_i, var_j, 1))
-    ret_2 = ShapeTracker((vm1,)) + ShapeTracker((vm2,)).reshape((var_i, var_j, 1))
-    assert ret == ret_2
-
-if __name__ == '__main__':
-  unittest.main()
@@ -1,5 +1,4 @@
 import unittest
-from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad import Variable
 from tinygrad.tensor import Tensor

@@ -7,115 +6,33 @@ class TestSymbolic(unittest.TestCase):
  def assert_tuple_equal(self, x, y):
    for a,b in zip(x,y): self.assertFalse(a != b)

-  def test_symbolic_st(self):
-    x = Variable("x", 1, 100)
-    st = ShapeTracker.from_shape((x, 3))
-    self.assert_tuple_equal(st.shape, (x, 3))
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
-
-  def test_is_expanded_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.assert_tuple_equal(st.is_expanded(), (False, False))
-
-  def test_is_expanded_1(self):
-    st = ShapeTracker(views=(View(shape=(3, (Variable('i', 1, 10)+2)), strides=(Variable('i', 1, 10), 1), offset=0, mask=((0, 3), (0, Variable('i', 1, 10))), contiguous=False),))   # noqa: E501
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
-
-  def test_is_expanded_2(self):
-    st = ShapeTracker(views=(View(shape=(3, (Variable('i', 1, 10)+Variable('j', 1, 10))), strides=(Variable('i', 1, 10), 1), offset=0, mask=((0, 3), (0, Variable('i', 1, 10))), contiguous=False),))   # noqa: E501
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
-
-  def test_merge_view_recursion_err(self):
-    vm2 = View(shape=(Variable('j', 1, 10),), strides=(0,), offset=0, mask=None, contiguous=False)
-    vm1 = View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True)
-    self.assertEqual(vm2+vm1, None)
-
-  def test_merge_view_recursion_err2(self):
-    vm2 = View(shape=(Variable('a', 1, 10).bind(4),), strides=(0,), offset=0, mask=None, contiguous=False)
-    # NOTE: vm1 is different from what create function would give, and this test vm2+vm1 halts
-    vm1 = View(shape=(Variable('a', 1, 10).bind(4),), strides=(1,), offset=0, mask=((0, Variable('a', 1, 10).bind(4)),), contiguous=False)
-    self.assertEqual(vm2+vm1, None)
-
-    vm3 = View.create(shape=(Variable('a', 1, 10).bind(4),))
-    self.assertEqual(vm3.shape, vm1.shape)
-    self.assertEqual(vm3.strides, vm1.strides)
-    self.assertEqual(vm2+vm3, vm2)
-
  def test_cat_dim0_is_expanded(self):
    i = Variable("i", 1, 5).bind(3)
    j = Variable("j", 1, 5).bind(3)
    k = Variable("k", 1, 5).bind(3)
    t = Tensor.rand(5, 4)[:i].cat(Tensor.rand(5, 4)[:j], dim=0).cat(Tensor.rand(5, 4)[:k], dim=0)
-    st = t.uop.st
-    self.assert_tuple_equal(st.shape, (i+j+k, 4))
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
+    self.assert_tuple_equal(t.shape, (i+j+k, 4))
    t = Tensor.rand(5, 3)[:i].cat(Tensor.rand(5, 3)[:i], dim=0).cat(Tensor.rand(3, 3), dim=0)
-    st = t.uop.st
-    self.assert_tuple_equal(st.shape, (2*i+3, 3))
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
+    self.assert_tuple_equal(t.shape, (2*i+3, 3))

  def test_cat_dim1_strides(self):
    i = Variable("i", 1, 5).bind(4)
    j = Variable("j", 1, 5).bind(4)
    k = Variable("k", 1, 5).bind(4)
    t = Tensor.rand(3, 5)[:, :i].cat(Tensor.rand(3, 5)[:, :j], dim=1).cat(Tensor.rand(3, 5)[:, :k], dim=1)
-    st = t.uop.st
-    self.assert_tuple_equal(st.shape, (3, i+j+k))
-    self.assert_tuple_equal(st.is_expanded(), (False, False))
+    self.assert_tuple_equal(t.shape, (3, i+j+k))

 class TestSymbolicVarVals(unittest.TestCase):
  def assert_equal(self, x, y): self.assertFalse(x != y)
-  def test_var_vals_empty(self):
-    assert ShapeTracker.from_shape((3, 4, 5)).var_vals == {}
-
-  def test_var_vals_shape(self):
-    x = Variable("x", 1, 100).bind(3)
-    assert ShapeTracker.from_shape((x, 3)).var_vals == {"x": 3}
-
-  def test_var_vals_offset(self):
-    x = Variable("x", 1, 100).bind(3)
-    st = ShapeTracker.from_shape((4, 3)).shrink(((x, x+1), (0, 3)))
-    self.assert_equal(st.views[-1].offset, x * 3)
-    assert st.var_vals == {"x": 3}
-
-  def test_var_vals_mask(self):
-    x = Variable("x", 1, 100).bind(3)
-    view = View.create(shape=(3,4), strides=(4,1), offset=0, mask=((0, x), (0, 4)))
-    st = ShapeTracker(views=(view,))
-    assert st.var_vals == {"x": 3}
-
-  def test_var_vals_complex(self):
-    x = Variable("x", 1, 100).bind(3)
-    y = Variable("y", 1, 100).bind(4)
-    z = Variable("z", 1, 100).bind(5)
-    st = ShapeTracker.from_shape((x, 5, y)).shrink(((0, x), (z, z+1), (0, 3)))
-    self.assert_equal(st.views[-1].offset, y * z)
-    assert st.var_vals == {"x": 3, "y": 4, "z": 5}
-
-  def test_shrink_reshape(self):
-    x = Variable("x", 1, 100).bind(3)
-    st = ShapeTracker.from_shape((10, 10, 10)).shrink(((x, x+3), (3, 7), (2, 5)))
-    st = st.reshape((3*4*3,))
-    assert st.var_vals == {"x": 3}
-
-class TestShapeTrackerUnbind(unittest.TestCase):
-  def test_view_unbind(self):
-    v = Variable("v", 1, 100)
-    bv = Variable("v", 1, 100).bind(3)
-    unbound_view, var_val = View.create(shape=(bv, 4)).unbind()
-    assert unbound_view == View.create(shape=(v, 4))
-    assert var_val == {v: 3}

  def test_shrink_unbind(self):
    v = Variable("v", 1, 100)
    bv = Variable("v", 1, 100).bind(2)
    t = Tensor.rand(3, 4).shrink(((0,bv),(0,4)))
-    unbound_st, var_val = t.uop.st.unbind()
-    assert unbound_st == ShapeTracker((View.create(shape=(v, 4)),))
+    unbound_st, var_val = t.uop.unbind_all()
    assert var_val == {v: 2}
    t = Tensor.rand(3, 4).shrink(((bv, bv+1), (0, 4)))
-    unbound_st, var_val = t.uop.st.unbind()
-    assert unbound_st == ShapeTracker((View.create(shape=(1, 4), offset=4*v),))
+    unbound_st, var_val = t.uop.unbind_all()
    assert var_val == {v: 2}

 class TestSymbolicReshape(unittest.TestCase):
@@ -145,17 +62,6 @@ class TestSymbolicReshape(unittest.TestCase):
        ret = ret.reshape(1, vi*vj)
        assert ret.shape == (1, vi*vj)

-  def test_symbolic_mask(self):
-    # taken from gpt2 single kvcache
-    # these two caused problems in gpt2 if reshape merged views
-    view = View(shape=(1, (Variable('start_pos', 1, 128).bind(2)+1), 16, 64), strides=(0, 0, 64, 1), offset=1024, mask=((0, 1), (Variable('start_pos', 1, 128).bind(2), (Variable('start_pos', 1, 128).bind(2)+1)), (0, 16), (0, 64)), contiguous=False)   # noqa: E501
-    new_shape = (1, 1, (Variable('start_pos', 1, 128).bind(2)+1), 16, 64)
-    assert view.reshape(new_shape) is None
-
-    view = View(shape=(2, 1, (Variable('start_pos', 1, 128)+1), 16, 64), strides=(0, 0, 1024, 64, 1), offset=131072, mask=((1, 2), (0, 1), (0, (Variable('start_pos', 1, 128)+1)), (0, 16), (0, 64)), contiguous=False)   # noqa: E501
-    new_shape = (2, (Variable('start_pos', 1, 128)+1), 16, 64)
-    assert view.reshape(new_shape) is None
-
 class TestSymbolicExpand(unittest.TestCase):
  def test_expand_into_symbols(self):
    vi = Variable("i", 1, 5).bind(3)
@@ -198,6 +104,5 @@ class TestSymbolicPad(unittest.TestCase):
    t = t[:9]
    assert t.tolist() == [0,0,0,0,1,1,1,1,1]

-
 if __name__ == '__main__':
  unittest.main()
@@ -4,8 +4,6 @@ from tinygrad.uop.ops import UPat, Ops, UOp

 # NOTE: unlike before base for a realized tensor is always a BUFFER
 realized_pattern = UPat(Ops.BUFFER)
-# after realization, base tensor uops become RESHAPE(BUFFER)
-buffer_view_pattern = UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER),))
 def is_pattern_uop(u:UOp, pat:UPat): assert pat.match(u, {}), f"{u}\nis not\n{pat}"
 def is_pattern(ten:Tensor, pat:UPat): is_pattern_uop(ten.uop, pat)

@@ -57,13 +55,5 @@ class TestTensorUopRepresentation(unittest.TestCase):
    is_pattern(c, UPat(Ops.ADD))
    for s in c.uop.src: is_pattern_uop(s.base, realized_pattern)

-  def test_empty_buf(self):
-    a = Tensor.empty(3, 3)
-    is_pattern(a, UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER),)))
-    vi = UOp.variable("i", 1, 3).bind(1)
-    a = Tensor.empty(3, vi)
-    is_pattern(a, UPat(Ops.RESHAPE, src=(UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),))),))
-    self.assertEqual(a.uop.base.buffer.size, 9)
-
 if __name__ == '__main__':
  unittest.main()
@@ -1,73 +0,0 @@
-#!/usr/bin/env python
-import unittest
-from tinygrad.shape.view import View, merge_dims
-# from tinygrad.shape.shapetracker import ShapeTracker
-
-class TestView(unittest.TestCase):
-  def test_canonicalize_empty_mask(self):
-    v = View.create(shape=(2,2,2), strides=(4,2,1), mask=((0,2),(0,2),(0,2)))
-    self.assertIsNone(v.mask)
-    v = View.create(shape=(4,3,2), strides=(1,4,10), mask=((0,4),(0,3),(0,2)))
-    self.assertIsNone(v.mask)
-
-  def test_empty_mask_contiguous(self):
-    v1 = View.create(shape=(2,2,2), strides=(4,2,1), mask=None)
-    v2 = View.create(shape=(2,2,2), strides=(4,2,1), mask=((0,2),(0,2),(0,2)))
-    self.assertEqual(v1.contiguous, v2.contiguous)
-    v1 = View.create(shape=(1,1,1,4), strides=(0,0,0,1), offset=0, mask=None)
-    v2 = View.create(shape=(1,1,1,4), strides=(0,0,0,1), offset=0, mask=((0,1),(0,1),(0,1),(0,4)))
-    self.assertEqual(v1.contiguous, v2.contiguous)
-    v = View.create(shape=(2,3,4), mask=((0,2),(0,3),(0,4)))
-    self.assertTrue(v.contiguous)
-
-  def test_reshape_all_invalid(self):
-    v = View.create((4,5), mask=((0,0), (0,0))).reshape((20,))
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View.create((20,), mask=((0,0),)))
-
-  def test_add_0(self):
-    v1 = View.create((2,3,4))
-    v2 = View.create((2,0,4))
-    self.assertEqual(v2, v1+v2)
-
-  def test_add_0_masked(self):
-    v1 = View.create((2,3,4), mask=((0, 0), (0, 0), (0, 0)))
-    v2 = View.create((2,0,4))
-    self.assertEqual(v2, v1+v2)
-
-class TestMergeDims(unittest.TestCase):
-  def test_contiguous(self):
-    shape = (2, 3, 4)
-    strides = (12, 4, 1) #=strides_for_shape(shape)
-    m = merge_dims(shape, strides)
-    self.assertEqual(m, ((24, 1, 24),))
-
-  def test_0_in_strides(self):
-    shape = (2, 3, 4)
-    self.assertEqual(merge_dims(shape, (0, 4, 1)), ((2, 0, 0), (12, 1, 12)))
-    self.assertEqual(merge_dims(shape, (0, 0, 1)), ((6, 0, 0), (4, 1, 4)))
-    self.assertEqual(merge_dims(shape, (3, 1, 0)), ((6, 1, 6), (4, 0, 4)))
-    self.assertEqual(merge_dims(shape, (0, 0, 0)), ((24, 0, 0),))
-
-  def test_pad(self):
-    # print(ShapeTracker.from_shape((1, 2)).pad(((1, 0), (0, 1))).views[-1])
-    self.assertEqual(merge_dims((2, 3), (0, 1), ((1, 2), (0, 2))), ((6, 1, 3),))
-
-    # print(f"{ShapeTracker.from_shape((1, 1, 2)).pad(((1, 0), (1, 0), (0, 1))).views[-1]}")
-    self.assertEqual(merge_dims((2, 2, 3), (0, 0, 1), ((1, 2), (1, 2), (0, 2))), ((12, 1, 3),))
-
-    # print(f"{ShapeTracker.from_shape((1, 1, 2, 2)).pad(((1, 0), (1, 0), (0, 1), (0, 1))).views[-1]}")
-    self.assertEqual(merge_dims((2, 2, 3, 3), (0, 0, 2, 1), ((1, 2), (1, 2), (0, 2), (0, 2))), ((12, 2, 3), (3, 1, 3)))
-
-    # print(f"{ShapeTracker.from_shape((2, 1, 2)).pad(((0, 0), (1, 0), (0, 1))).views[-1]}")
-    self.assertEqual(merge_dims((2, 2, 3), (2, 0, 1), ((0, 2), (1, 2), (0, 2))), ((2, 2, 2), (6, 1, 3)))
-
-  def test_different_1_pad(self):
-    # print(f"{ShapeTracker.from_shape((2, 2, 1)).pad(((0, 0), (0, 0), (0, 1))).views[-1]}")
-    self.assertEqual(merge_dims((2, 2, 2), (2, 1, 0), ((0, 2), (0, 2), (0, 1))), ((4, 1, 4), (2, 0, 2)))
-
-    # print(f"{ShapeTracker.from_shape((2, 1, 1)).pad(((0, 0), (0, 1), (0, 1))).views[-1]}")
-    self.assertEqual(merge_dims((2, 2, 2), (1, 0, 0), ((0, 2), (0, 2), (0, 1))), ((2, 1, 2), (4, 0, 4)))
-
-if __name__ == '__main__':
-  unittest.main()
@@ -2,8 +2,7 @@ import unittest, decimal, json, struct
 from dataclasses import dataclass
 from typing import Generator

-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher
-from tinygrad.uop.ops import graph_rewrite, track_rewrites, TRACK_MATCH_STATS
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, TRACK_MATCH_STATS
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
@@ -15,15 +14,16 @@ def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=Non
    sink = graph_rewrite(sink, TrackedPatternMatcher(pm.patterns), name=names[i] if names else None)
  return sink

-# real VIZ=1 pickles these tracked values
-from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, _name_cnt
-traces = [(tracked_keys, tracked_ctxs, uop_fields)]
-from tinygrad.viz.serve import get_metadata, uop_to_json, get_details
-def get_viz_list(): return get_metadata(traces)
+# real VIZ=1 loads the trace from a file, we just keep it in memory for tests
+from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, _name_cnt, RewriteTrace
+from tinygrad.viz import serve
+serve.trace = RewriteTrace(tracked_keys, tracked_ctxs, uop_fields)
+from tinygrad.viz.serve import get_rewrites, get_full_rewrite, uop_to_json
+def get_viz_list(): return get_rewrites(serve.trace)
 def get_viz_details(rewrite_idx:int, step:int) -> Generator[dict, None, None]:
  lst = get_viz_list()
  assert len(lst) > rewrite_idx, "only loaded {len(lst)} traces, expecting at least {idx}"
-  return get_details(tracked_ctxs[rewrite_idx][step])
+  return get_full_rewrite(tracked_ctxs[rewrite_idx][step])

 class BaseTestViz(unittest.TestCase):
  def setUp(self):
@@ -326,14 +326,14 @@ def load_profile(lst:list[ProfileEvent]) -> dict:
    event_type, event_count = u("<BI")
    if event_type == 0:
      for _ in range(event_count):
-        name, ref, st, dur, _ = u("<IIIfI")
-        v["events"].append({"name":strings[name], "ref":option(ref), "st":st, "dur":dur})
+        name, ref, key, st, dur, _ = u("<IIIIfI")
+        v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur})
    else:
      v["peak"] = u("<Q")[0]
      for _ in range(event_count):
        alloc, ts, key = u("<BII")
        if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-        else: v["events"].append({"event":"free", "ts":ts, "key":key})
+        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIBB") for _ in range(u("<I")[0])]}})
  return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}

 class TestVizProfiler(unittest.TestCase):
@@ -442,7 +442,7 @@ class TestVizMemoryLayout(BaseTestViz):
    profile_ret = load_profile(Buffer.profile_events)
    ret = profile_ret["layout"][f"{a.device} Memory"]
    self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["events"]), 2)
+    self.assertEqual(len(ret["events"]), 4)

  def test_del_once(self):
    a = _alloc(1)
@@ -451,7 +451,7 @@ class TestVizMemoryLayout(BaseTestViz):
    profile_ret = load_profile(Buffer.profile_events)
    ret = profile_ret["layout"][f"{b.device} Memory"]
    self.assertEqual(ret["peak"], 1)
-    self.assertEqual(len(ret["events"]), 3)
+    self.assertEqual(len(ret["events"]), 4)

  def test_alloc_free(self):
    a = _alloc(1)
@@ -461,7 +461,7 @@ class TestVizMemoryLayout(BaseTestViz):
    profile_ret = load_profile(Buffer.profile_events)
    ret = profile_ret["layout"][f"{c.device} Memory"]
    self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["events"]), 4)
+    self.assertEqual(len(ret["events"]), 6)

  def test_free_last(self):
    bufs = []
@@ -479,5 +479,46 @@ class TestVizMemoryLayout(BaseTestViz):
    self.assertEqual(len(ret["events"]), 6)
    self.assertEqual(len(profile["markers"]), 6)

+  def test_producer_simple(self):
+    a = Tensor.ones(10, device="NULL")
+    Tensor.realize(a.add(1).contiguous())
+    b = Tensor.ones(10, device="NULL")
+    Tensor.realize(b.add(1).contiguous())
+    profile = load_profile(cpu_events+Buffer.profile_events)
+    buffers = profile["layout"]["NULL Memory"]["events"]
+    programs = profile["layout"]["NULL"]["events"]
+    user_cnt = [len(b["arg"]["users"]) for b in buffers if b["arg"].get("users")]
+    self.assertEqual(len(user_cnt), len(programs))
+
+  def test_inflight_buf(self):
+    a = Tensor.empty(1, device="NULL")
+    n = 4
+    for i in range(n): (a+i).realize()
+    profile = load_profile(cpu_events+Buffer.profile_events)
+    buffers = profile["layout"]["NULL Memory"]["events"]
+    user_cnt = [len(b["arg"]["users"]) for b in buffers if b["arg"].get("users")]
+    self.assertEqual(max(user_cnt), n)
+    input_buf = buffers.pop()
+    assert all(u[3] == 0 for u in input_buf["arg"]["users"])
+
+  def test_annotate_read_write(self):
+    a = Tensor.ones(4, device="NULL").contiguous().realize()
+    b = a.assign(a+2)
+    c = a+1
+    Tensor.realize(b, c)
+    buf_events = load_profile(cpu_events+Buffer.profile_events)["layout"]["NULL Memory"]["events"]
+    users = next((b["arg"]["users"] for b in buf_events if len(b["arg"].get("users",[])) == 3))
+    self.assertEqual(users[0][3], 1) # write Tensor.ones
+    self.assertEqual(users[1][3], 2) # read+write Tensor.assign
+    self.assertEqual(users[2][3], 0) # readonly
+
+  def test_dedup_users(self):
+    a = Tensor.empty(1, device="NULL")
+    for _ in range(n:=4): a.add(1).realize()
+    profile = load_profile(cpu_events+Buffer.profile_events)
+    programs = profile["layout"][a.device]["events"]
+    users = profile["layout"][f"{a.device} Memory"]["events"].pop()["arg"]["users"]
+    self.assertEqual(len(programs), len(set(users)), n)
+
 if __name__ == "__main__":
  unittest.main()
@@ -42,7 +42,7 @@ class TestWinograd(unittest.TestCase):
    out = Tensor.conv2d(x,w, padding=1)
    out.mean().backward()
    backward_schedule = Tensor.schedule(x.grad, w.grad)
-    self.assertEqual(len(backward_schedule), 5)
+    self.assertEqual(len(backward_schedule), 4)

  def test_counters(self):
    IC, OC, X, Y = 4,4,9,9
@@ -9,7 +9,7 @@ from tinygrad.renderer import Renderer
 # import all pattern matchers here
 from tinygrad.codegen.quantize import pm_quant
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic
+from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load
 from tinygrad.uop.decompositions import get_late_rewrite_patterns
 from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_expander, pm_group_for_reduce
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
@@ -42,7 +42,7 @@ def get_rewrites_for_renderer(opts:Renderer, optimize:bool=True, linearizer:bool

@functools.cache
 def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
-  # ** lowerer (rewrite_shapetracker_with_index) **
+  # ** lowerer **
  ret: list[RewriteStep] = []

  if optimize:
@@ -62,7 +62,7 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
    ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))

  # ** expander (expand_rewrite) **
-  ret.append(RewriteStep(sym+migrate_indexing, name="postopt symbolic"))
+  ret.append(RewriteStep(sym+migrate_indexing+pm_move_where_on_load, name="postopt symbolic"))

  # expand
  ret.append(RewriteStep(sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander"))
@@ -20,8 +20,8 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
  # can drop valid if idx is out of bound when valid is False
  drop_stmt = []
  for stmt in valid.split_uop(Ops.AND):
-    try: X, is_upper_bound, c = parse_valid(stmt)
-    except ValueError: return None
+    if (res:=parse_valid(stmt)) is None: continue
+    X, is_upper_bound, c = res

    # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i
    if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in X.split_uop(Ops.ADD)):
@@ -5,7 +5,7 @@ from typing import cast, Final
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp
 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, get_single_element
+from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
 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
@@ -88,7 +88,14 @@ class Scheduler:

    self.ast = self.ast.substitute(dict(zip(self.rngs, rng)))

-  def colors(self) -> list[str]: return [axis_colors[x] if not self.dont_use_locals or not x == AxisType.GLOBAL else "BLUE" for x in self.axis_types]
+  def colors(self) -> list[str]:
+    store_rngs = flatten([x.src[2:] for x in self.ast.src])
+    ret = []
+    for x,r in zip(self.axis_types, self.rngs):
+      if self.dont_use_locals and x == AxisType.GLOBAL: ret.append("BLUE")
+      elif r not in store_rngs and x == AxisType.LOOP: ret.append("BLACK")
+      else: ret.append(axis_colors[x])
+    return ret
  def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())])

  def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input_new_rng=None):
@@ -97,7 +104,7 @@ class Scheduler:
    new_rng = UOp.range(amount, self.maxarg+1, new_type) if input_new_rng is None else input_new_rng
    replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),))
    sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
-    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[0]} {amount} {str(new_type).split('.')[1].lower()}")
+    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[:-1]} {amount} {str(new_type).split('.')[1].lower()}")
    return replaced_rng, new_rng

  def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
@@ -200,13 +207,14 @@ class Scheduler:
      self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}")
    elif opt.op is OptOps.SWAP:
      try:
-        altrng = self.rngs[opt.arg]
+        altrng:UOp = self.rngs[opt.arg]
      except IndexError:
        raise KernelOptError
      check(rng.arg[-1] == AxisType.GLOBAL and altrng.arg[-1] == AxisType.GLOBAL, "swap only for globals")
      self.ast = self.ast.substitute({rng:rng.replace(arg=(*altrng.arg[0:-1], rng.arg[-1]), tag=1),
-                                      altrng:altrng.replace(arg=(*rng.arg[0:-1], altrng.arg[-1]), tag=1)})
-      self.ast = graph_rewrite(self.ast, remove_tags)
+                                      altrng:altrng.replace(arg=(*rng.arg[0:-1], altrng.arg[-1]), tag=1)},
+                                      name=f"swap {rng.arg[:-1]} {altrng.arg[:-1]}")
+      self.ast = graph_rewrite(self.ast, remove_tags, name="swap remove tags")
    else:
      raise KernelOptError(f"unsupported opt {opt.op}")

@@ -26,14 +26,6 @@ pm_quant = symbolic+PatternMatcher([
  # x*c1 + y*c2 -> (x+y)*c1 (if c1 and c2 are close floats)
  (UPat.var("x")*UPat.cvar("c1", dtype=dtypes.floats) + UPat.var("y")*UPat.cvar("c2", dtype=dtypes.floats),
   lambda x,y,c1,c2: (x+y)*c1 if abs(c1.arg-c2.arg) < 1e-9 else None),
-  # mul 0 * c1 is 0
-  #(UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) *
-  # UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int).cast(dtypes.float).named("ld"), lambda ld,v,c1: ld*c1),
-  # mul (with plus) 0 * c1 is 0
-  #(UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c1"), UPat(Ops.CONST, arg=0)) *
-  # (UPat(Ops.LOAD, src=(UPat().view(name="v"),)).cast(dtypes.int) + \
-  #  UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar(), UPat(Ops.CONST, arg=0))).cast(dtypes.float).named("ld"),
-  #    lambda ld,v,c1: ld*c1),

  # const push through add
  ((UPat.var("x")*UPat.cvar("c1") + UPat.var("y")*UPat.cvar("c2")) * UPat.cvar("c3"), lambda x,y,c1,c2,c3: (x*c1*c3) + (y*c2*c3)),
@@ -1,5 +1,5 @@
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, ImageDType
-from tinygrad.uop.symbolic import symbolic_flat, sym, invalid_pat
+from tinygrad.uop.symbolic import symbolic_flat, sym
 from tinygrad.helpers import partition
 from tinygrad.dtype import dtypes

@@ -69,64 +69,6 @@ pm_split_ranges = PatternMatcher([

 def no_range(u:UOp) -> bool: return not any(x.op is Ops.RANGE for x in u.backward_slice_with_self)

-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
-  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
-
-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),
-  # lift x*y out of reduce
-  ((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
-   lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and y.vmin > 0 else None),
-  # lift x+y out of reduce on ne
-  ((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),
-  # fold the range
-  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.cvar("val")).reduce(arg=Ops.ADD, allow_any_len=True),
-   lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
-  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(UPat.cvar("val"), 0).reduce(arg=Ops.ADD, allow_any_len=True),
-   lambda r,cut,val: cut.maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
-  # REDUCE on ADD
-  ((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
-   lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)),
-  # MUL casted bool
-  ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast().or_broadcasted(name="b")),
-   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()),
-  (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()),
-  # 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])))),
-  # 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"),
-    lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)),
-  # remove REDUCEs that no longer have a RANGE in the src
-  (UPat(Ops.REDUCE, name="red"), reduce_rangeless),
-])+sym
-
-def reduce_collapse(red:UOp):
-  included, not_included = partition(red.backward_slice, lambda x: any(y in x.backward_slice_with_self for y in red.src[1:]))
-  if any(x.op in {Ops.STORE, Ops.REDUCE} for x in included): return None
-  replaces: dict[UOp, UOp] = {}
-  for u in included:
-    for s in u.src:
-      if s in not_included and s not in replaces and s.op not in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}:
-        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
-
 def reduce_unparented(red:UOp):
  if red.arg not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None
  assert all(x.op is Ops.RANGE for x in red.src[1:]), "some reduce srcs aren't ranges"
@@ -144,6 +86,46 @@ pm_reduce_unparented = PatternMatcher([
  (UPat(Ops.REDUCE, name="red"), reduce_unparented),
 ])

+pm_reduce_collapse = pm_reduce_unparented + 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),
+  # lift x*y out of reduce
+  ((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
+   lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and y.vmin > 0 else None),
+  # lift x+y out of reduce on ne
+  ((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),
+  # fold the range
+  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.cvar("val")).reduce(UPat.var("r"), arg=Ops.ADD),
+   lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
+  (((UPat.var("r")<UPat.var("lower")).logical_not()&(UPat(Ops.RANGE, name="r")<UPat.var("upper"))).where(UPat.cvar("val"), 0).reduce(UPat.var("r"),
+    arg=Ops.ADD), lambda r,lower,upper,val: (upper.minimum(r.src[0])-lower.maximum(0)).maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
+  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(UPat.cvar("val"), 0).reduce(UPat.var("r"), arg=Ops.ADD),
+   lambda r,cut,val: cut.maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
+  # REDUCE on ADD
+  ((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
+   lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)),
+  # MUL casted bool
+  ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)),
+  # reduce on gated load becomes can substitute the range and remove the reduce
+  ((UPat.var("idx")!=(UPat(Ops.RANGE, name="r").or_casted())).where(0, UPat.var("expr")).reduce(UPat.var("r"), arg=Ops.ADD),
+   lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),
+  # AND on WHERE
+  ((UPat(Ops.DEFINE_VAR, name="x") & UPat.var("y")).where(UPat.cvar("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
+    lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)),
+])+sym
+
+def reduce_collapse(red:UOp):
+  included, not_included = partition(red.backward_slice, lambda x: any(y in x.backward_slice_with_self for y in red.src[1:]))
+  if any(x.op in {Ops.STORE, Ops.REDUCE} for x in included): return None
+  replaces: dict[UOp, UOp] = {}
+  for u in included:
+    for s in u.src:
+      if s in not_included and s not in replaces and s.op not in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}:
+        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
+
 pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([
  # remove REDUCE without loads (generic arange opt / indexing). TODO: support multi range
  (UPat(Ops.REDUCE, src=(UPat(), UPat()), name="red"), reduce_collapse),
@@ -1,7 +1,7 @@
 from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
-from typing import Any, Generic, TypeVar, Iterator, Sequence, cast
+from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM
 from tinygrad.helpers import Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
@@ -327,8 +327,8 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
  if device is None: device = Device.DEFAULT
  if dtype == dtypes.bfloat16:
    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"}
+    if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX") and not getenv("NV_NAK")
+    if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not getenv("CPU_LVP")
    return device in {"AMD", "PYTHON", "NULL"}
  if dtype in dtypes.fp8s: return device in {"PYTHON", "NULL"}
  if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,
@@ -357,7 +357,7 @@ if PROFILE:
    from tinygrad.uop.ops import launch_viz
    launch_viz("PROFILE", fn)

-if __name__ == "__main__":
+def enumerate_devices_str() -> Generator[str, None, None]:
  from tinygrad import Tensor, Device

  for device in ALL_DEVICES:
@@ -376,4 +376,7 @@ if __name__ == "__main__":
      result = (colored('PASS', 'green') if any_works else f"{colored('FAIL', 'yellow')}") + ''.join([f'\n{" "*16} {x}' for x in compilers_results])
    except Exception as e:
      result = f"{colored('FAIL', 'red')} {e}"
-    print(f"{'*' if device == Device.DEFAULT else ' '} {device:10s}: {result}")
+    yield f"{'*' if device == Device.DEFAULT else ' '} {device:10s}: {result}"
+
+if __name__ == "__main__":
+  for s in enumerate_devices_str(): print(s)
@@ -5,10 +5,10 @@ from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv,
 from tinygrad.device import Buffer, Compiled, Device, MultiBuffer
 from tinygrad.dtype import DType
 from tinygrad.uop.ops import UOp, Variable, sym_infer, Ops
-from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.engine.realize import ExecItem, capturing, ViewOp, BufferCopy, BufferXfer, CompiledRunner, Runner, Estimates
 from tinygrad.engine.memory import _internal_memory_planner
 from tinygrad.nn.state import get_parameters
+from tinygrad.schedule.rangeify import mop_cleanup
 from dataclasses import dataclass
 from weakref import WeakKeyDictionary

@@ -158,7 +158,7 @@ class CapturedJit(Generic[ReturnType]):
  input_replace: dict[tuple[int, int], int]
  extra_view_inputs: list[tuple[int, int, str, int, DType]]
  expected_names: list[int|str]
-  expected_st_vars_dtype_device: list[tuple[ShapeTracker, tuple[Variable, ...], DType, str]]
+  expected_st_vars_dtype_device: list[tuple[UOp, tuple[Variable, ...], DType, str]]

  def __reduce__(self):
    # TODO: free_intermediates here? replan_buffers_memory_layout here?
@@ -224,7 +224,7 @@ def _prepare_jit_inputs(args, kwargs):
  input_buffers: list[Buffer] = flatten([rb.bufs if isinstance(rb:=lb.base.realized, MultiBuffer) else [rb]
                                         for lb in lbs if lb.base.realized is not None])
  assert len(set(input_buffers)) == len(input_buffers), "duplicate inputs to JIT"
-  st_varval_dtype_device = [(*unwrap(lb.st).unbind(), lb.dtype, lb.device) for lb in lbs]
+  st_varval_dtype_device = [(*(lb.substitute({lb.base:UOp(Ops.NOOP)}, extra_pm=mop_cleanup).unbind_all()), lb.dtype, lb.device) for lb in lbs]
  _var_vals = merge_dicts([x[1] for x in st_varval_dtype_device] + [dict(v.unbind() for v in (args + tuple(kwargs.values())) if isinstance(v, UOp))])
  var_vals = {k.expr:v for k,v in _var_vals.items()}
  st_vars_dtype_device = [(x[0], tuple(sorted(x[1].keys(), key=lambda v: v.expr)), x[2], x[3]) for x in st_varval_dtype_device]
@@ -166,8 +166,9 @@ class ExecItem:
    var_vals = self.fixedvars if _var_vals is None else (_var_vals|self.fixedvars)
    bufs = [cast(Buffer, x) for x in self.bufs] if jit else [cast(Buffer, x).ensure_allocated() for x in self.bufs]
    if PROFILE:
-      payload = {"metadata":self.metadata, "var_vals":var_vals, "bufs":[b.trace_num for b in bufs]}
-      cpu_events.append(ProfilePointEvent(self.prg.device, "exec", self.prg.display_name, payload))
+      payload = {"metadata":self.metadata, "var_vals":var_vals, "bufs":[b.trace_num for b in bufs], "name":self.prg.display_name}
+      payload["outputs"], payload["inputs"] = (self.prg.p.outs, self.prg.p.ins) if isinstance(self.prg, CompiledRunner) else ([0], [1])
+      cpu_events.append(ProfilePointEvent(self.prg.device, "exec", len(cpu_events), payload))
    et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2)
    if do_update_stats:
      GlobalCounters.kernel_count += 1
@@ -180,10 +181,11 @@ class ExecItem:
        header_color = 'magenta' if jit else ('green' if self.prg.first_run else None)
        ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
        flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
-        flops_str = f"{flops*1e-9:9.2f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:9.2f} TFLOPS", 'green')
-        mem_str = f"{membw*1e-9:6.1f}|{ldsbw*1e-9:<7.1f} GB/s" if membw < 1e13 else colored(f"{membw*1e-12:6.1f}|{ldsbw*1e-12:<7.1f} TB/s", 'green')
+        flops_str = f"{flops*1e-9:7.0f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:7.0f} TFLOPS", 'green')
+        mem_str = f"{membw*1e-9:4.0f}|{ldsbw*1e-9:<6.0f} GB/s" if membw < 1e13 and ldsbw < 1e15 else \
+          colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green')
        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
-          f" {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB"+
+          f" {self.prg.display_name+' '*(46-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
          ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})")+
          f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}")
      self.prg.first_run = False
@@ -4,13 +4,13 @@ from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata
 from tinygrad.helpers import argsort

 def reduce_gradient(ctx:UOp, ret:UOp):
-  def to_inp_shape(x): return x.reshape(x.shape+(1,)*(len(ret.src[0].shape)-len(x.shape))).expand(ret.src[0].shape)
-  if ret.arg[0] == Ops.ADD: return (to_inp_shape(ctx),)
+  def broadcast_to_input(x): return x.reshape(x.shape+(1,)*(len(ret.src[0].shape)-len(x.shape))).expand(ret.src[0].shape)
+  if ret.arg[0] == Ops.ADD: return (broadcast_to_input(ctx),)
  if ret.arg[0] == Ops.MAX:
-    max_is_1s = ret.src[0].eq(to_inp_shape(ret)).cast(ctx.dtype)
-    div = to_inp_shape(max_is_1s.r(Ops.ADD, ret.arg[1]))
-    return ((max_is_1s/div) * to_inp_shape(ctx),)
-  if ret.arg[0] == Ops.MUL: return (to_inp_shape(ctx * ret) / ret.src[0],)
+    mask = ret.src[0].eq(broadcast_to_input(ret)).cast(ctx.dtype)
+    count = mask.r(Ops.ADD, ret.arg[1])
+    return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
+  if ret.arg[0] == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)

 # ctx is grad_output
 pm_gradient = PatternMatcher([
@@ -31,12 +31,12 @@ pm_gradient = PatternMatcher([
  (UPat(Ops.REDUCE_AXIS, name="ret"), reduce_gradient),
  (UPat((Ops.CONTIGUOUS, Ops.FUSE)), lambda ctx: (ctx,)),
  (UPat(Ops.CONTIGUOUS_BACKWARD), lambda ctx: (ctx.contiguous(),)),
-  (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape),)),
-  (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.arg)),)),
-  (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.arg)])),)),
-  (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[1]) for s,p in zip(ret.src[0].shape, ret.arg)])),)),
-  (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip(ret.arg),)),
-  (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret: (ctx.r(Ops.ADD, tuple(i for i,(si,so) in enumerate(zip(ret.src[0].shape, ret.arg)) if si!=so)),)),
+  (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape), None)),
+  (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret: (ctx.r(Ops.ADD,tuple(i for i,(s,n) in enumerate(zip(ret.src[0].shape, ret.shape)) if s!=n)), None)),
+  (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
+  (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
+  (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)),
+  (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip(ret.marg),)),
  (UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
  # there's no gradient for bitcast
  (UPat(Ops.BITCAST), lambda: (None,)),
@@ -60,5 +60,9 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
      if v is None: continue
      if k in grads: grads[k] = grads[k] + v
      else: grads[k] = v
-      if len(forward_metadata:=all_metadata.get(t0, ())): all_metadata[v] = tuple(dataclasses.replace(x, backward=True) for x in forward_metadata)
+      if len(forward_metadata:=all_metadata.get(t0, ())):
+        backward_metadata = tuple(dataclasses.replace(x, backward=True) for x in forward_metadata)
+        # we add the backward metadata to everything new in the graph
+        for bw_uop in v.toposort(lambda x: x not in (t0, *t0.src, grads[t0])):
+          all_metadata[bw_uop] = all_metadata.get(bw_uop, ())+backward_metadata
  return grads
@@ -2,7 +2,7 @@ from __future__ import annotations
 import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass
 import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools
 from dataclasses import dataclass, field
-from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator
+from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast

 T = TypeVar("T")
 U = TypeVar("U")
@@ -84,6 +84,17 @@ def word_wrap(x, wrap=80):
  i = 0
  while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1
  return x[:i] + "\n" + word_wrap(x[i:], wrap)
+def pad_bytes(b:bytes, align:int) -> bytes: return b + b'\x00' * ((align - (len(b) % align)) % align)
+
+@functools.cache
+def canonicalize_strides(shape:tuple[T, ...], strides:tuple[T, ...]) -> tuple[T, ...]:
+  return tuple(cast(T, 0) if s == 1 else st for s, st in zip(shape, strides))
+
+@functools.cache
+def strides_for_shape(shape:tuple[T, ...]) -> tuple[T, ...]:
+  if not shape: return ()
+  strides = tuple(itertools.accumulate(reversed(shape[1:]), operator.mul, initial=1))[::-1]
+  return canonicalize_strides(shape, strides)

 # returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
 def get_contraction(old_shape:tuple[T, ...], new_shape:tuple[T, ...]) -> list[list[int]]|None: # T is sint
@@ -147,18 +158,19 @@ SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), Conte
 PICKLE_BUFFERS, LRU = ContextVar("PICKLE_BUFFERS", 1), 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)
 QUANTIZE, VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("QUANTIZE", 0), ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
 FUSE_ATTENTION = ContextVar("FUSE_ATTENTION", 0)
 EMULATE = ContextVar("EMULATE", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
-CPU_LLVM, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("AMD_LLVM", 1)
+CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 1)
 VIZ = PROFILE = ContextVar("VIZ", 0)
 SPEC = ContextVar("SPEC", 0)
 # TODO: disable by default due to speed
 IGNORE_OOB = ContextVar("IGNORE_OOB", 1)
+PCONTIG = ContextVar("PCONTIG", 0)  # partial contiguous in rangeify
+REAL_SUBSTITUTE = ContextVar("REAL_SUBSTITUTE", 0)

@dataclass(frozen=True)
 class Metadata:
@@ -1242,7 +1242,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    G, V, H = G.detach(), V.detach(), H.detach()
    X.grad = norm_coefficient * X.detach() + G
    opt = TinyAdam([X], b1=alpha, b2=beta, eps=epsilon)
-    opt.m, opt.v, opt.lr = [V], [H], R
+    # NOTE: FUSE_OPTIM can change shapes of m and v
+    opt.m, opt.v, opt.lr = [V.reshape(opt.m[0].shape)], [H.reshape(opt.v[0].shape)], R
    # need no-op for m_hat and v_hat if T == 0
    if T == 0: opt.b1_t, opt.b2_t = opt.b1_t.zeros_like(), opt.b2_t.zeros_like()
    else:
@@ -50,7 +50,7 @@ class Optimizer:
    if self.fused:
      # optimizer fusion just concatenates all the buffers, runs the _step, then splits them back up
      out, extra = self._step([Tensor.cat(*[t.flatten() for t in self.params], dim=0)],
-                              [Tensor.cat(*[unwrap(t.grad).flatten() for t in self.params], dim=0)])
+                              [Tensor.cat(*[unwrap(t.grad).contiguous().flatten() for t in self.params], dim=0)])
      updated_params = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
    else:
      updated_params, extra = self._step(self.params, [unwrap(t.grad) for t in self.params])
@@ -1,10 +1,9 @@
 import json, pathlib, zipfile, pickle, tarfile, struct, functools, io
 from collections import OrderedDict
-from typing import Any, Callable, BinaryIO, Iterable
+from typing import Any, Callable, BinaryIO, Iterable, cast
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import prod, argsort, DEBUG, Timing, CI, unwrap, GlobalCounters, tqdm, round_up, T
-from tinygrad.shape.view import strides_for_shape
+from tinygrad.helpers import prod, argsort, DEBUG, Timing, CI, unwrap, GlobalCounters, tqdm, round_up, T, strides_for_shape

 class TensorIO(io.RawIOBase, BinaryIO):
  def __init__(self, t: Tensor):
@@ -210,7 +209,7 @@ def torch_load(t:Tensor) -> dict[str, Tensor]:
      assert tuple([shape_strides[i][1] for i in argsort(permute_indexes)]) == strides_for_shape(intermediate_shape), "nonpermutable strides"
      if DEBUG >= 3: print(f"WARNING: this torch load is slow. to permute {intermediate_shape} with {permute_indexes}")
      assert storage[1] != dtypes.bfloat16, "can't permute BF16"
-      # TODO: find a nice way to support all shapetracker on disktensors
+      # TODO: find a nice way to support all movement ops on disktensors
      ret = ret.to(None).reshape(intermediate_shape).permute(permute_indexes)

    return ret.reshape(size)
@@ -238,11 +237,18 @@ def torch_load(t:Tensor) -> dict[str, Tensor]:

  if passthrough_reset(zipfile.is_zipfile(fobj)): # NOTE: passthrough_reset required to support python < 3.14
    myzip = zipfile.ZipFile(fobj, 'r')
-    base_name = myzip.namelist()[0].split('/', 1)[0]
-    for n in myzip.namelist():
-      if n.startswith(f'{base_name}/data/'):
-        with myzip.open(n) as myfile:
-          offsets[n.split("/")[-1]] = myfile._orig_compress_start # type: ignore
+    base_name = None
+    header_offsets = {}
+    for zi in myzip.filelist:
+      if base_name is None: base_name = zi.filename.split('/', 1)[0]
+      if zi.filename.startswith(f'{base_name}/data/'): header_offsets[zi.filename.split("/")[-1]] = zi.header_offset
+    # sadly there's no way to get the start of the file in the zip without reading the header
+    # at least here we read them in parallel
+    header_contents = [t[v+26:v+30].bitcast(dtypes.uint16).to('CPU') for v in header_offsets.values()]
+    Tensor.realize(*header_contents)
+    for (n,o),c in zip(header_offsets.items(), header_contents):
+      # header_offset + sizeFileHeader + File name length + Extra field length : https://en.wikipedia.org/wiki/ZIP_(file_format)
+      offsets[n] = o+30+sum(cast(list[int], c.tolist()))
    with myzip.open(f'{base_name}/data.pkl') as myfile:
      return TorchPickle(myfile).load()
  elif passthrough_reset(tarfile.is_tarfile(fobj)): # NOTE: passthrough_reset required to support python < 3.11
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	70a1126830	ugh	2025-10-17 23:14:36 +08:00
geohot	b22790eb0f	fix tests	2025-10-17 23:05:01 +08:00
geohot	dad778564c	reset ending ranges	2025-10-17 22:52:43 +08:00
George HotzandGitHub	c5617ed8cf	Merge branch 'master' into test_fa	2025-10-17 22:41:31 +08:00
George HotzandGitHub	33025b99f6	small changes from fa backward (#12769 )	2025-10-17 22:41:18 +08:00
chenyuandGitHub	e0d0d4372d	fix shape of m and v in onnx Adam with FUSE_OPTIM (#12768 ) value is still slightly off but that's not onnx specific	2025-10-17 10:32:41 -04:00
qazalandGitHub	bd662bea67	viz: light up program runs (#12764 ) * basics work * fix the color * light up program events * swap a with p * better	2025-10-17 19:33:18 +08:00
geohot	28efb4395c	multiout at every level	2025-10-17 19:32:38 +08:00
geohot	bc9048ccca	very big	2025-10-17 19:11:19 +08:00
geohot	7c80285fa8	render colors	2025-10-17 18:58:44 +08:00
geohot	05f69b48e9	end ranges	2025-10-17 18:25:31 +08:00
geohot	5fa053a5ee	TODO: fix pcontig	2025-10-17 17:51:36 +08:00
geohot	eb5070786a	test flash attention backward	2025-10-17 17:28:59 +08:00
George HotzandGitHub	c9a3464f76	those decimals never mattered (#12760 ) * those decimals never mattered * this * improve debug * real substitute fixes pcontig * locals are different buffers	2025-10-17 17:16:24 +08:00
qazalandGitHub	0160f034d6	viz: show display name for copy runners (#12761 ) * viz: show display name for copy runners * more u32	2025-10-17 16:59:51 +08:00
qazalandGitHub	253d32b065	viz: add metadata to buffer user list (#12758 ) * simple failing test * encodings * test passing * key is deduped	2025-10-17 16:28:54 +08:00
George HotzandGitHub	935a60db72	bring back partial contig and flash attention (#12756 ) * bring back partial contig and flash attention * why not 2 * work * that * fix pcontig	2025-10-17 16:19:05 +08:00
Sieds LyklesandGitHub	f6bc620169	UOp.prod and UOp.sum methods (#12755 )	2025-10-17 10:02:01 +02:00
Sieds LyklesandGitHub	d1bb5c0426	slightly flatter symbolic (#12757 )	2025-10-17 09:58:45 +02:00
qazalandGitHub	5417e4b099	viz helper cleanups (#12754 )	2025-10-17 15:20:24 +08:00
qazalandGitHub	3196a7aae3	viz: pre reqs for lighting up programs (#12753 )	2025-10-17 15:03:21 +08:00
qazalandGitHub	dfb8f9fc9e	viz: annotate buffer mutability in the memory graph (#12750 )	2025-10-17 11:53:02 +08:00
Sieds LyklesandGitHub	79c2f1ae26	remove reduce_rangless and replace with reduce_unparented (#12749 )	2025-10-17 04:46:05 +02:00
chenyuandGitHub	9561803cb0	fix assert in test_schedule (#12745 ) * fix assert in test_schedule updated kernel counts and some old tests * fix	2025-10-16 15:39:50 -04:00
chenyuandGitHub	285534ce64	delete DONT_REALIZE_EXPAND and DONT_GROUP_REDUCES (#12744 ) does nothing now	2025-10-16 14:11:33 -04:00
chenyuandGitHub	98239f1156	few shapetracker cleanups (#12741 )	2025-10-16 12:43:27 -04:00
chenyuandGitHub	53478c741d	relax ASSERT_MIN_STEP_TIME for space lab policy (#12742 )	2025-10-16 11:40:36 -04:00
geohotstanandGitHub	5d209ee7ec	onnx helper intermediate node output validation (#12740 ) * start * update comments * good * add comments and better printing * done	2025-10-16 11:17:47 -04:00
sirhcmandGitHub	bce2bc0465	Revert "use RTLD_GLOBAL on macos" (#12738 ) This reverts commit `89fe3e574d`.	2025-10-16 10:07:21 -04:00
chenyuandGitHub	f34f26bca0	fix gpt2 with benchmark (#12736 ) `CPU=1 python3 examples/gpt2.py --benchmark 128` works now	2025-10-16 09:55:20 -04:00
Sieds LyklesandGitHub	55db1b0e0e	reduce where that is cut from two sides (#12733 ) * better rule * correct pattern * shorten line	2025-10-16 15:25:15 +02:00
nimlgenandGitHub	cf9baeea61	Revert "nv: check if jitlink is avail (#12731 )" (#12735 ) This reverts commit `a069a45d14`.	2025-10-16 20:41:49 +08:00
George HotzandGitHub	8be7844b2e	use apply uop for assign to fix assign metadata (#12732 ) * use apply uop for assign * fix metadata for assign * fix backward metadata * those aren't real tests	2025-10-16 20:34:12 +08:00
nimlgenandGitHub	3aa2277b8f	nv: usb4 (#12696 ) * hackish * prog * match * l * simpler * refactor * not osx * apple things * tiny changes * fix mask * match fix * nn	2025-10-16 20:11:19 +08:00
nimlgenandGitHub	a069a45d14	nv: check if jitlink is avail (#12731 )	2025-10-16 19:58:50 +08:00
George HotzandGitHub	a498ec9c18	cleanup names of postrange + fast FUSE_OPTIM (#12730 ) * cleanup names of postrange * make FUSE_OPTIM not slow * delete junk in def r	2025-10-16 19:38:31 +08:00
Sieds LyklesandGitHub	8f740e07ff	no broadcasting/vectors in reduce collapse (#12729 )	2025-10-16 13:22:57 +02:00
qazalandGitHub	533f18b22c	viz: add trace data for inflight buffers (#12728 ) * viz: add trace data for inflight buffers * add test_inflight_buf * temp stores the keys * update tests / use Tensor.ones	2025-10-16 19:15:03 +08:00
George HotzandGitHub	af4479c169	faster stable diffusion load (#12725 ) * faster stable diffusion load * failing tests	2025-10-16 18:31:59 +08:00
nimlgenandGitHub	e7c057d5dc	system: alloc_sysmem return view (#12724 ) * system: alloc_sysmem return view * e	2025-10-16 17:55:01 +08:00
nimlgenandGitHub	b86a33a312	ptx: support bw (#12722 )	2025-10-16 15:38:08 +08:00
nimlgenandGitHub	b8cd66c7a2	nv: support all gb20x and small bar (#12721 )	2025-10-16 15:37:54 +08:00
George HotzandGitHub	1d1e1d9d88	delete the ShapeTracker (#12720 ) * delete the ShapeTracker * fix tests * fix more * fix gc test	2025-10-16 15:36:22 +08:00
George HotzandGitHub	592e86f6f5	remove UOp.st (#12716 ) * remove UOp.st * fix tests * torch backend disable	2025-10-16 14:44:09 +08:00
wozeparrotandGitHub	cc2dfe22f5	tinyfs: fetch file utility (#12719 )	2025-10-15 23:38:56 -07:00
nimlgenandGitHub	3ed543f956	system: reorder funcs + barrier on macos (#12714 )	2025-10-16 14:38:01 +08:00
qazalandGitHub	b77bdbbc62	viz: count unpickle in server startup time (#12715 ) * viz: count unpickle in server startup time * type checking	2025-10-16 13:07:46 +08:00
George HotzandGitHub	7c19db00f1	remove st from jit/split_reduceop (#12713 ) * remove st from jit * fix by merging reshapes * no st usage in rangeify * hmm, stop early works * fix speed regressions	2025-10-16 12:50:58 +08:00
qazalandGitHub	069177c1be	trace buffer producer and consumers (#12639 ) * trace buffer producer and consumers * work * generic colored util * fix batched * basic clicking works * generic javascript that works for producer and consumers * keep focused shape * idle time * timings for producer and consumers dedup * from sd test * tiny cleanups * timeline * work * up to here * assert * list it * work	2025-10-16 11:11:31 +08:00
George HotzandGitHub	4a151e7533	make xcode signing happy, waiting for entitlement (#12712 )	2025-10-16 10:20:34 +08:00
chenyuandGitHub	c3278e5622	clean up old tests (#12708 )	2025-10-15 17:53:17 -04:00
chenyuandGitHub	b8cf35fb77	print macOS version in CI (#12705 )	2025-10-15 15:05:33 -04:00
Daniel GitHub chenyu	d65bd669f8	update tiny torch backend hook (#12575 ) * update the backend to fix torch deprecation warning * use param_hook to avoid full backward hook needlessly firing on inputs which do not require gradients * fix indentation --------- Co-authored-by: chenyu <[email protected]>	2025-10-15 14:02:33 -04:00
nimlgenandGitHub	db5ae846aa	nv: do not use va_addr for cpu accesses (#12697 ) * nv: do not use va_addr for cpu accesses * mypy	2025-10-15 22:48:12 +08:00
nimlgenandGitHub	3ab23af829	nv: copy prog with copyin (#12701 ) * nv: copy prog with copyin * to bytes * fix test	2025-10-15 22:48:01 +08:00
nimlgenandGitHub	fafbf3daea	memory: reserve ptable (#12702 )	2025-10-15 22:47:50 +08:00
geohot	85a907605c	hotfix: only 20 steps of beautiful_mnist_torch, some CI machines are slow	2025-10-15 22:29:34 +08:00
sirhcmandGitHub	e1996d358c	use RTLD_GLOBAL on macos (#12699 )	2025-10-15 22:24:50 +08:00
chenyuandGitHub	312c622d35	support None in pad_to and shrink_to (#12700 )	2025-10-15 09:25:31 -04:00
George HotzandGitHub	612e3d6143	replace mop arg with vectorized index (#12695 ) * replace mop arg with vectorized index * tests passing * better viz * no compile4	2025-10-15 20:50:06 +08:00
wozeparrotandGitHub	9ec4c06d7d	feat: one request per device (#12698 )	2025-10-15 05:22:07 -07:00
Sieds LyklesandGitHub	99aa3bd5f9	reduce collapse reduce only the cut range (#12687 )	2025-10-15 13:57:41 +02:00
Sieds LyklesandGitHub	91ac4f1f92	late merging of where and load (#12694 )	2025-10-15 13:33:06 +02:00
qazalandGitHub	768dc952de	viz ui cleanups / renaming (#12691 ) * better viz names * delete unused * don't use opacity, it's multiplicative * keep styles * scrollbar coloring * pyrender doesn't work here beautiful_mnist r_64_16_32_36@lower all index dtypes	2025-10-15 18:40:22 +08:00
chenyuandGitHub	2e50ed0767	increase timeout of resnet cron (#12693 ) does not finish in 6 hours now	2025-10-15 06:08:58 -04:00
sirhcm GitHub George Hotz nimlgen	0aabc1e938	Mesa NIR backend (NAK/LLVMpipe) (#12089 ) * nak works * TestOps::test_add works * testop has no crashes * fix bool casts * fix typo * add disassemble * RANGE and locals/regs * simplify NAKCompiler * disass cleanup * cleanup nir codegen * almost all tests passing * cleanup notes in extra/ * old notes * only import nak if NIR=1 * fix new SPECIAL syntax * fix local/shared memory * more tests passing * add DEFINE_VAR support * llvmpipe kinda works * diskcache * some mypy stuff * lvp passing test_ops.py * fix imports * actually fix imports * remove 'stdout' * fix llvm import * fix mypy issues * nicer errors * simpler test_dtype skips * test lvp in CI * fix github action syntax * fix more actions typos * switch to mesa 25.1.0 * diskcache_put * better generation for lvp nir_options * b64encode shader blobs * Revert diskcache changes This reverts commits `930fa3de8a` and `8428c694b3`. * general cleanup * better error messages * fix llvm import * fix windows tests * link with libm and libgcc_s * fix some errors * dont check for 'float4' * NIR uses pointer arithmetic * use tinymesa * bump tinymesa * bump tinymesa again * update lvp nir_options * print nir shader with DEBUG * simplify LVPCompiler * more tests * "gated" STORE * NAK is cacheable * more tests * all tests pass locally for NAK * test autogen in CI * autogen deps * more deps * fix uop_gc * fix macos * mypy * save 2 lines * save two more lines * save 1 line * save 4 lines * save more lines * Revert "save more lines" This reverts commit `dd3a720c5a`. * save more lines * fix LVP on windows * refactor * reorganize some code * refactor lib_gpu * move LVP check * out of order loads * remove support.mesa * bump tinymesa version * simplify LVP jit * macos * macos ci * shell: bash * testing * more testing * compute brew prefix * stupid typo * actually fix * lib * stdout on macos * inline gallivm_compile_module * Revert "inline gallivm_compile_module" This reverts commit `b65983b151`. * elf macos * semicolon * inherit from CPULLVMCompiler * ruff * disas test * fix libm linking * default is fine actually * arm works * add elf loader link test * fix NAK beam * pylint is too smart by half --------- Co-authored-by: George Hotz <[email protected]> Co-authored-by: nimlgen <[email protected]>	2025-10-15 17:38:33 +08:00
qazalandGitHub	f0268d13f6	cleanup viz server (#12688 )	2025-10-15 15:58:36 +08:00
nimlgenandGitHub	aa81bde150	amd: usb4/thunderbolt on macs (#12641 ) * tbgpu * works * cleaner * this * zero size * h * fix * simpler * prio over usb * c * not needed * linter * this way * mappings * mypy * mypy * mypy 2 * nn	2025-10-15 13:02:01 +08:00
George HotzandGitHub	236c4590c3	use margs as intermediate for new style mops (#12686 ) * use marg to prepare for movement op change * clean up forced reshape * move marg * more marg * more	2025-10-15 12:43:00 +08:00
qazalandGitHub	7597e1dcac	pyrender in viz (#12682 ) * pyrender in viz * keep profile still print_tree * keep special in render	2025-10-15 11:53:30 +08:00
qazalandGitHub	60e03eec37	viz: add View Program option (#12683 )	2025-10-15 11:37:51 +08:00
George HotzandGitHub	a59439d013	use UOp.shape property instead of UOp.st (#12664 ) * work on shape property * reshape causing issues * more mops * all mops * need to cache it * _shape is like _device * mostly works * shape is good * const uses _shape * fix tests * size doesn't use st * close * test is broken * one less st * hack for 3 op assign * oops, i didn't mean to change that * support emulate in the NullDevice * reproed failure in emulation * fix wmma	2025-10-15 10:01:34 +08:00
chenyuandGitHub	89df6f611d	reenable sdxl mac benchmark (#12680 ) also updated faster sd step times	2025-10-14 17:36:17 -04:00
chenyuandGitHub	d25ceffe8d	update padto opts tests (#12679 )	2025-10-14 17:00:42 -04:00
chenyuandGitHub	e8380968f2	add venv_sd_mlperf to gitignore (#12676 ) training stable diffusion stuff	2025-10-14 12:51:36 -04:00
wozeparrotandGitHub	f228c03f9f	fetch raid from cloud (#10799 ) * feat: initial tinyfs device * feat: don't allow compute on tinyfs device * feat: tensor helpers to load and store * feat: bufferview for tinyfs * fix: keep copy sizes correct * fix: recv large * clean: unneeded * feat: comment * clean: unneeded * clean: remove * clean: remove * feat: get request tag * feat: rename to cloud * feat: send request_id * feat: start computing tree * feat: compute store tree on this side * feat: jank chunked load * feat: more debugging * feat: rename to just load and store * feat: correct chunk count * fix: fix load for < 1mb * feat: comments * feat: don't truncate on block devices * feat: better way of testing block device * feat: don't need to pad that much * feat: connect to nodes directly on load * feat: cache connections * feat: don't hard code chunk size * feat: close mmap when closing file handle * feat: don't overwrite stuff on disk if storing from disk * clean: debug print * fix: close mmap * feat: await workers * feat: fast copy from tinyfs to disk * feat: don't copy to device on last * feat: use single socket per device * feat: raid in tinyfs * clean: remove import * clean: type * feat: maintain single event loop * feat: lower worker count * feat: use connection pool * feat: fetch mapping in its own process * fix: release lock * feat: don't fetch if exists * feat: req id only on stores * feat: always fetch * fix: rangeify * feat: allow specifying raid root * fix: dealloc buffer * feat: start support non 0 offset * clean: use cleaner * feat: don't pass to threadpool * clean: typing	2025-10-14 07:53:55 -07:00
chenyuandGitHub	70dd297a05	BS=96 for bert (#12675 ) 96 trains fine now	2025-10-14 09:07:43 -04:00
Sieds LyklesandGitHub	852d80dff9	better where on load folding (#12651 ) * move where clauses to load * shorten line * drop clauses if they are duplicated * add rule for swapped where branch * where on ungated load * dont move clause if load is in the clause * parse_valid returns None * no data dependent branches * fix rule * enable swapped rule * remove those	2025-10-14 13:30:47 +02:00
nimlgenandGitHub	c7e63601fd	gfx1200 tc for AMD_LLVM (#12673 )	2025-10-14 19:17:48 +08:00
George HotzandGitHub	db4a359374	fix up some slow tests that launch python (#12672 ) * fix up some slow tests that launch python * svd nonfull in parallel * split test_advancedindex	2025-10-14 19:13:55 +08:00
nimlgenandGitHub	4918c827c2	amd: lib_gpu does not need cpu_access (#12670 )	2025-10-14 18:34:34 +08:00
nimlgenandGitHub	0c9d47deab	hcq: add alignment to kernargs (#12669 )	2025-10-14 18:33:12 +08:00
qazalandGitHub	d3bfcd3277	minor patches for SQTT over usb on gfx12 (#12627 ) * disable cpu_access in the sqtt buffer allocation not sure if this is required, it results in a very slow call to pcie_mem_write over USB GPU, removing it worked fine. * fix itrace_se_mask on gfx12 on gfx11 it gave 6 se, on gfx11 this value is 2 so no instructions were traced. * Revert "fix itrace_se_mask on gfx12" This reverts commit `0644adbcd1`.	2025-10-14 18:07:46 +08:00
Sieds LyklesandGitHub	1e6e5a0efd	`parse_valid` returns None instead of raising (#12663 ) * parse_valid returns None * change there too	2025-10-14 11:57:38 +02:00
qazalandGitHub	471bd30d16	cleanup viz/serve.py (#12665 ) * use load_pickle * update comment	2025-10-14 17:50:39 +08:00
George HotzandGitHub	fb61f3519f	remove assign contiguous hack (#12659 ) * remove assign contiguous hack * remove bad contiguous usage in torch backend * assign	2025-10-14 16:42:14 +08:00
George HotzandGitHub	30ee7c4c26	cleanup Device usage in Tensor (#12662 )	2025-10-14 16:22:22 +08:00
Sieds LyklesandGitHub	e06cbfcb8a	combine `pm_drop_and_clauses` (#12660 ) * combine those * wino kernels decreased	2025-10-14 10:09:41 +02:00
George HotzandGitHub	84d4589ed4	remove pylint from pre-commit and CI (#12658 ) * remove pylint from pre-commit and CI * multidevice test is fast * faster pre-commit * 8 is faster than 4 * better name * how did that typecheck?	2025-10-14 15:39:59 +08:00
qazalandGitHub	8ecaf839e2	cleanup UOp tracing [pr] (#12657 )	2025-10-14 14:50:59 +08:00