improve DEBUG=2 string with TB/s and TFLOPS [pr]

Add mxfp4 support (#11873 )
* bump ggml url * map mxfp4 to tensor * tests
2025-08-27 11:24:09 -07:00 · 2025-08-27 10:56:56 -07:00 · 2025-08-27 15:50:43 +03:00 · 2025-08-27 03:38:15 +02:00 · 2025-08-26 15:57:50 -07:00 · 2025-08-27 01:30:29 +03:00
232 changed files with 6970 additions and 224647 deletions
@@ -121,7 +121,7 @@ runs:
        echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
        echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' | sudo tee -a /etc/apt/apt.conf.d/99keep-debs
-    
+
    - name: Add OpenCL Repo
      if: inputs.opencl == 'true' && runner.os == 'Linux'
      shell: bash
@@ -174,7 +174,7 @@ runs:
        if [[ "${{ inputs.llvm }}" == "true" ]]; then
          pkgs+=" libllvm20 clang-20 lld-20"
        fi
-        
+
        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"

@@ -183,21 +183,21 @@ runs:
      uses: actions/cache@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.APT_CACHE_VERSION }}

    - name: Run apt Update + Install
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
      shell: bash
      run: |
        sudo apt -qq update || true
-        
+
        # ******** do install ********
        if [[ -n "${{ steps.apt-pkgs.outputs.pkgs }}" ]]; then
          sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
        fi
-        
+
        sudo chown -R $USER:$USER /var/cache/apt/archives/
-    
+
    # **** AMD ****
    - name: Setup AMD (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
@@ -234,7 +234,7 @@ runs:
        cache-name: cache-gpuocelot-build
      with:
        path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-0
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.BUILD_CACHE_VERSION }}
    - name: Clone/compile gpuocelot
      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
      shell: bash
@@ -62,10 +62,8 @@ jobs:
      run: BENCHMARK_LOG=stable_diffusion_xl 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: Run huggingface_onnx test
-      run: METAL=1 python3.11 extra/huggingface_onnx/run_models.py test --debug FacebookAI/xlm-roberta-large
    - name: Test speed vs torch
-      run: BIG=2 MPS=1 python3.11 test/test_speed_v_torch.py | tee torch_speed.txt
+      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test tensor cores
      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test AMX tensor cores
@@ -189,7 +187,7 @@ jobs:
    - name: Run model inference benchmark
      run: NV=1 CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
    - name: Test speed vs torch
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
+      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
      run: NV=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test benchmark allreduce
@@ -325,7 +323,7 @@ jobs:
      run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA Training)
@@ -391,7 +389,7 @@ jobs:
    #- name: Test speed vs torch
    #  run: |
    #    python3 -c "import torch; print(torch.__version__)"
-    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
+    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
      run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test tensor cores
@@ -576,7 +574,7 @@ jobs:
      run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD MLPerf)
@@ -605,18 +603,22 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    - name: validate openpilot 0.9.7
-      run: PYTHONPATH=. FLOAT16=0 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
-    - name: benchmark openpilot 0.9.7
-      run: BENCHMARK_LOG=openpilot_0_9_7 PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_7.txt
-    - name: benchmark openpilot w IMAGE=2 0.9.7
-      run: BENCHMARK_LOG=openpilot_0_9_7_image PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
-    - name: openpilot compile3 0.9.7
-      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx
-    - name: openpilot compile3 0.9.7+ tomb raider
-      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/e8bea2c78ffa92685ece511e9b554122aaf1a79d/selfdrive/modeld/models/supercombo.onnx
-    - name: openpilot dmonitoring compile3 0.9.7
-      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: benchmark openpilot 0.9.9 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    - name: benchmark openpilot 0.9.9 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    - name: benchmark openpilot 0.9.9 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 0.9.9 driving_vision
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.9.9 driving_policy
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.9.9 dmonitoring
+      run: PYTHONPATH="." 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="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
+        PYTHONPATH="." 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: |
        # generate quantized weights
@@ -649,10 +651,10 @@ jobs:
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
-    - name: Remove amdgpu
-      run: sudo rmmod amdgpu || true
-    - name: Cleanup running AM processes
-      run: python extra/amdpci/am_smi.py --pids --kill
+    - name: Remove amd modules
+      run: ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -688,11 +690,12 @@ jobs:
      run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
    - name: Run full CIFAR training w 1 GPU
      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
+    # TODO: enable
+    # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #   run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
+      run: BENCHMARK_LOG=bert_10steps AMD=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 | tee am_train_bert_one_gpu.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AM Driver)
@@ -716,7 +719,9 @@ jobs:
    - name: Checkout Code
      uses: actions/checkout@v4
    - name: Remove nv modules
-      run: ./extra/nvpci/nv_smi.py rmmod
+      run: ./extra/hcq/hcq_smi.py nv rmmod
+    - name: Kill stale pids
+      run: ./extra/hcq/hcq_smi.py nv kill_pids
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -743,12 +748,12 @@ jobs:
    - name: Test LLAMA-3
      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
+      run: BENCHMARK_LOG=bert_10steps NV=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 | tee nv_train_bert_one_gpu.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NV Driver)
@@ -12,7 +12,7 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 240
+    timeout-minutes: 360

    steps:
    - name: Checkout Code
@@ -27,4 +27,4 @@ jobs:
      run: |
        rm "~/.cache/tinygrad/cache_mlperf.db" || true
        BENCHMARK_LOG=mlpert_train_resnet LOGMLPERF=0 CACHEDB="~/.cache/tinygrad/cache_mlperf.db" examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-        rm "~/.cache/tinygrad/cache_mlperf.db"
+        rm "~/.cache/tinygrad/cache_mlperf.db"
@@ -1,8 +1,10 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  DOWNLOAD_CACHE_VERSION: '10'
-  PYTHON_CACHE_VERSION: '2'
+  DOWNLOAD_CACHE_VERSION: '12'
+  PYTHON_CACHE_VERSION: '3'
+  APT_CACHE_VERSION: '1'
+  BUILD_CACHE_VERSION: '1'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}

@@ -30,9 +32,9 @@ jobs:
    - name: External Benchmark Schedule
      run: PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
    - name: Speed Test
-      run: LLVM=1 python3 test/test_speed_v_torch.py
+      run: LLVM=1 python3 test/speed/external_test_speed_v_torch.py
    - name: Speed Test (BEAM=2)
-      run: BEAM=2 LLVM=1 python3 test/test_speed_v_torch.py
+      run: BEAM=2 LLVM=1 python3 test/speed/external_test_speed_v_torch.py

  docs:
    name: Docs
@@ -46,6 +48,11 @@ jobs:
      with:
        deps: docs
        pydeps: "capstone"
+    - name: Build wheel and show size
+      run: |
+        pip install build
+        python -m build --wheel --outdir dist
+        ls -lh dist/*.whl
    - name: Use as an external package
      run: |
        mkdir $HOME/test_external_dir
@@ -329,7 +336,6 @@ jobs:
      run: |
        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check .
-        python3 -m ruff check extra/onnx.py extra/onnx_parser.py
        python3 -m ruff check examples/mlperf/ --ignore E501
    - name: Lint tinygrad with pylint
      run: python -m pylint tinygrad/
@@ -337,8 +343,8 @@ jobs:
      run: |
        python -m mypy --strict-equality --lineprecision-report .
        cat lineprecision.txt
-        python -m mypy --strict-equality extra/onnx_parser.py
-        python -m mypy --strict-equality extra/onnx.py
+    - name: Run TYPED=1
+      run: TYPED=1 python -c "import tinygrad"

  unittest:
    name: Unit Tests
@@ -376,8 +382,8 @@ jobs:
        PYTHONPATH=. python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
        DEBUG=1 MIN_ASTS=1 PYTHONPATH=. python extra/optimization/get_action_space.py
-    - name: Repo line count < 15500 lines
-      run: MAX_LINE_COUNT=15500 python sz.py
+    - name: Repo line count < 17500 lines
+      run: MAX_LINE_COUNT=17500 python sz.py

  fuzzing:
    name: Fuzzing
@@ -456,7 +462,7 @@ jobs:
  testopenpilot:
    name: 'openpilot Compile Tests'
    runs-on: ubuntu-22.04
-    timeout-minutes: 10
+    timeout-minutes: 15
    env:
      IGNORE_OOB: 0
    steps:
@@ -543,8 +549,8 @@ jobs:
        run: PYTHONPATH="." GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
      - name: Test MLPerf stuff
        run: GPU=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
-      - name: Run handcode_opt
-        run: PYTHONPATH=. MODEL=resnet GPU=1 DEBUG=1 BS=4 HALF=0 python3 examples/handcode_opt.py
+      - name: Test llama 3 training
+        run: MAX_BUFFER_SIZE=0 PYTHONPATH="." DEV=NULL SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -587,6 +593,56 @@ jobs:
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

+  testrangeify:
+    name: Linux (rangeify)
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: rangeify-minimal-llvm
+        deps: testing_minimal
+        llvm: "true"
+    - name: Test CPU=1 RANGEIFY=1
+      # TODO: add more passing tests here
+      # test_symbolic_arange_sym_step is passing now
+      # test_threefry_doesnt_use_long is because there's a contig after the long now
+      run: |
+        CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \
+          -k "not test_symbolic_arange_sym_step and not test_threefry_doesnt_use_long" \
+          test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_tensor_variable.py \
+          test/test_outerworld_range.py test/test_sample.py test/test_randomness.py test/test_tensor_data.py
+    - name: Test CPU=1 RANGEIFY=2
+      run: CPU=1 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
+    - name: Test LLVM=1 RANGEIFY=1 (slow tests)
+      run: LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20
+
+  testdevectorize:
+    name: Linux (devectorize)
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: devectorize-minimal
+        deps: testing_minimal
+        pydeps: "pillow"
+        llvm: "true"
+    - name: Test LLVM=1 DEVECTORIZE=0
+      run: LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+    - name: Test LLVM=1 DEVECTORIZE=0 for model
+      run: PYTHONPATH="." LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
+    - name: Test CPU=1 DEVECTORIZE=0
+      run: CPU=1 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+
  testdsp:
    name: Linux (DSP)
    runs-on: ubuntu-24.04
@@ -622,12 +678,6 @@ jobs:
      run: CC=clang-20 PYTHONPATH="." DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
    - name: Test quantize onnx
      run: PYTHONPATH="." DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py
-    - name: Test LLVM=1 DEVECTORIZE=0
-      run: LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
-    - name: Test LLVM=1 DEVECTORIZE=0 for model
-      run: PYTHONPATH="." LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
-    - name: Test CPU=1 DEVECTORIZE=0
-      run: CPU=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"

  testwebgpu:
    name: Linux (WebGPU)
@@ -687,9 +737,9 @@ jobs:
          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
      - name: Run LLVM test
        if: matrix.backend=='amdllvm'
-        run: python test/test_amd_llvm.py
+        run: python test/device/test_amd_llvm.py
      - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/test_hcq.py --durations=20
+        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py --durations=20
      - name: Run pytest (amd)
        run: python -m pytest test/external/external_test_am.py --durations=20
      - name: Run TRANSCENDENTAL math
@@ -814,14 +864,14 @@ jobs:
        AMD: 1
        FORWARD_ONLY: 1
      run: |
-        python3 -m pytest -n=auto test/test_hcq.py test/test_tiny.py --durations=20
+        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
    - name: Run pytest (amd with llvm backend)
      env:
        MOCKGPU: 1
        AMD: 1
        FORWARD_ONLY: 1
      run: |
-        python -m pytest -n=auto test/test_hcq.py test/test_tiny.py test/test_amd_llvm.py --durations=20
+        python -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py test/device/test_amd_llvm.py --durations=20
    - name: Run pytest (ptx)
      env:
        MOCKGPU: 1
@@ -829,7 +879,7 @@ jobs:
        NV: 1
        FORWARD_ONLY: 1
      run: |
-        python3 -m pytest -n=auto test/test_hcq.py test/test_tiny.py --durations=20
+        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -871,28 +921,28 @@ jobs:
      run: WEBGPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py

  osxremote:
-    name: MacOS (remote metal)
-    runs-on: macos-15
-    timeout-minutes: 10
-    env:
-      REMOTE: 1
-      REMOTEDEV: METAL
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: macos-remote
-          deps: testing_minimal
-      - name: Check Device.DEFAULT and print some source
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
-          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'METAL', Device.default.properties.real_device"
-          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-      - name: Run REMOTE=1 Test
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
+   name: MacOS (remote metal)
+   runs-on: macos-15
+   timeout-minutes: 10
+   env:
+     REMOTE: 1
+     REMOTEDEV: METAL
+   steps:
+     - name: Checkout Code
+       uses: actions/checkout@v4
+     - name: Setup Environment
+       uses: ./.github/actions/setup-tinygrad
+       with:
+         key: macos-remote
+         deps: testing_minimal
+     - name: Check Device.DEFAULT and print some source
+       run: |
+         python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
+         python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'METAL', Device.default.properties.real_device"
+         DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
+     - name: Run REMOTE=1 Test
+       run: |
+         python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py

  amdremote:
    name: Linux (remote)
@@ -940,18 +990,18 @@ jobs:
        env:
          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py --durations 20
      - name: Run REMOTE=1 Test (GPU)
        env:
          HOST: 127.0.0.1:7667*6
        run: |
-          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
+          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py --durations 20
          IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
      - name: Run REMOTE=1 Test (CPU)
        env:
          HOST: 127.0.0.1:8667*6
        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py --durations 20
      - name: Show remote server logs
        if: always()
        run: |
@@ -978,6 +1028,7 @@ jobs:
        with:
          key: macos-${{ matrix.backend }}-minimal
          deps: testing_minimal
+          pydeps: "capstone"
          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
@@ -198,11 +198,7 @@ generate_amd() {
  clang2py -k cdefstum \
    extra/hip_gpu_driver/sdma_registers.h \
    extra/hip_gpu_driver/nvd.h \
-    extra/hip_gpu_driver/kfd_pm4_headers_ai.h \
-    extra/hip_gpu_driver/soc21_enum.h \
-    extra/hip_gpu_driver/sdma_v6_0_0_pkt_open.h \
    extra/hip_gpu_driver/gc_11_0_0_offset.h \
-    extra/hip_gpu_driver/gc_10_3_0_offset.h \
    extra/hip_gpu_driver/sienna_cichlid_ip_offset.h \
    --clang-args="-I/opt/rocm/include -x c++" \
    -o $BASE/amd_gpu.py
@@ -376,26 +372,6 @@ generate_am() {
    -o $BASE/am/pm4_nv.py
  fixup $BASE/am/pm4_nv.py

-  clang2py -k cdefstum \
-    $AMKERN_INC/vega10_enum.h \
-    -o $BASE/am/vega10.py
-  fixup $BASE/am/vega10.py
-
-  clang2py -k cdefstum \
-    $AMKERN_INC/navi10_enum.h \
-    -o $BASE/am/navi10.py
-  fixup $BASE/am/navi10.py
-
-  clang2py -k cdefstum \
-    $AMKERN_INC/soc21_enum.h \
-    -o $BASE/am/soc21.py
-  fixup $BASE/am/soc21.py
-
-  clang2py -k cdefstum \
-    $AMKERN_INC/soc24_enum.h \
-    -o $BASE/am/soc24.py
-  fixup $BASE/am/soc24.py
-
  clang2py -k cdefstum \
    extra/hip_gpu_driver/sdma_registers.h \
    $AMKERN_AMD/amdgpu/vega10_sdma_pkt_open.h \
@@ -18,11 +18,11 @@ Group UOps into kernels.

 ---

-## tinygrad/opt
+## tinygrad/codegen/opt

 Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.

-::: tinygrad.opt.get_optimized_ast
+::: tinygrad.codegen.opt.get_optimized_ast
    options:
        members: false
        show_labels: false
@@ -126,7 +126,7 @@ print(t_log_grad.uop)
 """
 void E_(float* restrict data0, float* restrict data1) {
  float val0 = *(data1+0);
-  *(data0+0) = (0.6931471805599453f*(1/(val0*0.6931471805599453f)));
+  *(data0+0) = (1/val0);
 }
 """
 # the derivative is close to 1/3
@@ -78,6 +78,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.minimum
 ::: tinygrad.Tensor.where
 ::: tinygrad.Tensor.copysign
+::: tinygrad.Tensor.logaddexp

 ## Casting Ops

@@ -6,7 +6,7 @@ If you don't have a tinybox and you want one, see [tinygrad.org](https://tinygra

 ## Welcome

-Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, and the green box includes six 4090 GPUs. Whether you bought a red one or a green one, we want you to love it.
+Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, the green box includes six 4090 GPUs, and the green v2 box includes four 5090 GPUs. Whether you bought a red one or a green one, we want you to love it.

 We don't have a stupid cloud service, you don't have to create a tiny account to set it up, and we aren't tracking how you use the box. We're just happy you bought one. This petaflop is your petaflop.

@@ -47,8 +47,8 @@ Reboot after making these changes or restart the `displayservice.service` servic

 The [default tinybox image](https://github.com/tinygrad/tinyos) ships with tinygrad and PyTorch. While we develop tinygrad, the box is universal hardware. Use whatever framework you desire, run notebooks, download demos, install more things, train, inference, live, laugh, love, you aren't paying per hour for this box so the only limit is your imagination.

-## tinychat
+## Building the OS image

-Since LLMs are so popular, we ship with a built in tinygrad based chatbot using a LLaMA-3 finetune. Visit the IP (not the BMC IP) of your tinybox in a web browser on your computer or phone, and you'll find a friendly looking chat interface. This chatbot also provides an OpenAI compatible LLM API on that port, so you can script it.
+The OS image is built using `ubuntu-image` from <https://github.com/tinygrad/tinyos>.

-The conversations you have with this chatbot are between you and your tinybox. Also, the history in the web app is saved on the client, not the tinybox.
+After cloning, run `make green` or `make red` to build a tinybox green or tinybox red image respectively.
@@ -1,12 +1,12 @@
 # model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
-from typing import List, Callable
+from typing import Callable
 from tinygrad import Tensor, TinyJit, nn, GlobalCounters
 from tinygrad.helpers import getenv, colored, trange
 from tinygrad.nn.datasets import mnist

 class Model:
  def __init__(self):
-    self.layers: List[Callable[[Tensor], Tensor]] = [
+    self.layers: list[Callable[[Tensor], Tensor]] = [
      nn.Conv2d(1, 32, 5), Tensor.relu,
      nn.Conv2d(32, 32, 5), Tensor.relu,
      nn.BatchNorm(32), Tensor.max_pool2d,
@@ -21,17 +21,15 @@ if __name__ == "__main__":
  X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))

  model = Model()
-  opt = nn.optim.Adam(nn.state.get_parameters(model))
+  opt = (nn.optim.Adam if not getenv("MUON") else nn.optim.Muon)(nn.state.get_parameters(model))

  @TinyJit
  @Tensor.train()
  def train_step() -> Tensor:
    opt.zero_grad()
    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
-    # TODO: this "gather" of samples is very slow. will be under 5s when this is fixed
    loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
-    opt.step()
-    return loss
+    return loss.realize(*opt.schedule_step())

  @TinyJit
  def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
@@ -1,134 +0,0 @@
-from extra.models.resnet import ResNet50
-from extra.mcts_search import mcts_search
-from examples.mlperf.helpers import get_mlperf_bert_model
-from tinygrad import Tensor, Device, dtypes, nn
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.heuristic import hand_coded_optimizations
-from tinygrad.uop.ops import Ops, sym_infer
-from tinygrad.device import Compiled
-from tinygrad.opt.search import beam_search, bufs_from_lin
-from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
-from extra.optimization.helpers import time_linearizer
-from tinygrad.engine.realize import get_program
-
-def get_sched_resnet():
-  mdl = ResNet50()
-  optim = (nn.optim.LARS if getenv("LARS") else nn.optim.SGD)(nn.state.get_parameters(mdl))
-  BS = getenv("BS", 64)
-
-  # run model twice to get only what changes, these are the kernels of the model
-  for _ in range(2):
-    out = mdl(Tensor.empty(BS, 3, 224, 224))
-    targets = [out]
-    if getenv("BACKWARD"):
-      optim.zero_grad()
-      out.sparse_categorical_crossentropy(Tensor.empty(BS, dtype=dtypes.int)).backward()
-      targets += [x for x in optim.schedule_step()]
-    sched = Tensor.schedule(*targets)
-    print(f"schedule length {len(sched)}")
-  return sched
-
-def get_sched_bert():
-  mdl = get_mlperf_bert_model()
-  optim = nn.optim.LAMB(nn.state.get_parameters(mdl))
-
-  # fake data
-  BS = getenv("BS", 9)
-  input_ids = Tensor.empty((BS, 512), dtype=dtypes.float32)
-  segment_ids = Tensor.empty((BS, 512), dtype=dtypes.float32)
-  attention_mask = Tensor.empty((BS, 512), dtype=dtypes.default_float)
-  masked_positions = Tensor.empty((BS, 76), dtype=dtypes.float32)
-  masked_lm_ids = Tensor.empty((BS, 76), dtype=dtypes.float32)
-  masked_lm_weights = Tensor.empty((BS, 76), dtype=dtypes.float32)
-  next_sentence_labels = Tensor.empty((BS, 1), dtype=dtypes.float32)
-
-  # run model twice to get only what changes, these are the kernels of the model
-  for _ in range(2):
-    lm_logits, seq_relationship_logits = mdl(input_ids, attention_mask, masked_positions, segment_ids)
-    targets = [lm_logits, seq_relationship_logits]
-    if getenv("BACKWARD"):
-      optim.zero_grad()
-      loss = mdl.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
-      # ignore grad norm and loss scaler for now
-      loss.backward()
-      targets += [x for x in optim.schedule_step()]
-    sched = Tensor.schedule(*targets)
-    print(f"schedule length {len(sched)}")
-  return sched
-
-if __name__ == "__main__":
-  if getenv("HALF", 1):
-    dtypes.default_float = dtypes.half
-
-  # the device we are optimizing for
-  device: Compiled = Device[Device.DEFAULT]
-  if getenv("BACKWARD"): Tensor.training = True
-  print(f"optimizing for {Device.DEFAULT}")
-
-  sched = globals()[f"get_sched_{getenv('MODEL', 'resnet')}"]()
-  sched = [x for x in sched if x.ast.op is Ops.SINK]
-
-  # focus on one kernel
-  if getenv("KERNEL", -1) >= 0: sched = sched[getenv("KERNEL", -1):getenv("KERNEL", -1)+1]
-
-  # work with the schedule
-  total_tm = 0
-  running_gflops = 0
-  usage = {}
-  for i,si in enumerate(sched):
-    if DEBUG >= 3: print(si.ast)
-
-    rawbufs = bufs_from_lin(Kernel(si.ast))
-
-    # "linearize" the op into uops in different ways
-    lins: list[tuple[Kernel, str]] = []
-
-    # always try hand coded opt
-    lin = Kernel(si.ast, opts=device.renderer)
-    lin.apply_opts(hand_coded_optimizations(lin))
-    lins.append((lin, "HC"))
-
-    # maybe try tensor cores
-    lin = Kernel(si.ast, opts=device.renderer)
-    if lin.apply_tensor_cores():
-      lins.append((lin, "TC"))
-
-    # try a beam search
-    if beam:=getenv("BEAM"):
-      lin = Kernel(si.ast, opts=device.renderer)
-      lin = beam_search(lin, rawbufs, beam, bool(getenv("BEAM_ESTIMATE", 1)))
-      lins.append((lin, "BEAM"))
-
-    # try MCTS
-    if mcts:=getenv("MCTS"):
-      lin = Kernel(si.ast, opts=device.renderer)
-      lin = mcts_search(lin, rawbufs, mcts)
-      lins.append((lin, "MCTS"))
-
-    # benchmark the programs
-    choices = []
-    for lin, nm in lins:
-      tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
-      ops = (prg:=get_program(lin.get_optimized_ast(), lin.opts)).estimates.ops
-      gflops = sym_infer(ops, {k:k.min for k in lin.ast.variables()})*1e-9/tm
-      choices.append((tm, gflops, lin, prg, nm))
-
-    sorted_choices = sorted(choices, key=lambda x: x[0])
-    if DEBUG >= 1: # print all kernels
-      for tm, gflops, lin, prg, nm in choices:
-        print(f"                 kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(prg.global_size):18s} {str(prg.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS -- {colored(nm, 'green') if lin is sorted_choices[0][2] else nm}")
-
-    tm, gflops, lin, prg, nm = sorted_choices[0]
-    if getenv("SRC"):
-      print(si.ast)
-      print(lin.applied_opts)
-      print(get_program(lin.get_optimized_ast(), lin.opts).src)
-    total_tm += tm
-    running_gflops += gflops * tm
-    if (key := str([str(m) for m in si.metadata])) not in usage: usage[key] = (0, 0)
-    usage[key] = (usage[key][0] + tm, usage[key][1] + 1)
-    print(f"*** {total_tm*1000:7.2f} ms : kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(prg.global_size):18s} {str(prg.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS {[repr(m) if TRACEMETA >= 2 else str(m) for m in si.metadata]}")
-  print(f"******* total {total_tm*1000:.2f} ms, {running_gflops/total_tm:6.0f} GFLOPS")
-  print("usage:")
-  for k in sorted(usage, key=lambda x: -usage[x][0])[:10]:
-    print(f"{usage[k][0]*1000:.2f} ms: {k} ({usage[k][1]} times)")
@@ -118,7 +118,7 @@ class SpeedyResNet:
 # hyper-parameters were exactly the same as the original repo
 bias_scaler = 58
 hyp = {
-  'seed' : 200,
+  'seed' : 201,
  'opt': {
    'bias_lr':            1.76 * bias_scaler/512,
    'non_bias_lr':        1.76 / 512,
@@ -1,4 +1,4 @@
-import os, random, pickle, queue
+import os, random, pickle, queue, struct, math, functools, hashlib, time
 from typing import List
 from pathlib import Path
 from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
@@ -6,6 +6,7 @@ from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu
 import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX
+from tinygrad.nn.state import TensorIO

 ### ResNet

@@ -510,6 +511,253 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
      # happens with BENCHMARK set
      pass

+# llama3
+
+class BinIdxDataset:
+  def __init__(self, base_path:Path):
+    self.idx_t = Tensor(base_path.with_name(f"{base_path.name}.idx"))
+    self.idx = TensorIO(self.idx_t)
+
+    # parse idx file
+    magic = self.idx.read(9)
+    assert magic == b"MMIDIDX\x00\x00", "invalid index file format"
+    version, = struct.unpack("<Q", self.idx.read(8))
+    assert version == 1, "unsupported index version"
+    dtype_code, = struct.unpack("<B", self.idx.read(1))
+    self.dtype = {1:dtypes.uint8, 2:dtypes.int8, 3:dtypes.int16, 4:dtypes.int32, 5:dtypes.int64, 6:dtypes.float64, 7:dtypes.double, 8:dtypes.uint16}[dtype_code]
+    self.count, = struct.unpack("<Q", self.idx.read(8))
+    doc_count, = struct.unpack("<Q", self.idx.read(8))
+
+    start = self.idx.tell()
+    end = start + self.count * dtypes.int32.itemsize
+    self.sizes = self.idx_t[start:end].bitcast(dtypes.int32).numpy()
+
+    start = end
+    end = start + self.count * dtypes.int64.itemsize
+    self.pointers = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
+
+    start = end
+    end = start + doc_count * dtypes.int64.itemsize
+    self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
+
+    # bin file
+    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))
+
+  def _index(self, idx) -> tuple[int, int]:
+    return int(self.pointers[idx]), int(self.sizes[idx])
+
+  def get(self, idx, offset:int=0, length:int|None=None):
+    ptr, size = self._index(idx)
+    if length is None: length = size - offset
+    ptr += offset * self.dtype.itemsize
+    return self.bin_t[ptr:ptr+length*self.dtype.itemsize].bitcast(self.dtype).to(None)
+
+# https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
+class GPTDataset:
+  def __init__(self, base_path:Path, samples:int, seqlen:int, seed:int, shuffle:bool):
+    self.samples, self.seqlen = samples, seqlen
+    self.shuffle = shuffle
+    self.rng = np.random.RandomState(seed)
+
+    self.indexed_dataset = BinIdxDataset(base_path)
+
+    # check for cache
+    cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
+    cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
+    print(f"try loading GPTDataset from {cache_path}...")
+    if cache_path.exists():
+      print("cache found, loading...")
+      with open(cache_path, "rb") as f:
+        self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
+    else:
+      print("cache not found, building index...")
+      self.doc_idx = self._build_doc_idx()
+      self.sample_idx = self._build_sample_idx()
+      self.shuffle_idx = self._build_shuffle_idx()
+      # save cache
+      with open(cache_path, "wb") as f:
+        pickle.dump((self.doc_idx, self.sample_idx, self.shuffle_idx), f)
+
+  def __getitem__(self, idx):
+    if idx is None:
+      text = self._get(0)
+    else:
+      text = self._get(idx)
+
+    return text
+
+  def _get(self, idx):
+    idx = self.shuffle_idx[idx]
+
+    doc_idx_beg, doc_idx_beg_offset = self.sample_idx[idx]
+    doc_idx_end, doc_idx_end_offset = self.sample_idx[idx + 1]
+
+    doc_ids, sample_parts = [], []
+
+    if doc_idx_beg == doc_idx_end:
+      doc_ids.append(self.doc_idx[doc_idx_beg])
+
+      sample_parts.append(
+          self.indexed_dataset.get(
+            int(self.doc_idx[doc_idx_beg]), offset=int(doc_idx_beg_offset), length=int(doc_idx_end_offset - doc_idx_beg_offset + 1)))
+    else:
+      for i in range(doc_idx_beg, doc_idx_end + 1):
+        doc_ids.append(self.doc_idx[i])
+
+        offset = 0 if i > doc_idx_beg else doc_idx_beg_offset
+        length = None if i < doc_idx_end else int(doc_idx_end_offset + 1)
+        sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))
+
+    # concat all parts
+    text = Tensor.cat(*sample_parts)
+
+    return text
+
+  @functools.cached_property
+  def tokens_per_epoch(self) -> int:
+    return sum(self.indexed_dataset.sizes.tolist())
+
+  @functools.cached_property
+  def num_epochs(self) -> int:
+    # we need enough epochs to cover the requested amount of tokens
+    num_epochs = 1
+    num_tokens = self.tokens_per_epoch
+    while num_tokens < self.samples * self.seqlen:
+      num_epochs += 1
+      num_tokens += self.tokens_per_epoch
+    return num_epochs
+
+  # https://github.com/NVIDIA/Megatron-LM/blob/94bd476bd840c2fd4c3ebfc7448c2af220f4832b/megatron/core/datasets/gpt_dataset.py#L558
+  def _build_doc_idx(self):
+    print(f"building doc_idx for {self.num_epochs=}, {self.indexed_dataset.count=}")
+    st = time.perf_counter()
+    # doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
+    doc_idx = np.arange(self.indexed_dataset.count).reshape(1, -1).repeat(self.num_epochs, axis=0).flatten()
+    doc_idx = doc_idx.astype(np.int32)
+    at = time.perf_counter()
+    if self.shuffle: self.rng.shuffle(doc_idx)
+    print(f"doc_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
+    return doc_idx
+
+  def _build_sample_idx(self):
+    print(f"building sample_idx for {self.samples=}, {self.seqlen=}, {self.doc_idx.shape[0]=}")
+    sample_idx_max = max(self.doc_idx.shape[0], self.indexed_dataset.sizes.max())
+    sample_idx = np.empty((self.samples + 1, 2), dtype=np.int64 if sample_idx_max > dtypes.int32.max else np.int32)
+
+    sample_idx_idx, doc_idx_idx, doc_offset = 0, 0, 0
+    sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
+    sample_idx_idx += 1
+
+    for _ in tqdm(range(1, self.samples + 1)):
+      remaining_seqlen = self.seqlen + 1
+      while remaining_seqlen > 0:
+        doc_idx = int(self.doc_idx[doc_idx_idx])
+        doc_len = int(self.indexed_dataset.sizes[doc_idx]) - doc_offset
+        remaining_seqlen -= doc_len
+        if remaining_seqlen <= 0:
+          doc_offset += remaining_seqlen + doc_len - 1
+          remaining_seqlen = 0
+        else:
+          if doc_idx_idx == len(self.doc_idx) - 1:
+            assert sample_idx_idx == self.samples
+            doc_idx = int(self.doc_idx[doc_idx_idx])
+            doc_offset = int(self.indexed_dataset.sizes[doc_idx]) - 1
+            break
+          doc_idx_idx += 1
+          doc_offset = 0
+
+      sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
+      sample_idx_idx += 1
+
+    return sample_idx
+
+  def _build_shuffle_idx(self):
+    print(f"building shuffle_idx for {self.samples=}")
+    st = time.perf_counter()
+    shuffle_idx = np.arange(self.samples, dtype=np.int32)
+    at = time.perf_counter()
+    if self.shuffle: self.rng.shuffle(shuffle_idx)
+    print(f"shuffle_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
+    return shuffle_idx
+
+class BlendedGPTDataset:
+  def __init__(self, paths:list[Path], weights:list[float], samples:int, seqlen:int, seed:int, shuffle:bool):
+    self.shuffle = shuffle
+    self.rng = np.random.RandomState(seed)
+
+    # normalize weights
+    total_weight = sum(weights)
+    self.weights = [w / total_weight for w in weights]
+
+    self.samples = samples
+    surplus = 0.005
+    samples_per_blend = [math.ceil(math.ceil(self.samples * w) * (1 + surplus)) for w in self.weights]
+
+    self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]
+
+    # check for cache
+    cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
+    cache_path = paths[0].with_name(f"{paths[0].name}.{cache_hash}.blend_cache")
+    print(f"try loading BlendedGPTDataset from {cache_path}...")
+    if cache_path.exists():
+      print("cache found, loading...")
+      with open(cache_path, "rb") as f:
+        self.dataset_idx, self.dataset_sample_idx = pickle.load(f)
+    else:
+      print("cache not found, building index...")
+      self.dataset_idx, self.dataset_sample_idx = self._build_blend_idx()
+      # save cache
+      with open(cache_path, "wb") as f:
+        pickle.dump((self.dataset_idx, self.dataset_sample_idx), f)
+
+  def get(self, idx:int):
+    tokens = self.datasets[self.dataset_idx[idx]][self.dataset_sample_idx[idx]]
+    return tokens
+
+  def _build_blend_idx(self):
+    dataset_idx = np.zeros(self.samples, dtype=np.int16)
+    dataset_sample_idx = np.zeros(self.samples, dtype=np.int64)
+
+    unspent_datasets = set(range(len(self.datasets)))
+    dataset_sample_counts = [0] * len(self.datasets)
+
+    for i in tqdm(range(self.samples)):
+      error_argmax, error_max = 0, 0.0
+      for di in unspent_datasets:
+        error = self.weights[di] * max(i, 1) - dataset_sample_counts[di]
+        if error > error_max:
+          error_max = error
+          error_argmax = di
+
+      dataset_idx[i] = error_argmax
+      dataset_sample_idx[i] = dataset_sample_counts[error_argmax]
+
+      dataset_sample_counts[error_argmax] += 1
+
+    return dataset_idx, dataset_sample_idx
+
+def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
+  if val:
+    dataset = BlendedGPTDataset([
+      base_dir / "validation" / "c4-validationn-91205-samples.en_text_document",
+    ], [
+      1.0
+    ], samples, seqlen, seed, False)
+  else:
+    dataset = BlendedGPTDataset([
+      base_dir / "c4-train.en_6_text_document",
+      base_dir / "c4-train.en_7_text_document",
+    ], [
+      1.0, 1.0
+    ], samples, seqlen, seed, True)
+
+  for b in range(math.ceil(samples / bs)):
+    batch = []
+    for i in range(bs):
+      tokens = dataset.get(b * bs + i)
+      batch.append(tokens)
+    yield Tensor.stack(batch, dim=0)
+
 if __name__ == "__main__":
  def load_unet3d(val):
    assert not val, "validation set is not supported due to different sizes on inputs"
@@ -538,6 +786,18 @@ if __name__ == "__main__":
      for x in batch_load_retinanet(dataset, val, base_dir):
        pbar.update(x[0].shape[0])

+  def load_llama3(val):
+    bs = 24
+    samples = 5760 if val else 1_200_000 * 1152
+    seqlen = 8192
+
+    max_, min_ = 0, math.inf
+    for tokens in tqdm(batch_load_llama3(bs, samples, seqlen, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=5760, val=bool(val)), total=samples//bs):
+      max_ = max(max_, tokens.shape[1])
+      min_ = min(min_, tokens.shape[1])
+    print(f"max seq length: {max_}")
+    print(f"min seq length: {min_}")
+
  load_fn_name = f"load_{getenv('MODEL', 'resnet')}"
  if load_fn_name in globals():
    globals()[load_fn_name](getenv("VAL", 1))
@@ -1,4 +1,4 @@
-import time
+import time, math
 start = time.perf_counter()
 from pathlib import Path
 import numpy as np
@@ -241,6 +241,34 @@ def eval_mrcnn():
  evaluate_predictions_on_coco(bbox_output, iou_type='bbox')
  evaluate_predictions_on_coco(mask_output, iou_type='segm')

+def eval_llama3():
+  from extra.models.llama import Transformer
+  from examples.llama3 import MODEL_PARAMS
+  from tinygrad.helpers import tqdm
+
+  bs = 4
+  sequence_length = 512
+
+  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
+
+  @TinyJit
+  def eval_step(model, tokens):
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    return loss.flatten()
+
+  from examples.mlperf.dataloader import batch_load_llama3
+  iter = batch_load_llama3(bs, 5760, sequence_length, Path(getenv("BASEDIR", "/raid/datasets/c4/")), True)
+
+  losses = []
+  for tokens in tqdm(iter, total=5760//bs):
+    GlobalCounters.reset()
+    losses += eval_step(model, tokens).tolist()
+    tqdm.write(f"loss: {np.mean(losses)}")
+
+  log_perplexity = Tensor(losses).mean()
+  print(f"Log Perplexity: {log_perplexity.item()}")
+
 if __name__ == "__main__":
  # inference only
  Tensor.training = False
@@ -1290,9 +1290,19 @@ def train_llama3():
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup

  config = {}
-  BS                 = config["BS"]                     = getenv("BS", 4)
+  BS                 = config["BS"]                     = getenv("BS", 16)
  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
+  SEED               = config["SEED"]                   = getenv("SEED", 5760)
+  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
+  TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
+  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
+  EVAL_FREQ          = config["EVAL_FREQ"]              = getenv("EVAL_FREQ", 46080)
+  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 16)
+  EVAL_TARGET        = config["EVAL_TARGET"]            = getenv("EVAL_TARGET", 5.6)
+
+  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
+  # trains to 7

  opt_adamw_beta_1 = 0.9
  opt_adamw_beta_2 = 0.95
@@ -1300,7 +1310,6 @@ def train_llama3():
  opt_adamw_weight_decay = 0.1

  opt_gradient_clip_norm = 1.0
-  sequence_length = 8192
  opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
  opt_learning_rate_decay_steps = getenv("DECAY_STEPS", math.ceil(1_200_000 * 1152 / GBS) - opt_learning_rate_warmup_steps)
  opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
@@ -1308,7 +1317,37 @@ def train_llama3():

  # TODO: confirm weights are in bf16
  # vocab_size from the mixtral tokenizer
-  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
+  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
+  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+
+  if getenv("FAKEDATA"):
+    for v in get_parameters(model):
+      v = v.assign(Tensor.empty(v.shape))
+
+  if (DP := getenv("DP", 1)) > 1:
+    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+    for v in get_parameters(model):
+      v.shard_(device, axis=None)
+
+  if (MP := getenv("MP", 1)) > 1:
+    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+    for k,v in get_state_dict(model).items():
+      if 'scale' in k: v.shard_(device, axis=None)  # from quantized
+      elif '.attention.wq' in k: v.shard_(device, axis=0)
+      elif '.attention.wk' in k: v.shard_(device, axis=0)
+      elif '.attention.wv' in k: v.shard_(device, axis=0)
+      elif '.attention.wo' in k: v.shard_(device, axis=1)
+      elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
+      elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
+      elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
+      elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
+      elif 'output.weight' in k: v.shard_(device, axis=0)
+      else:
+        # attention_norm, ffn_norm, norm
+        v.shard_(device, axis=None)
+      # prevents memory spike on device 0
+      v.realize()

  optim = AdamW(get_parameters(model), lr=0.0,
                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
@@ -1316,12 +1355,20 @@ def train_llama3():

  @TinyJit
  @Tensor.train()
-  def train_step(model, x, y):
+  def train_step(model, tokens:Tensor, grad_acc:int):
    optim.zero_grad()
-    logits:Tensor = model(x, start_pos=0, temperature=math.nan)
-    loss = logits.cross_entropy(y)
-    loss.backward()
-
+    # grad acc
+    for batch in tokens.split(tokens.shape[0]//grad_acc):
+      if (DP := getenv("DP", 1)) > 1:
+        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+        batch = batch.shard(device, 0)
+      if (MP := getenv("MP", 1)) > 1:
+        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+        batch = batch.shard(device)
+      logits:Tensor = model(batch[:, :-1], start_pos=0, temperature=math.nan)
+      loss = logits.sparse_categorical_crossentropy(batch[:, 1:])
+      loss.backward()
+      Tensor.realize(*[p.grad for p in optim.params])
    # L2 norm grad clip
    # https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
    # https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
@@ -1331,7 +1378,7 @@ def train_llama3():
        total_norm += p.grad.float().square().sum()
      total_norm = total_norm.sqrt().contiguous()
      for p in optim.params:
-        p.grad = p.grad * opt_gradient_clip_norm / (total_norm + 1e-6)
+        p.grad = p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)

    optim.step()
    scheduler.step()
@@ -1340,19 +1387,81 @@ def train_llama3():
    loss.realize(lr)
    return loss, lr

-  # overfitting this example should give cross_entropy log(BS)
-  fake_input = Tensor([list(range(getenv("SEQLEN", 10)))], dtype="int16").expand(BS, -1)
-  fake_label = Tensor(list(range(BS)), dtype="int16")
+  @TinyJit
+  @Tensor.train(False)
+  def eval_step(model, tokens:Tensor):
+    if (DP := getenv("DP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+      tokens = tokens.shard(device, 0)
+    if (MP := getenv("MP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+      tokens = tokens.shard(device)
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    return loss.flatten().float()

-  for _ in range(100):
+  # ** data iters **
+  def fake_data(bs, samples):
+    for _ in range(samples // bs):
+      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
+
+  def get_train_iter():
+    if getenv("FAKEDATA", 0):
+      return fake_data(GBS, SAMPLES)
+    else:
+      from examples.mlperf.dataloader import batch_load_llama3
+      return batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))
+
+  def get_eval_iter():
+    if getenv("FAKEDATA", 0):
+      return fake_data(EVAL_BS, 5760)
+    else:
+      from examples.mlperf.dataloader import batch_load_llama3
+      return batch_load_llama3(EVAL_BS, 5760, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=True)
+
+  iter = get_train_iter()
+  i, sequences_seen = 0, 0
+  for tokens in tqdm(iter, total=SAMPLES//GBS):
+    t = time.perf_counter()
    GlobalCounters.reset()
-    loss, lr = train_step(model, fake_input, fake_label)
-    # BS=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
-    # uses 43% ~= 83GB
-    # 8B bf16 = 16GB. model + grad + optim m and v = 64GB
-    # TODO: this OOM
-    # BS=1 SEQLEN=4000 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
-    print(loss.item(), lr.item(), f"{GlobalCounters.global_mem//10**9=}")
+    loss, lr = train_step(model, tokens, grad_acc)
+    loss = loss.float().item()
+    # above as tqdm.write f-string
+    tqdm.write(f"{loss:.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
+    if (fname:=getenv("LOSS_FILE", "")):
+      with open(fname, "a") as f:
+        f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
+
+    if getenv("CKPT") and (i % 200 == 0 or i == 10):
+      tqdm.write("saving checkpoint")
+      if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
+      fn = f"{ckpt_dir}/llama3_{i}.safe"
+      safe_save(get_state_dict(model), fn)
+
+    i += 1
+    sequences_seen += tokens.shape[0]
+
+    if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
+      tqdm.write(f"evaluating after {sequences_seen} sequences")
+
+      # run eval
+      eval_losses = []
+      eval_iter = get_eval_iter()
+      tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
+
+      for tokens in tqdm(eval_iter, total=5760//EVAL_BS):
+        eval_losses += eval_step(model, tokens).tolist()
+      log_perplexity = Tensor(eval_losses).mean().float().item()
+
+      tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
+
+      if log_perplexity < EVAL_TARGET:
+        tqdm.write(f"target achieved after {sequences_seen} sequences")
+        if getenv("CKPT"):
+          if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
+          fn = f"{ckpt_dir}/llama3.safe"
+          safe_save(get_state_dict(model), fn)
+        break

 if __name__ == "__main__":
  multiprocessing.set_start_method('spawn')
@@ -4,6 +4,8 @@ export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128

+export IGNORE_OOB=1
+
 export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 # export BEAM_LOG_SURPASS_MAX=1
@@ -5,6 +5,8 @@ export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1

+export IGNORE_OOB=1
+
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
 export BASEDIR="/raid/datasets/wiki"
@@ -8,6 +8,8 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
 export TRAIN_STEPS=3900

+export IGNORE_OOB=1
+
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
 export BASEDIR="/raid/datasets/wiki"
@@ -11,6 +11,8 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
 export TRAIN_STEPS=3900

+export IGNORE_OOB=1
+
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
 export BASEDIR="/raid/datasets/wiki"
@@ -2,9 +2,9 @@

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

-export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+export IGNORE_OOB=1

 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -2,9 +2,9 @@

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

-export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+export IGNORE_OOB=1

 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -5,9 +5,9 @@ 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=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

-export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+export IGNORE_OOB=1

 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -2,9 +2,9 @@

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

-export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+export IGNORE_OOB=1

 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -2,9 +2,9 @@

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

-export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+export IGNORE_OOB=1

 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -5,9 +5,9 @@ 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=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

-export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+export IGNORE_OOB=1

 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -1,8 +1,7 @@
 # https://arxiv.org/pdf/2409.02060
-import time
+import time, functools
 import numpy as np
 np.set_printoptions(suppress=True, linewidth=1000)
-import functools
 from tinygrad import Tensor, nn, Device, GlobalCounters
 from tinygrad.helpers import Timing, getenv
 from extra.models.llama import Transformer, convert_from_huggingface
@@ -17,7 +16,7 @@ class MixtureFeedForward:
  def __call__(self, x:Tensor) -> Tensor:
    assert x.shape[0] == 1, "only BS=1"
    assert x.shape[1] == 1, "only length=1"
-    g = self.gate(x).float().softmax(-1)
+    g = self.gate(x).softmax(-1)

    g = g.squeeze() # (BS, length, num_experts) -> (num_experts,)
    probs, sel = g.topk(self.activated_experts)
@@ -25,7 +24,7 @@ class MixtureFeedForward:
    # run MoE
    x_up_gate = x.dot(self.gate_proj[sel].permute(0,2,1)).silu() * x.dot(self.up_proj[sel].permute(0,2,1))
    x_down = x_up_gate.dot(self.down_proj[sel].permute(0,2,1))
-    return (x_down.float() * probs.reshape(self.activated_experts, 1, 1)).sum(axis=0)
+    return (x_down * probs.reshape(self.activated_experts, 1, 1)).sum(axis=0)

 # model is bf16, 1.3B active, 6.9B total
 # M3 Max is 400 GB/s, so 400/2.6 = ~154 tok/s
@@ -71,8 +71,8 @@ def bbox_iou(box1, box2):
  # get the coordinates of the intersection rectangle
  inter_rect_x1 = np.maximum(b1_x1, b2_x1)
  inter_rect_y1 = np.maximum(b1_y1, b2_y1)
-  inter_rect_x2 = np.maximum(b1_x2, b2_x2)
-  inter_rect_y2 = np.maximum(b1_y2, b2_y2)
+  inter_rect_x2 = np.minimum(b1_x2, b2_x2)
+  inter_rect_y2 = np.minimum(b1_y2, b2_y2)
  #Intersection area
  inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, 99999) * np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, 99999)
  #Union Area
@@ -297,13 +297,13 @@ class Darknet:
          # Get the number of weights of batchnorm
          num_bn_biases = math.prod(bn.bias.shape)
          # Load weights
-          bn_biases = Tensor(weights[ptr:ptr + num_bn_biases])
+          bn_biases = Tensor(weights[ptr:ptr + num_bn_biases].astype(np.float32))
          ptr += num_bn_biases
-          bn_weights = Tensor(weights[ptr:ptr+num_bn_biases])
+          bn_weights = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
          ptr += num_bn_biases
-          bn_running_mean = Tensor(weights[ptr:ptr+num_bn_biases])
+          bn_running_mean = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
          ptr += num_bn_biases
-          bn_running_var = Tensor(weights[ptr:ptr+num_bn_biases])
+          bn_running_var = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
          ptr += num_bn_biases
          # Cast the loaded weights into dims of model weights
          bn_biases = bn_biases.reshape(shape=tuple(bn.bias.shape))
@@ -319,7 +319,7 @@ class Darknet:
          # load biases of the conv layer
          num_biases = math.prod(conv.bias.shape)
          # Load weights
-          conv_biases = Tensor(weights[ptr: ptr+num_biases])
+          conv_biases = Tensor(weights[ptr: ptr+num_biases].astype(np.float32))
          ptr += num_biases
          # Reshape
          conv_biases = conv_biases.reshape(shape=tuple(conv.bias.shape))
@@ -327,7 +327,7 @@ class Darknet:
          conv.bias = conv_biases
        # Load weighys for conv layers
        num_weights = math.prod(conv.weight.shape)
-        conv_weights = Tensor(weights[ptr:ptr+num_weights])
+        conv_weights = Tensor(weights[ptr:ptr+num_weights].astype(np.float32))
        ptr += num_weights
        conv_weights = conv_weights.reshape(shape=tuple(conv.weight.shape))
        conv.weight = conv_weights
@@ -371,7 +371,7 @@ class Darknet:
 if __name__ == "__main__":
  model = Darknet(fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg').read_bytes())
  print("Loading weights file (237MB). This might take a while…")
-  model.load_weights('https://pjreddie.com/media/files/yolov3.weights')
+  model.load_weights('https://github.com/shadiakiki1986/yolov3.weights/releases/download/3.0.1/yolov3.weights')
  if len(sys.argv) > 1:
    url = sys.argv[1]
  else:
@@ -1,6 +1,16 @@
-import re, ctypes, sys
+import re, ctypes, sys, importlib

-from tinygrad.runtime.autogen.am import am, mp_11_0, mp_13_0_0, nbio_4_3_0, mmhub_3_0_0, gc_11_0_0, osssys_6_0_0
+from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
+class AMDFake(AMDev):
+  def __init__(self, devfmt, vram, doorbell, mmio, dma_regions=None):
+    self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
+    self._run_discovery()
+    self._build_regs()
+
+amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
+amdev.AMDev = AMDFake
+
+from tinygrad.runtime.ops_amd import PCIIface

 def parse_amdgpu_logs(log_content, register_names=None):
  register_map = register_names
@@ -23,16 +33,11 @@ def parse_amdgpu_logs(log_content, register_names=None):
  return processed_log

 def main():
-  regs_offset = {13: {0: [3072, 37784576]}, 28: {0: [93184, 37754880], 1: [201327616, 201461760], 2: [209716224, 209850368], 3: [218104832, 218238976], 4: [226493440, 226627584], 5: [234882048, 235016192], 6: [243270656, 243404800]}, 21: {0: [28672, 12582912, 37795840, 130023424, 306184192], 1: [201326592, 201463808, 201465856, 204210176, 204472320], 2: [209715200, 209852416, 209854464, 212598784, 212860928], 3: [218103808, 218241024, 218243072, 220987392, 221249536], 4: [226492416, 226629632, 226631680, 229376000, 229638144], 5: [234881024, 235018240, 235020288, 237764608, 238026752], 6: [243269632, 243406848, 243408896, 246153216, 246415360]}, 22: {0: [18, 192, 13504, 36864, 37764096]}, 1: {0: [4704, 40960, 114688, 37760000]}, 2: {0: [3872, 37790720]}, 11: {0: [70656, 38103040]}, 12: {0: [106496, 37783552]}, 15: {0: [90112, 14417920, 14680064, 14942208, 38009856]}, 16: {0: [90112, 14417920, 14680064, 14942208, 38009856]}, 14: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 26: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 23: {0: [4256, 37789696]}, 33: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 25: {0: []}, 3: {0: [4704, 40960, 114688, 37760000]}, 4: {0: [4704, 40960, 114688, 37760000]}, 24: {0: [92160, 92672, 37752832, 54788096]}, 27: {0: [91648, 37751808], 1: [201339904, 201458176], 2: [209728512, 209846784], 3: [218117120, 218235392], 4: [226505728, 226624000], 5: [234894336, 235012608], 6: [243282944, 243401216]}, 29: {0: [201342976, 201344000, 205520896, 205537280], 1: [209731584, 209732608, 213909504, 213925888], 2: [218120192, 218121216, 222298112, 222314496], 3: [226508800, 226509824, 230686720, 230703104], 4: [234897408, 234898432, 239075328, 239091712], 5: [243286016, 243287040, 247463936, 247480320]}, 17: {0: [30720, 32256], 1: [31488, 73728]}}
-
  reg_names = {}
-  def _prepare_registers(modules):
-    for base, m in modules:
-      for k, regval in m.__dict__.items():
-        if k.startswith("reg") and not k.endswith("_BASE_IDX") and (base_idx:=getattr(m, f"{k}_BASE_IDX", None)) is not None:
-          reg_names[regs_offset[am.__dict__.get(f"{base}_HWIP")][0][base_idx] + regval] = k
-
-  _prepare_registers([("MP0", mp_13_0_0), ("NBIO", nbio_4_3_0), ("MMHUB", mmhub_3_0_0), ("GC", gc_11_0_0), ("OSSSYS", osssys_6_0_0)])
+  dev = PCIIface(None, 0)
+  for x, y in dev.dev_impl.__dict__.items():
+    if isinstance(y, AMRegister):
+      for inst, addr in y.addr.keys(): reg_names[addr] = f"{x}, xcc={inst}"

  with open(sys.argv[1], 'r') as f:
    log_content = log_content_them = f.read()
@@ -1,5 +1,5 @@
 from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
-from tinygrad.opt.kernel import Ops, MemOp, UOp
+from tinygrad.codegen.opt.kernel import Ops, MemOp, UOp
 from tinygrad.uop.ops import BinaryOps, UnaryOps
 from tinygrad.dtype import DType, dtypes
 from tinygrad.helpers import DEBUG
@@ -3,7 +3,7 @@ from platform import system
 from typing import Tuple, Dict, List, Optional
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.opt.kernel import Ops, UOp
+from tinygrad.codegen.opt.kernel import Ops, UOp
 from tinygrad.helpers import CI
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage

@@ -1,7 +1,7 @@
 from typing import List
 import struct
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
-from tinygrad.opt.kernel import Ops, UOp
+from tinygrad.codegen.opt.kernel import Ops, UOp
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.runtime.ops_cuda import arch
@@ -2,7 +2,7 @@ import yaml
 from typing import Tuple, Set, Dict
 from tinygrad import dtypes
 from tinygrad.codegen.assembly import AssemblyCodegen, Register
-from tinygrad.opt.kernel import Ops
+from tinygrad.codegen.opt.kernel import Ops
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH

@@ -2,7 +2,7 @@ from typing import Dict, List, Final, Callable, DefaultDict
 from collections import defaultdict
 from tinygrad.uop.ops import UnaryOps, BinaryOps, TernaryOps, Op
 from tinygrad.helpers import DType, PtrDType, dtypes, ImageDType, DEBUG, getenv
-from tinygrad.opt.kernel import  UOp, Ops
+from tinygrad.codegen.opt.kernel import  UOp, Ops
 from triton.compiler import compile as triton_compile
 import linecache
 import math
@@ -2,11 +2,11 @@ from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, graph_rewrite, AxisType, PatternMatcher, UPat
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from tinygrad.dtype import AddrSpace
-from tinygrad.schedule.kernelize import merge_views, view_left
 from tinygrad.helpers import getenv, colored, prod, unwrap
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.shape.view import strides_for_shape
-from tinygrad.opt.kernel import axis_colors
+from tinygrad.codegen.opt.kernel import axis_colors, Opt, OptOps
+from tinygrad.codegen.opt.swizzler import merge_views, view_left

 def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])

@@ -44,6 +44,21 @@ pm = PatternMatcher([
  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
 ])

+def rangeify_kernel3():
+  a = Tensor.empty(N,N)
+  b = Tensor.empty(N,N)
+  c = a@b
+  #c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
+  with Context(RANGEIFY=1):
+    sink = c.schedule()[-1].ast
+  #print(sink)
+
+  opts  = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
+  opts += [Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 1, 16), Opt(OptOps.UPCAST, 1, 2)]
+  opts += [Opt(OptOps.UNROLL, 0, 8)]
+
+  return sink.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
+
 def top_spec_kernel3():
  a = Tensor.empty(N,N)
  b = Tensor.empty(N,N)
@@ -309,10 +324,15 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):

 if __name__ == "__main__":
  HL = getenv("HL")
-  if HL == 2: hprg = top_spec_kernel3()
+  if HL == 3: hprg = rangeify_kernel3()
+  elif HL == 2: hprg = top_spec_kernel3()
  elif HL == 1: hprg = hl_spec_kernel3()
  else: hprg = hand_spec_kernel3()
-  prg = get_program(hprg, Device.default.renderer)
+  if HL == 3:
+    with Context(RANGEIFY=1, BLOCK_REORDER=0):
+      prg = get_program(hprg, Device.default.renderer)
+  else:
+    prg = get_program(hprg, Device.default.renderer)
  print(prg.src)
  if getenv("SRC"): exit(0)
  hrunner = CompiledRunner(prg)
@@ -5,9 +5,9 @@ from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
 from tinygrad.engine.realize import get_program

 # for copied uops
-from tinygrad.opt.kernel import Kernel, KernelOptError
+from tinygrad.codegen.opt.kernel import Kernel, KernelOptError
 from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
-from tinygrad.opt.search import Opt, OptOps
+from tinygrad.codegen.opt.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.dtype import PtrDType, DType, DTYPES_DICT
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -2,7 +2,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
-from tinygrad.opt.kernel import OptOps
+from tinygrad.codegen.opt.kernel import OptOps
 from tinygrad.engine.realize import lower_schedule

 dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
@@ -1,6 +1,6 @@
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.helpers import getenv, DEBUG
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from dataclasses import replace

@@ -37,7 +37,7 @@ B = Tensor.rand(K, N, device="CPU")
 C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2)

 sched = C.schedule()
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.device import CompilerOptions
 lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
 lin.to_program()
@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+
+from tinygrad.runtime.support.system import System
+import argparse, glob, os, re, time, subprocess, sys
+
+def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
+  target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
+
+  devs = []
+  for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
+    dev_id = dev[8:-5]
+    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
+  return devs
+
+def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
+def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
+
+def cmd_remove_module(args):
+  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
+  to_unload = [m for m in modules if _is_module_loaded(m)]
+  if not to_unload: print("Kernel modules are not loaded")
+  else:
+    print("Removing kernel modules:", ", ".join(to_unload))
+    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
+    except subprocess.CalledProcessError as e:
+      print("Failed to unload all modules — they may be in use.", file=sys.stderr)
+      sys.exit(e.returncode)
+
+def cmd_insert_module(args):
+  cmd_remove_module(args)
+  cmd_reset_devices(args)
+
+  module = "nvidia" if args.backend == "nv" else "amdgpu"
+  if _is_module_loaded(module):
+    print(f"{module} kernel module already loaded")
+    return
+
+  print(f"Inserting kernel module: {module}")
+  if args.backend == "nv":
+    subprocess.run(["nvidia-smi"], check=True)
+  elif args.backend == "amd":
+    subprocess.run(["sudo", "modprobe", "amdgpu"], check=True)
+
+def cmd_reset_devices(args):
+  devs = scan_devs_based_on_lock({"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    print(f"Resetting device {dev}")
+    if args.backend != "amd": _do_reset_device(dev)
+    time.sleep(0.2)
+
+def cmd_show_pids(args):
+  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    try:
+      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+      print(f"{dev}: {pid}")
+    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
+
+def cmd_kill_pids(args):
+  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    try:
+      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+      print(f"{dev}: {pid}")
+    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
+
+def cmd_kill_pids(args):
+  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    for i in range(128):
+      if i > 0: time.sleep(0.2)
+
+      try:
+        try: pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+        except subprocess.CalledProcessError: break
+
+        print(f"Killing process {pid} (which uses {dev})")
+        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
+      except subprocess.CalledProcessError as e:
+        print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)
+
+def add_common_commands(parent_subparsers):
+  p_insmod = parent_subparsers.add_parser("insmod", help="Insert a kernel module")
+  p_insmod.set_defaults(func=cmd_insert_module)
+
+  p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
+  p_rmmod.set_defaults(func=cmd_remove_module)
+
+  p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
+  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device to reset")
+  p_reset.set_defaults(func=cmd_reset_devices)
+
+  p_reset = parent_subparsers.add_parser("pids", help="Show pids of processes using the device")
+  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
+  p_reset.set_defaults(func=cmd_show_pids)
+
+  p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
+  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
+  p_reset.set_defaults(func=cmd_kill_pids)
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  backend_subparsers = parser.add_subparsers(dest="backend", required=True, metavar="{nv,amd}", help="Hardware backend to target")
+
+  nv_parser = backend_subparsers.add_parser("nv", help="NVIDIA GPUs")
+  nv_commands = nv_parser.add_subparsers(dest="command", required=True)
+  add_common_commands(nv_commands)
+
+  amd_parser = backend_subparsers.add_parser("amd", help="AMD GPUs")
+  amd_commands = amd_parser.add_subparsers(dest="command", required=True)
+  add_common_commands(amd_commands)
+
+  args = parser.parse_args()
+  if args.command is None:
+    parser.print_help(sys.stderr)
+    sys.exit(1)
+
+  args.func(args)
@@ -8,7 +8,6 @@ bert_train_params = {
  "BS": 96,
  "EVAL_BS": 96,
  "FUSE_ARANGE": 1,
-  "FUSE_ARANGE_UINT": 0,
  "BASEDIR": "/raid/datasets/wiki",
 }

@@ -0,0 +1,61 @@
+# HuggingFace ONNX
+
+Tool for discovering, downloading, and validating ONNX models from HuggingFace.
+
+## Extra Dependencies
+
+```bash
+pip install huggingface_hub pyyaml requests onnx onnxruntime numpy
+```
+
+## Huggingface Manager (discovering and downloading)
+
+The `huggingface_manager.py` script discovers top ONNX models from HuggingFace, collects metadata, and optionally downloads them.
+
+```bash
+# Download top 50 models sorted by downloads
+python huggingface_manager.py --limit 50 --download
+
+# Just collect metadata (no download)
+python huggingface_manager.py --limit 100
+
+# Sort by likes instead of downloads
+python huggingface_manager.py --limit 20 --sort likes --download
+
+# Custom output file
+python huggingface_manager.py --limit 10 --output my_models.yaml
+```
+
+### Output Format
+
+The tool generates a YAML file with the following structure:
+
+```yaml
+repositories:
+  "model-name":
+    url: "https://huggingface.co/model-name"
+    download_path: "/path/to/models/..."  # when --download used
+    files:
+      - file: "model.onnx"
+        size: "90.91MB"
+total_size: "2.45GB"
+created_at: "2024-01-15T10:30:00Z"
+```
+
+## Run Models (validation)
+
+The `run_models.py` script validates ONNX models against ONNX Runtime for correctness.
+
+```bash
+# Validate models from a YAML configuration file
+python run_models.py --validate huggingface_repos.yaml
+
+# Debug specific repository (downloads and validates all ONNX models)
+python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2
+
+# Debug specific model file
+python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2/onnx/model.onnx
+
+# Debug with model truncation for debugging and validating intermediate results
+DEBUGONNX=1 python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2/onnx/model.onnx --truncate 10
+```
@@ -1,85 +0,0 @@
-import yaml, time, requests, argparse
-from pathlib import Path
-from huggingface_hub import list_models, HfApi
-from tinygrad.helpers import tqdm
-
-HUGGINGFACE_URL = "https://huggingface.co"
-SKIPPED_FILES = [
-  "fp16", "int8", "uint8", "quantized",      # numerical accuracy issues
-  "avx2", "arm64", "avx512", "avx512_vnni",  # numerical accuracy issues
-  "q4", "q4f16", "bnb4",                     # unimplemented quantization
-  "model_O4",                                # requires non cpu ort runner and MemcpyFromHost op
-  "merged",                                  # TODO implement attribute with graph type and Loop op
-]
-SKIPPED_REPO_PATHS = [
-  # Invalid model-index
-  "AdamCodd/vit-base-nsfw-detector",
-  # TODO: implement attribute with graph type and Loop op
-  "minishlab/potion-base-8M", "minishlab/M2V_base_output", "minishlab/potion-retrieval-32M",
-  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, GroupQueryAttention
-  "HuggingFaceTB/SmolLM2-360M-Instruct",
-  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, RotaryEmbedding, MultiHeadAttention
-  "HuggingFaceTB/SmolLM2-1.7B-Instruct",
-  # TODO: implmement RandomNormalLike
-  "stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", 'SimianLuo/LCM_Dreamshaper_v7',
-  # TODO: implement NonZero
-  "mangoapps/fb_zeroshot_mnli_onnx",
-  # TODO huge Concat in here with 1024 (1, 3, 32, 32) Tensors, and maybe a MOD bug with const folding
-  "briaai/RMBG-2.0",
-]
-
-def get_top_repos(n: int, sort: str) -> list[str]: # list["FacebookAI/xlm-roberta-large", ...]
-  print(f"** Getting top {n} models sorted by {sort} **")
-  repos = []
-  i = 0
-  for model in list_models(filter="onnx", sort=sort):
-    if model.id in SKIPPED_REPO_PATHS: continue
-    print(f"{i+1}/{n}: {model.id} ({getattr(model, sort)})")
-    repos.append(model.id)
-    i += 1
-    if i == n: break
-  return repos
-
-def get_metadata(repos:list[str]) -> dict:
-  api = HfApi()
-  repos_metadata = {"repositories": {}}
-  total_size = 0
-
-  # TODO: speed head requests up with async?
-  for repo in tqdm(repos, desc="Getting metadata"):
-    files_metadata = []
-    model_info = api.model_info(repo)
-
-    for file in model_info.siblings:
-      filename = file.rfilename
-      if not (filename.endswith('.onnx') or filename.endswith('.onnx_data')): continue
-      if any(skip_str in filename for skip_str in SKIPPED_FILES): continue
-      head = requests.head(f"{HUGGINGFACE_URL}/{repo}/resolve/main/{filename}", allow_redirects=True)
-      file_size = file.size or int(head.headers.get('Content-Length', 0))
-      files_metadata.append({"file": filename, "size": f"{file_size/1e6:.2f}MB"})
-      total_size += file_size
-
-    repos_metadata["repositories"][repo] = {
-      "url": f"{HUGGINGFACE_URL}/{repo}",
-      "download_path": None,
-      "files": files_metadata,
-    }
-  repos_metadata['total_size'] = f"{total_size/1e9:.2f}GB"
-  repos_metadata['created_at'] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
-  return repos_metadata
-
-if __name__ == "__main__":
-  sort = "downloads" # recent 30 days downloads
-  huggingface_onnx_dir = Path(__file__).parent
-
-  parser = argparse.ArgumentParser(description="Produces a YAML file with metadata of top huggingface onnx models")
-  parser.add_argument("--limit", type=int, required=True, help="Number of top repositories to process (e.g., 100)")
-  parser.add_argument("--output", type=str, default="huggingface_repos.yaml", help="Output YAML file name to save the report")
-  args = parser.parse_args()
-
-  top_repos = get_top_repos(args.limit, sort)
-  metadata = get_metadata(top_repos)
-  yaml_path = huggingface_onnx_dir / args.output
-  with open(yaml_path, 'w') as f:
-    yaml.dump(metadata, f, sort_keys=False)
-    print(f"YAML saved to: {str(yaml_path)}")
@@ -1,29 +0,0 @@
-import yaml, argparse
-from pathlib import Path
-from huggingface_hub import snapshot_download
-
-def download_models(yaml_file: str, download_dir: str) -> None:
-  with open(yaml_file, 'r') as f: metadata = yaml.safe_load(f)
-  n = len(metadata["repositories"])
-
-  for i, (model_id, model_data) in enumerate(metadata["repositories"].items()):
-    print(f"Downloading {i+1}/{n}: {model_id}...")
-    allow_patterns = [file_info["file"] for file_info in model_data["files"]]
-    root_path = Path(snapshot_download(repo_id=model_id, allow_patterns=allow_patterns, cache_dir=download_dir))
-    # download configs too (the sizes are small)
-    snapshot_download(repo_id=model_id, allow_patterns=["*config.json"], cache_dir=download_dir)
-    print(f"Downloaded model files to: {root_path}")
-    model_data["download_path"] = str(root_path)
-
-  # Save the updated metadata back to the YAML file
-  with open(yaml_file, 'w') as f: yaml.dump(metadata, f, sort_keys=False)
-  print("Download completed according to YAML file.")
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Download models from Huggingface Hub based on a YAML configuration file.")
-  parser.add_argument("input", type=str, help="Path to the input YAML configuration file containing model information.")
-  args = parser.parse_args()
-
-  models_folder = Path(__file__).parent / "models"
-  models_folder.mkdir(parents=True, exist_ok=True)
-  download_models(args.input, str(models_folder))
@@ -0,0 +1,230 @@
+import yaml
+import time
+import requests
+import argparse
+from pathlib import Path
+from huggingface_hub import list_models, HfApi, snapshot_download
+from tinygrad.helpers import _ensure_downloads_dir
+DOWNLOADS_DIR = _ensure_downloads_dir() / "models"
+from tinygrad.helpers import tqdm
+
+def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, cache_dir: str|Path|None=None,
+                                 tries: int=2, **kwargs) -> Path:
+  for attempt in range(tries):
+    try:
+      return Path(snapshot_download(
+        repo_id=repo_id,
+        allow_patterns=allow_patterns,
+        cache_dir=str(cache_dir) if cache_dir is not None else None,
+        **kwargs
+      ))
+    except Exception as e:
+      if attempt == tries-1: raise
+      time.sleep(1)
+
+# Constants for filtering models
+HUGGINGFACE_URL = "https://huggingface.co"
+SKIPPED_FILES = [
+  "fp16", "int8", "uint8", "quantized",      # numerical accuracy issues
+  "avx2", "arm64", "avx512", "avx512_vnni",  # numerical accuracy issues
+  "q4", "q4f16", "bnb4",                     # unimplemented quantization
+  "model_O4",                                # requires non cpu ort runner and MemcpyFromHost op
+  "merged",                                  # TODO implement attribute with graph type and Loop op
+]
+
+SKIPPED_REPO_PATHS = [
+  # Invalid model-index
+  "AdamCodd/vit-base-nsfw-detector",
+  # TODO: implement attribute with graph type and Loop op
+  "minishlab/potion-base-8M", "minishlab/M2V_base_output", "minishlab/potion-retrieval-32M",
+  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, GroupQueryAttention
+  "HuggingFaceTB/SmolLM2-360M-Instruct",
+  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, RotaryEmbedding, MultiHeadAttention
+  "HuggingFaceTB/SmolLM2-1.7B-Instruct",
+  # TODO: implement RandomNormalLike
+  "stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", 'SimianLuo/LCM_Dreamshaper_v7',
+  # TODO: implement NonZero
+  "mangoapps/fb_zeroshot_mnli_onnx",
+  # TODO huge Concat in here with 1024 (1, 3, 32, 32) Tensors, and maybe a MOD bug with const folding
+  "briaai/RMBG-2.0",
+]
+
+
+class HuggingFaceONNXManager:
+  def __init__(self):
+    self.base_dir = Path(__file__).parent
+    self.models_dir = DOWNLOADS_DIR
+    self.api = HfApi()
+
+  def discover_models(self, limit: int, sort: str = "downloads") -> list[str]:
+    print(f"Discovering top {limit} ONNX models sorted by {sort}...")
+    repos = []
+    i = 0
+
+    for model in list_models(filter="onnx", sort=sort):
+      if model.id in SKIPPED_REPO_PATHS:
+        continue
+
+      print(f"  {i+1}/{limit}: {model.id} ({getattr(model, sort)})")
+      repos.append(model.id)
+      i += 1
+      if i == limit:
+        break
+
+    print(f"Found {len(repos)} suitable ONNX models")
+    return repos
+
+  def collect_metadata(self, repos: list[str]) -> dict:
+    print(f"Collecting metadata for {len(repos)} repositories...")
+    metadata = {"repositories": {}}
+    total_size = 0
+
+    for repo in tqdm(repos, desc="Collecting metadata"):
+      try:
+        files_metadata = []
+        model_info = self.api.model_info(repo)
+
+        for file in model_info.siblings:
+          filename = file.rfilename
+          if not (filename.endswith('.onnx') or filename.endswith('.onnx_data')):
+            continue
+          if any(skip_str in filename for skip_str in SKIPPED_FILES):
+            continue
+
+          # Get file size from API or HEAD request
+          try:
+            head = requests.head(
+              f"{HUGGINGFACE_URL}/{repo}/resolve/main/{filename}",
+              allow_redirects=True,
+              timeout=10
+            )
+            file_size = file.size or int(head.headers.get('Content-Length', 0))
+          except requests.RequestException:
+            file_size = file.size or 0
+
+          files_metadata.append({
+            "file": filename,
+            "size": f"{file_size/1e6:.2f}MB"
+          })
+          total_size += file_size
+
+        if files_metadata:  # Only add repos with valid ONNX files
+          metadata["repositories"][repo] = {
+            "url": f"{HUGGINGFACE_URL}/{repo}",
+            "download_path": None,
+            "files": files_metadata,
+          }
+
+      except Exception as e:
+        print(f"WARNING: Failed to collect metadata for {repo}: {e}")
+        continue
+
+    metadata['total_size'] = f"{total_size/1e9:.2f}GB"
+    metadata['created_at'] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+
+    print(f"Collected metadata for {len(metadata['repositories'])} repositories")
+    print(f"Total estimated download size: {metadata['total_size']}")
+
+    return metadata
+
+  def download_models(self, metadata: dict) -> dict:
+    self.models_dir.mkdir(parents=True, exist_ok=True)
+
+    repos = metadata["repositories"]
+    n = len(repos)
+
+    print(f"Downloading {n} repositories to {self.models_dir}...")
+
+    for i, (model_id, model_data) in enumerate(repos.items()):
+      print(f"  Downloading {i+1}/{n}: {model_id}...")
+
+      try:
+        # Download ONNX model files
+        allow_patterns = [file_info["file"] for file_info in model_data["files"]]
+        root_path = snapshot_download_with_retry(
+          repo_id=model_id,
+          allow_patterns=allow_patterns,
+          cache_dir=str(self.models_dir)
+        )
+
+        # Download config files (usually small)
+        snapshot_download_with_retry(
+          repo_id=model_id,
+          allow_patterns=["*config.json"],
+          cache_dir=str(self.models_dir)
+        )
+
+        model_data["download_path"] = str(root_path)
+        print(f"    Downloaded to: {root_path}")
+
+      except Exception as e:
+        print(f"    ERROR: Failed to download {model_id}: {e}")
+        model_data["download_path"] = None
+        continue
+
+    successful_downloads = sum(1 for repo in repos.values() if repo["download_path"] is not None)
+    print(f"Successfully downloaded {successful_downloads}/{n} repositories")
+    print(f"All models saved to: {self.models_dir}")
+
+    return metadata
+
+  def save_metadata(self, metadata: dict, output_file: str):
+    yaml_path = self.base_dir / output_file
+    with open(yaml_path, 'w') as f:
+      yaml.dump(metadata, f, sort_keys=False)
+    print(f"Metadata saved to: {yaml_path}")
+
+  def discover_and_download(self, limit: int, output_file: str = "huggingface_repos.yaml",
+                          sort: str = "downloads", download: bool = True):
+    print(f"Starting HuggingFace ONNX workflow...")
+    print(f"   Limit: {limit} models")
+    print(f"   Sort by: {sort}")
+    print(f"   Download: {'Yes' if download else 'No'}")
+    print(f"   Output: {output_file}")
+    print("-" * 50)
+
+    repos = self.discover_models(limit, sort)
+
+    metadata = self.collect_metadata(repos)
+
+    if download:
+      metadata = self.download_models(metadata)
+
+    self.save_metadata(metadata, output_file)
+
+    print("-" * 50)
+    print("Workflow completed successfully!")
+    if download:
+      successful = sum(1 for repo in metadata["repositories"].values()
+                     if repo["download_path"] is not None)
+      print(f"{successful}/{len(metadata['repositories'])} models downloaded")
+
+    return metadata
+
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(
+    description="HuggingFace ONNX Model Manager - Discover, collect metadata, and download ONNX models",
+  )
+
+  parser.add_argument("--limit", type=int, help="Number of top repositories to process")
+  parser.add_argument("--output", type=str, default="huggingface_repos.yaml",
+                     help="Output YAML file name (default: huggingface_repos.yaml)")
+  parser.add_argument("--sort", type=str, default="downloads",
+                     choices=["downloads", "likes", "created", "modified"],
+                     help="Sort criteria for model discovery (default: downloads)")
+
+  parser.add_argument("--download", action="store_true", default=False,
+                     help="Download models after collecting metadata")
+
+  args = parser.parse_args()
+
+  if not args.limit: parser.error("--limit is required")
+
+  manager = HuggingFaceONNXManager()
+  manager.discover_and_download(
+    limit=args.limit,
+    output_file=args.output,
+    sort=args.sort,
+    download=args.download
+  )
@@ -1,10 +1,11 @@
-import onnx, yaml, tempfile, time, collections, pprint, argparse, json
+import onnx, yaml, tempfile, time, argparse, json
 from pathlib import Path
+from typing import Any
 from tinygrad.frontend.onnx import OnnxRunner
-from extra.onnx import get_onnx_ops
 from extra.onnx_helpers import validate, get_example_inputs
+from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry

-def get_config(root_path: Path):
+def get_config(root_path: Path) -> dict[str, Any]:
  ret = {}
  for path in root_path.rglob("*config.json"):
    config = json.load(path.open())
@@ -12,19 +13,19 @@ def get_config(root_path: Path):
      ret.update(config)
  return ret

-def run_huggingface_validate(onnx_model_path, config, rtol, atol):
-  onnx_runner = OnnxRunner(onnx_model_path)
-  inputs = get_example_inputs(onnx_runner.graph_inputs, config)
-  validate(onnx_model_path, inputs, rtol=rtol, atol=atol)
-
-def get_tolerances(file_name): # -> rtol, atol
+def get_tolerances(file_name: str) -> tuple[float, float]:
  # TODO very high rtol atol
  if "fp16" in file_name: return 9e-2, 9e-2
  if any(q in file_name for q in ["int8", "uint8", "quantized"]): return 4, 4
  return 4e-3, 3e-2

+def run_huggingface_validate(onnx_model_path: str | Path, config: dict[str, Any], rtol: float, atol: float):
+  onnx_runner = OnnxRunner(onnx_model_path)
+  inputs = get_example_inputs(onnx_runner.graph_inputs, config)
+  validate(onnx_model_path, inputs, rtol=rtol, atol=atol)
+
 def validate_repos(models:dict[str, tuple[Path, Path]]):
-  print(f"** Validating {len(model_paths)} models **")
+  print(f"** Validating {len(models)} models **")
  for model_id, (root_path, relative_path) in models.items():
    print(f"validating model {model_id}")
    model_path = root_path / relative_path
@@ -36,25 +37,6 @@ def validate_repos(models:dict[str, tuple[Path, Path]]):
    et = time.time() - st
    print(f"passed, took {et:.2f}s")

-def retrieve_op_stats(models:dict[str, tuple[Path, Path]]) -> dict:
-  ret = {}
-  op_counter = collections.Counter()
-  unsupported_ops = collections.defaultdict(set)
-  supported_ops = get_onnx_ops()
-  print(f"** Retrieving stats from {len(model_paths)} models **")
-  for model_id, (root_path, relative_path) in models.items():
-    print(f"examining {model_id}")
-    model_path = root_path / relative_path
-    onnx_runner = OnnxRunner(model_path)
-    for node in onnx_runner.graph_nodes:
-      op_counter[node.op] += 1
-      if node.op not in supported_ops:
-        unsupported_ops[node.op].add(model_id)
-    del onnx_runner
-  ret["unsupported_ops"] = {k:list(v) for k, v in unsupported_ops.items()}
-  ret["op_counter"] = op_counter.most_common()
-  return ret
-
 def debug_run(model_path, truncate, config, rtol, atol):
  if truncate != -1:
    model = onnx.load(model_path)
@@ -71,12 +53,9 @@ def debug_run(model_path, truncate, config, rtol, atol):
    run_huggingface_validate(model_path, config, rtol, atol)

 if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Huggingface ONNX Model Validator and Ops Checker")
-  parser.add_argument("input", type=str, help="Path to the input YAML configuration file containing model information.")
-  parser.add_argument("--check_ops", action="store_true", default=False,
-                      help="Check support for ONNX operations in models from the YAML file")
-  parser.add_argument("--validate", action="store_true", default=False,
-                      help="Validate correctness of models from the YAML file")
+  parser = argparse.ArgumentParser(description="Huggingface ONNX Model Validator")
+  parser.add_argument("--validate", type=str, default="",
+                      help="Validate correctness of models from the specified YAML configuration file")
  parser.add_argument("--debug", type=str, default="",
                      help="""Validates without explicitly needing a YAML or models pre-installed.
                      provide repo id (e.g. "minishlab/potion-base-8M") to validate all onnx models inside the repo
@@ -85,13 +64,13 @@ if __name__ == "__main__":
  parser.add_argument("--truncate", type=int, default=-1, help="Truncate the ONNX model so intermediate results can be validated")
  args = parser.parse_args()

-  if not (args.check_ops or args.validate or args.debug):
-    parser.error("Please provide either --validate, --check_ops, or --debug.")
+  if not (args.validate or args.debug):
+    parser.error("Please provide either --validate <yaml_file> or --debug <repo_id>.")
  if args.truncate != -1 and not args.debug:
    parser.error("--truncate and --debug should be used together for debugging")

-  if args.check_ops or args.validate:
-    with open(args.input, 'r') as f:
+  if args.validate:
+    with open(args.validate, 'r') as f:
      data = yaml.safe_load(f)
      assert all(repo["download_path"] is not None for repo in data["repositories"].values()), "please run `download_models.py` for this yaml"
      model_paths = {
@@ -101,22 +80,16 @@ if __name__ == "__main__":
        if model["file"].endswith(".onnx")
      }

-    if args.check_ops:
-      pprint.pprint(retrieve_op_stats(model_paths))
-
-    if args.validate:
-      validate_repos(model_paths)
+    validate_repos(model_paths)

  if args.debug:
-    from huggingface_hub import snapshot_download
-    download_dir = Path(__file__).parent / "models"
    path:list[str] = args.debug.split("/")
    if len(path) == 2:
      # repo id
      # validates all onnx models inside repo
      repo_id = "/".join(path)
-      root_path = Path(snapshot_download(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=download_dir))
-      snapshot_download(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=download_dir)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=DOWNLOADS_DIR)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
      config = get_config(root_path)
      for onnx_model in root_path.rglob("*.onnx"):
        rtol, atol = get_tolerances(onnx_model.name)
@@ -128,8 +101,8 @@ if __name__ == "__main__":
      onnx_model = path[-1]
      assert path[-1].endswith(".onnx")
      repo_id, relative_path = "/".join(path[:2]), "/".join(path[2:])
-      root_path = Path(snapshot_download(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=download_dir))
-      snapshot_download(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=download_dir)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=DOWNLOADS_DIR)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
      config = get_config(root_path)
      rtol, atol = get_tolerances(onnx_model)
      print(f"validating {relative_path} with truncate={args.truncate}, {rtol=}, {atol=}")
@@ -4,9 +4,9 @@ import numpy as np
 np.set_printoptions(suppress=True)
 import math, functools, time, random, statistics
 from tinygrad.helpers import DEBUG, getenv, CACHELEVEL, diskcache_get, diskcache_put, colored, Profiling
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.device import Buffer, Device, CompileError
-from tinygrad.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
+from tinygrad.codegen.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
 from tinygrad.engine.realize import get_program

 class MCTSNode:
@@ -99,7 +99,9 @@ class FeedForward:
    self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit

  def __call__(self, x:Tensor) -> Tensor:
-    return self.w2(self.w1(x).silu() * self.w3(x)) # SwiGLU [arxiv/2002.05202, eq (5)]
+    w1 = self.w1(x).silu()
+    w3 = self.w3(x.contiguous_backward())  # this fixes a strange fusion that makes tensor cores miss
+    return self.w2(w1 * w3)

 class TransformerBlock:
  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, max_context:int, linear=nn.Linear,
@@ -111,7 +113,7 @@ class TransformerBlock:

  def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]):
    h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
-    return (h + self.feed_forward(self.ffn_norm(h))).contiguous()
+    return (h + self.feed_forward(self.ffn_norm(h))).contiguous().contiguous_backward()

 # standard openai sampling
 def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):
@@ -179,16 +181,14 @@ class Transformer:
  def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):
    _bsz, seqlen = tokens.shape
    h = self.tok_embeddings(tokens)
-
-    self.freqs_cis = self.freqs_cis.cast(h.dtype).contiguous()
-    freqs_cis = self.freqs_cis[:, start_pos:start_pos+seqlen, :, :, :]
+    freqs_cis = self.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :]

    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None
    for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
-    logits = self.output(self.norm(h)).float()[:, -1, :]
+    logits = self.output(self.norm(h))
    if math.isnan(temperature): return logits

-    return sample(logits.flatten(), temperature, top_k, top_p, alpha_f, alpha_p)
+    return sample(logits[:, -1, :].flatten(), temperature, top_k, top_p, alpha_f, alpha_p)

  def __call__(self, tokens:Tensor, start_pos:int, temperature:float=0.0, top_k:int=0, top_p:float=0.8, alpha_f:float=0.0, alpha_p:float=0.0):
    # TODO: better way to handle the first call v.s. the rest?
@@ -1,2 +0,0 @@
-GPU="$1"
-echo 1 | sudo tee /sys/bus/pci/devices/$GPU/reset 2>/dev/null
@@ -1,65 +0,0 @@
-#!/usr/bin/env python3
-
-from tinygrad.runtime.support.system import System
-import argparse, glob, os, re, time, subprocess, sys
-
-def scan_devs_based_on_lock(prefix:str) -> list[str]:
-  devs = []
-  for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
-    dev_id = dev[8:-5]
-    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}"): devs.append(dev_id)
-  return devs
-
-def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
-def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
-
-def cmd_remove_module(args):
-  to_unload = [m for m in ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if _is_module_loaded(m)]
-  if not to_unload:
-    print("NVIDIA kernel modules are not loaded")
-  else:
-    print("Removing NVIDIA kernel modules:", ", ".join(to_unload))
-    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
-    except subprocess.CalledProcessError as e:
-      print("Failed to unload all modules — they may be in use.", file=sys.stderr)
-      sys.exit(e.returncode)
-
-def cmd_insert_module(args):
-  cmd_remove_module(args)
-  cmd_reset_devices(args)
-
-  if not os.path.exists("/sys/module/nvidia"):
-    print("Inserting nvidia kernel module")
-    subprocess.run(["nvidia-smi"], check=True)
-  else: print("Nvidia kernel module already loaded")
-
-def cmd_reset_devices(args):
-  devs = scan_devs_based_on_lock("nv")
-  dev_to_reset = args.pci_bus if 'pci_bus' in args.__dir__() else ""
-
-  for dev in devs:
-    if dev.startswith(dev_to_reset):
-      print(f"Resetting device {dev}")
-      _do_reset_device(dev)
-      time.sleep(0.2)
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  subparsers = parser.add_subparsers(required=True, dest="cmd")
-
-  parser_insmod = subparsers.add_parser('insmod', help='Insert a nvidia kernel module')
-  parser_insmod.set_defaults(func=cmd_insert_module)
-
-  parser_rmmod = subparsers.add_parser('rmmod', help='Remove a nvidia kernel module')
-  parser_rmmod.set_defaults(func=cmd_remove_module)
-
-  parser_reset = subparsers.add_parser('reset', help='Reset a nvidia device')
-  parser_reset.add_argument('--pci_bus', type=str, default="", help='PCI bus ID of the device to reset')
-  parser_reset.set_defaults(func=cmd_reset_devices)
-
-  args = parser.parse_args()
-  if args.cmd is None:
-    parser.print_help(sys.stderr)
-    sys.exit(1)
-
-  args.func(args)
@@ -1,954 +0,0 @@
-# mypy: disable-error-code="misc, list-item, assignment, operator, index, arg-type"
-from types import SimpleNamespace
-from typing import Any, Sequence, cast, Literal, Callable, get_args, NamedTuple
-import dataclasses, functools, io, math, types, warnings, pathlib, sys, enum
-from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr
-from tinygrad.helpers import getenv, DEBUG, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element
-from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype
-from tinygrad.device import is_dtype_supported, Device
-from extra.onnx_parser import onnx_load
-
-# https://github.com/onnx/onnx/blob/rel-1.17.0/onnx/onnx.proto3#L500-L544
-data_types: dict[int, DType] = {
-  1:dtypes.float32, 2:dtypes.uint8, 3:dtypes.int8, 4:dtypes.uint16, 5:dtypes.int16, 6:dtypes.int32, 7:dtypes.int64,
-  9:dtypes.bool, 10:dtypes.float16, 11:dtypes.double, 12:dtypes.uint32, 13:dtypes.uint64, 16:dtypes.bfloat16,
-}
-
-# https://github.com/onnx/onnx/blob/rel-1.17.0/onnx/onnx.proto3#L128-L145
-attribute_types: dict[int, Callable] = {
-  1: lambda a: float(a.f),
-  2: lambda a: int(a.i),
-  3: lambda a: a.s.data().tobytes().decode("utf8") if isinstance(a.s, Tensor) else a.s.decode("utf8"),
-  4: lambda a: buffer_parse(a.t),
-  6: lambda a: tuple(float(x) for x in a.floats),
-  7: lambda a: tuple(int(x) for x in a.ints),
-  8: lambda a: tuple(x.data().tobytes().decode("utf8") for x in a.strings)
-}
-
-# ***** protobuf parsing ******
-from onnx import AttributeProto, TensorProto, TypeProto
-
-def has_field(onnx_type: TypeProto|SimpleNamespace, field):
-  if isinstance(onnx_type, TypeProto): return onnx_type.HasField(field)
-  return hasattr(onnx_type, field)
-
-def dtype_parse(onnx_dtype: int, fallback_context: str | None = None) -> DType:
-  if onnx_dtype not in data_types: raise NotImplementedError(f"onnx dtype id {onnx_dtype} is not supported")
-  if is_dtype_supported(dtype := data_types[onnx_dtype]): return dtype
-  # if fallback_context is provided, we can fall back to a default dtype
-  if fallback_context is not None:
-    default_dtype = dtypes.default_int if dtypes.is_int(dtype) else dtypes.default_float
-    warnings.warn(f"dtype {dtype} on {Device.DEFAULT} from {fallback_context} is not supported, falling back to {default_dtype}")
-    assert is_dtype_supported(default_dtype), f"dtype {default_dtype} must be supported on {Device.DEFAULT}"
-    return default_dtype
-  raise RuntimeError(f"dtype {dtype} on device {Device.DEFAULT} is not supported")
-
-def attribute_parse(onnx_attribute: AttributeProto):
-  if onnx_attribute.type not in attribute_types: raise NotImplementedError(f"attribute type {onnx_attribute.type} is not supported")
-  return attribute_types[onnx_attribute.type](onnx_attribute)
-
-def buffer_parse(onnx_tensor: TensorProto) -> Tensor:
-  if onnx_tensor.string_data: raise NotImplementedError("Parsing for buffer with string data is not implemented.")
-  to_dtype, true_dtype = dtype_parse(onnx_tensor.data_type, "buffer parse"), data_types[onnx_tensor.data_type]
-  shape = tuple(onnx_tensor.dims)
-  keys = ['float_data', 'int32_data', 'int64_data', 'double_data', 'uint64_data', "raw_data"]
-  data = next((val for k in keys if (val := getattr(onnx_tensor, k)) is not None), None)
-  if data is None: raise RuntimeError("empty buffer")
-  if not isinstance(data, Tensor): return Tensor(data, dtype=to_dtype).reshape(shape)
-  assert data.dtype is dtypes.uint8, data.dtype
-  data = data.bitcast(true_dtype).reshape(shape)
-  data = data.to(Device.DEFAULT) if true_dtype is to_dtype else data.to("cpu").cast(to_dtype).to(Device.DEFAULT)
-  if shape == ():
-    if data.dtype is dtypes.float16 and sys.version_info < (3, 12): data = data.cast(dtypes.float32)
-    return Tensor(data.item(), dtype=to_dtype).reshape(shape)
-  return data
-
-def type_parse(onnx_type: TypeProto):
-  elem_type = onnx_type
-  if has_field(elem_type, "map_type") or has_field(elem_type, "sparse_tensor_type") or has_field(elem_type, "opaque_type"):
-    raise NotImplementedError("parsing for map_type, sparse_tensor_type and opaque_type are not implemented")
-  if is_optional := has_field(elem_type, "optional_type"): elem_type = elem_type.optional_type.elem_type
-  if is_sequence := has_field(elem_type, "sequence_type"): elem_type = elem_type.sequence_type.elem_type
-  if has_field(elem_type, "tensor_type"):
-    shape = tuple(getattr(d, "dim_param", None) or getattr(d, "dim_value") for d in elem_type.tensor_type.shape.dim) \
-      if has_field(elem_type.tensor_type, "shape") else None # test_identity_sequence_cpu
-    dtype = data_types[elem_type.tensor_type.elem_type]
-    return OnnxValue(shape, dtype, is_optional, is_sequence)
-  raise RuntimeError(f"TypeProto was not parsed properly: {onnx_type=}")
-
-# ***** onnx spec *****
-@dataclasses.dataclass(frozen=True)
-class OnnxValue:
-  shape: tuple[str|int, ...]
-  dtype: DType
-  is_optional: bool
-  is_sequence: bool
-
-class Domain(enum.Enum):
-  ONNX = "ai.onnx"
-  ONNX_ML = "ai.onnx.ml"
-  AI_ONNX_TRAINING = "ai.onnx.training"
-  AI_ONNX_PREVIEW_TRAINING = "ai.onnx.preview.training"
-  MICROSOFT_CONTRIB_OPS = "com.microsoft"
-  @classmethod
-  def from_onnx(cls, domain: str | None) -> "Domain": return cls.ONNX if domain is None or domain == "" else cls(domain)
-
-class OpSetId(NamedTuple):
-  domain: Domain
-  version: int
-
-@dataclasses.dataclass(frozen=True)
-class OnnxNode:
-  num: int
-  op: str
-  opset_id: OpSetId
-  inputs: tuple[str, ...]
-  outputs: tuple[str, ...]
-  opts: dict[str, Any]
-
-# ***** python const *****
-required_input_python_consts: dict[str, tuple[int, ...]] = {
-  "Tile": (1,), "Range": (0,1,2), "Expand": (1,), "Reshape": (1,), "Squeeze": (1,), "Unsqueeze": (1,), "Trilu": (1,), "ConstantOfShape": (0,),
-  "CumSum": (1,), "TopK": (1,), "Pad": (1,2,3), "MaxUnpool": (2,), "Dropout": (1,2), "CenterCropPad": (1,), "OneHot": (1,), "Compress": (1,),
-  "ImageDecoder": (0,), "AffineGrid": (1,), "Resize": (1,2,3), "Upsample": (1,), "Split": (1,), "Slice": (1,2,3,4),
-  **{"Reduce"+r: (1,) for r in ("Max", "Min", "Sum", "Mean", "SumSquare", "Prod", "L1", "L2", "LogSum", "LogSumExp")},
-  **{optim: (1,) for optim in ("Adam", "Adagrad", "Momentum")}
-}
-
-cache_misses = 0
-@functools.cache
-def _cached_to_python_const(t:Tensor):
-  if t.dtype is dtypes.uint8: return t.data().tobytes()
-  if 0 in t.shape: return []
-  return t.tolist()
-
-# Tensor -> python value cache for parameters
-def to_python_const(t:Any, op:str, idx:int) -> list[ConstType]|ConstType|bytes:
-  if idx not in required_input_python_consts.get(op, ()) or not isinstance(t, Tensor): return t
-  global cache_misses
-  ret = _cached_to_python_const(t)
-  if (info := _cached_to_python_const.cache_info()).misses > cache_misses and DEBUG >= 3:
-    print(f"Cache miss for {t}")
-    cache_misses = info.misses
-  return ret
-
-# ***** runner ******
-debug = int(getenv("DEBUGONNX", "0"))
-limit = int(getenv("ONNXLIMIT", "-1"))
-class OnnxRunner:
-  """
-  `OnnxRunner` executes an ONNX model using Tinygrad.
-
-  Args:
-    model_path: The ONNX model, provided as a file path (a string or Path object) or a Tensor.
-  """
-  def __init__(self, model_path: Tensor | str | pathlib.Path):
-    model = onnx_load(model_path)
-    self.is_training = any(n.domain in {Domain.AI_ONNX_TRAINING, Domain.AI_ONNX_PREVIEW_TRAINING} for n in model.graph.node)
-    self.old_training = Tensor.training
-    Tensor.training = True if self.is_training else False
-    self.graph_values = {"": None, **{x.name:buffer_parse(x) for x in model.graph.initializer}}
-    self.graph_inputs = {x.name:type_parse(x.type) for x in model.graph.input if x.name not in self.graph_values}
-    self.graph_outputs = tuple(x.name for x in model.graph.output)
-    opset_imports = {Domain.from_onnx(getattr(x, "domain", "")):x.version for x in model.opset_import}
-    self.graph_nodes = []
-    for num, n in enumerate(model.graph.node):
-      domain = Domain.from_onnx(n.domain)
-      opset_id = OpSetId(domain, opset_imports.get(domain, 1))
-      self.graph_nodes.append(OnnxNode(num, n.op_type, opset_id, tuple(n.input), tuple(n.output), {x.name:attribute_parse(x) for x in n.attribute}))
-    self.graph_nodes = tuple(self.graph_nodes)
-    self.variable_dims: dict[str, int] = {}
-
-    self.onnx_ops = onnx_ops
-
-  def _parse_input(self, name: str, value: Any, spec: OnnxValue):
-    if spec.is_optional and value is None: return None
-    if spec.is_sequence:
-      if not isinstance(value, Sequence): raise RuntimeError(f"input {name} received {value}, expected a sequence type")
-      sequence = [Tensor(v, dtype=spec.dtype, requires_grad=self.is_training) if not isinstance(v, Tensor) else v for v in value]
-      if not all_same(tuple(t.shape for t in sequence)): raise RuntimeError(f"Shapes for input {name} sequence must be homogeneous")
-      if not all(t.dtype is spec.dtype for t in sequence): warnings.warn(f"Dtypes for input {name} sequence aren't all {spec.dtype}")
-      return sequence
-    dtype = _from_np_dtype(value.dtype) if is_numpy_ndarray(value) else spec.dtype
-    tensor = Tensor(value, dtype=dtype, requires_grad=self.is_training) if not isinstance(value, Tensor) else value
-    if tensor.dtype is not spec.dtype: warnings.warn(f"input {name} has mismatch on dtype. Expected {spec.dtype}, received {tensor.dtype}.")
-    for dim, (onnx_dim, user_dim_input) in enumerate(zip(spec.shape, tensor.shape, strict=True)):
-      if isinstance(onnx_dim, str):
-        onnx_dim = self.variable_dims[onnx_dim] if onnx_dim in self.variable_dims else self.variable_dims.setdefault(onnx_dim, int(user_dim_input))
-      if user_dim_input != onnx_dim: raise RuntimeError(f"input {name} has mismatch on {dim=}. Expected {onnx_dim}, received {user_dim_input}.")
-    return tensor
-
-  def _select_op(self, op:str, required_opset:OpSetId) -> types.FunctionType:
-    if op not in self.onnx_ops: raise NotImplementedError(f"{op=} is not supported")
-    # return default implementation if no opset_id is specified
-    if isinstance(impl := self.onnx_ops[op], types.FunctionType): return impl
-    # match domain and select implementation with latest compatible version
-    eligible_ops = {impl_opset.version:impl_fxn for impl_opset,impl_fxn in impl.items()
-                    if impl_opset.domain == required_opset.domain and impl_opset.version <= required_opset.version}
-    if not eligible_ops: raise NotImplementedError(f"{op=} is not supported for domain {required_opset.domain} and version {required_opset.version}")
-    return eligible_ops[max(eligible_ops.keys())]
-
-  def get_empty_input_data(self, device:str|None=None, dtype:DType|None=None) -> dict[str, Tensor]:
-    return {name:Tensor.empty(*spec.shape, device=device, dtype=dtype or spec.dtype) for name, spec in self.graph_inputs.items()}
-
-  def to(self, device:str|None):
-    self.graph_values = {k:v.to(device) if isinstance(v, Tensor) else v for k,v in self.graph_values.items()}
-    self.graph_nodes = tuple(OnnxNode(n.num, n.op, n.opset_id, tuple(n.inputs), tuple(n.outputs),
-                                      {k:v.to(device) if isinstance(v, Tensor) else v for k,v in n.opts.items()}) for n in self.graph_nodes)
-    return self
-
-  def __call__(self, inputs:dict[str, Any], debug=debug):
-    for name, input_spec in self.graph_inputs.items():
-      if name not in inputs: raise RuntimeError(f"Please provide input data for {name}")
-      self.graph_values[name] = self._parse_input(name, inputs[name], input_spec)
-
-    for node in self.graph_nodes:
-      inps = [to_python_const(self.graph_values[name], node.op, i) for i,name in enumerate(node.inputs)]
-      opts = node.opts
-
-      # provide additional opts
-      if node.op == "Split" and 'num_outputs' not in opts: opts['num_outputs'] = len(node.outputs)
-      if node.op == "Gradient": opts['intermediate_tensors'] = self.graph_values
-
-      if debug >= 1: print(f"{node.num}: op '{node.op}' opt {opts}")
-      if debug >= 2 and node.inputs: print("\tinputs:\n" + "\n".join(f"\t\t{x} - {i!r}" for x,i in zip(node.inputs, inps)))
-      ret = self._select_op(node.op, node.opset_id)(*inps, **opts)
-      ret = ret if isinstance(ret, tuple) else (ret,)
-      if debug >= 2: print("\toutputs:\n" + "\n".join(f"\t\t{x} - {o!r}" for x,o in zip(node.outputs, ret)))
-
-      self.graph_values.update(dict(zip(node.outputs, ret[:len(node.outputs)], strict=True)))
-
-      if node.num == limit:
-        Tensor.training = self.old_training
-        return {name:self.graph_values[name] for name in node.outputs}
-    Tensor.training = self.old_training
-    return {name:self.graph_values[name] for name in self.graph_outputs}
-
-####################
-##### ONNX OPS #####
-####################
-def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionType]]:
-  # ***** helper functions *****
-  def _resolve_const(x: Sequence[ConstType]|ConstType): return x if isinstance(x, get_args(ConstType)) else get_single_element(x)
-
-  def _axes(axes, noop_with_empty_axes): return axes or ([] if noop_with_empty_axes else None)
-
-  # (padding_top, padding_left, ..., padding_bottom, padding_right, ...) -> (padding_left, padding_right, padding_top, padding_bottom, ...)
-  def _onnx_pads_to_tiny_pads(pads): return tuple(flatten(reversed(list(zip(pads, pads[len(pads)//2:])))))
-
-  AUTO_PAD_OPTIONS = Literal["NOTSET", "SAME_UPPER", "SAME_LOWER", "VALID"]
-  # (padding_height, padding_width) -> (padding_top, padding_left, padding_bottom, padding_right)
-  def _auto_pad(pads, auto_pad: AUTO_PAD_OPTIONS):
-    if auto_pad == "SAME_UPPER": return [pads[i]//2 for i in range(len(pads))] + [pads[i]-pads[i]//2 for i in range(len(pads))]
-    return [pads[i]-pads[i]//2 for i in range(len(pads))] + [pads[i]//2 for i in range(len(pads))]
-
-  def _resolve_pool_pads(x:Tensor, p_, k_, d_, s_, auto_pad:AUTO_PAD_OPTIONS):
-    if auto_pad == "VALID": return [0]*(len(k_)*2)
-    i_, (s_,d_,p_) = x.shape[-len(k_):], (make_tuple(x, len(k_)*2) for x in (s_, d_, p_))
-    if auto_pad == "NOTSET": return _onnx_pads_to_tiny_pads(p_ if len(p_)==len(k_)*2 else p_*2)
-    o_ = [((i - (1 if auto_pad in ("SAME_UPPER", "SAME_LOWER") else k)) // s + 1) for i,k,s in zip(i_, k_, s_)]
-    return _onnx_pads_to_tiny_pads(_auto_pad([(o-1)*s+k-i for o,i,k,s in zip(o_, i_, k_, s_)], auto_pad))
-
-  def _clamp_cast(x:Tensor, dtype:DType): return x.clamp(dtypes.min(dtype), dtypes.max(dtype)).cast(dtype)
-
-  def _prepare_quantize(x:Tensor, scale:Tensor, zero_point:Tensor|int, axis=1, block_size=0):
-    if axis < 0: axis += x.ndim
-    # https://github.com/onnx/onnx/blob/main/onnx/reference/ops/op_quantize_linear.py#L31
-    def reshape(val:Tensor):
-      if val.numel() == 1: return val
-      if block_size == 0: return val.reshape([val.shape[0] if dim == axis else 1 for dim in range(x.ndim)])
-      return val.repeat_interleave(block_size, axis)
-    return (reshape(scale), reshape(zero_point) if isinstance(zero_point, Tensor) else zero_point)
-
-  def _op_integer(op, inputs:list[Tensor], zero_points:list[Tensor], **opts):
-    adjusted_inputs = [inp.int() - zp for inp, zp in zip(inputs, zero_points)]
-    return op(*adjusted_inputs, **opts)
-
-  def _qlinearop_quantized(op, inputs:list[Tensor], zero_points:list[Tensor], scales:list[Tensor], out_scale:Tensor, out_zero_point:Tensor, **opts):
-    # op execution is done in quantized int
-    out = _op_integer(op, inputs, zero_points, **opts)
-    assert dtypes.is_int(out.dtype), "quantized op should've done math in int"
-    out_quantized = (out * prod(scales) / out_scale).round() + out_zero_point
-    return _clamp_cast(out_quantized, out_zero_point.dtype)
-
-  def _qlinearop_float(op, inputs:list[Tensor], zero_points:list[Tensor], scales:list[Tensor], out_scale:Tensor, out_zero_point:Tensor, **opts):
-    # op execution is done in float32
-    dequantized_inputs = [(inp.int() - zp) * scale for inp, zp, scale in zip(inputs, zero_points, scales)]
-    out = op(*dequantized_inputs, **opts)
-    assert dtypes.is_float(out.dtype), "op should've done math in float"
-    out_quantized = (out / out_scale).round() + out_zero_point
-    return _clamp_cast(out_quantized, out_zero_point.dtype)
-
-  def _onnx_training(input_group_size):
-    def __decorator(func):
-      def ___wrapper(R:Tensor, T:int, *inputs:Tensor, **kwargs):
-        R = R.detach()
-        groups = len(inputs) // input_group_size
-        ret = [func(R, T, *inps, **kwargs) for inps in (inputs[i::groups] for i in range(groups))]
-        return tuple(flatten(zip(*ret)))
-      return ___wrapper
-    return __decorator
-
-  # ***** Property/Graph Ops *****
-  def Identity(x:Tensor): return x
-  def Constant(sparse_value:Tensor|None=None, value:Tensor|None=None, value_float:float|None=None, value_floats:list[float]|None=None,
-              value_int:int|None=None, value_ints:list[int]|None=None, value_string:str|None=None, value_strings:list[str]|None=None):
-    if value is not None: return value
-    if value_float is not None: return Tensor(value_float, dtype=dtypes.float32, requires_grad=False)
-    if value_floats is not None: return Tensor(list(value_floats), dtype=dtypes.float32, requires_grad=False)
-    if value_int is not None: return Tensor(value_int, dtype=dtypes.int64, requires_grad=False)
-    if value_ints is not None: return Tensor(list(value_ints), dtype=dtypes.int64, requires_grad=False)
-    if value_string is not None or value_strings is not None and sparse_value is not None:
-      raise NotImplementedError('Constant OP not implemented for value_string, value_strings and sparse_value')
-
-  def Range(start:float|int|list[float|int], limit:float|int|list[float|int], delta:float|int|list[float|int]):
-    return Tensor.arange(start=_resolve_const(start), stop=_resolve_const(limit), step=_resolve_const(delta))
-
-  def ImageDecoder(encoded_stream:bytes, pixel_format="RGB"):
-    try: import PIL.Image
-    except ImportError as e: raise ImportError("Pillow must be installed for the ImageDecoder operator") from e
-    img = PIL.Image.open(io.BytesIO(encoded_stream))
-    if pixel_format == "BGR": return Tensor(img.tobytes(), dtype=dtypes.uint8).reshape(*img.size, 3).flip(-1)
-    if pixel_format == "RGB": return Tensor(img.tobytes(), dtype=dtypes.uint8).reshape(*img.size, 3)
-    if pixel_format == "Grayscale": return Tensor(img.convert("L").tobytes(), dtype=dtypes.uint8).reshape(*img.size, 1)
-    raise ValueError(f"pixel_format={pixel_format!r} is not supported.")
-
-  def EyeLike(x:Tensor, dtype:int|None=None, k:int=0):
-    ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_parse(dtype, "EyeLike op") if dtype is not None else x.dtype)
-    return ret if x.size(0) == x.size(1) else ret.pad(tuple(None if d == ret.size(0) else (k, d-ret.shape[0]-k) for d in x.shape))
-
-  def OptionalHasElement(x:Tensor|None=None): return Tensor(x is not None and x.numel() > 0)
-  def OptionalGetElement(x:Tensor|None=None): return x if x is not None else Tensor([])
-  def ConstantOfShape(shape:list[int], value:Tensor|None=None):
-    if value is None: value = Tensor(0, dtype=dtypes.float32)
-    if shape == [0]: return Tensor([], dtype=value.dtype)
-    return value.expand(shape)
-
-  def Size(data:Tensor): return data.numel()
-  def Shape(data:Tensor, end:int|None=None, start:int=0): return Tensor(data.shape[start:end], dtype=dtypes.int64)
-
-  # ***** Unary Ops (math) *****
-  def Not(x:Tensor): return x.logical_not()
-  def Clip(x: Tensor, min:Tensor|None=None, max:Tensor|None=None): return x if min is None and max is None else x.clip(min, max)  # noqa: A002
-  def IsInf(x:Tensor, detect_negative:int=1, detect_positive:int=1): return x.isinf(bool(detect_positive), bool(detect_negative))
-
-  # ***** Unary Ops (activation) *****
-  def softmax_1(x:Tensor, axis:int=1): return x.softmax(axis)
-  def softmax_13(x:Tensor, axis:int=-1): return x.softmax(axis)
-  Softmax = {OpSetId(Domain.ONNX, 1):softmax_1, OpSetId(Domain.ONNX, 13):softmax_13}
-  def HardSigmoid(x:Tensor, alpha:float=0.2, beta:float=0.5): return (alpha*x + beta).clip(0, 1)
-  def Gelu(x:Tensor, approximate:str|None=None): return x.gelu() if approximate == "tanh" else 0.5 * x * (1 + (x/math.sqrt(2)).erf())
-  def BiasGelu(x: Tensor, bias: Tensor, approximate: str | None = None) -> Tensor: return Gelu(x + bias, approximate)
-  def FastGelu(x:Tensor, bias:Tensor|None=None): return (x + bias).gelu() if bias is not None else x.gelu() # this is tanh approximated
-  def PRelu(X:Tensor, slope:Tensor): return (X > 0).where(X, X * slope)
-  def LeakyRelu(X:Tensor, alpha:float=0.01): return X.leaky_relu(alpha)
-  def ThresholdedRelu(X:Tensor, alpha:float=1.0): return (X > alpha).where(X, 0)
-  def LogSoftmax(x: Tensor, axis:int=-1): return x.log_softmax(axis)
-  def Binarizer(x:Tensor, threshold:float=0.0): return (x > threshold).float()
-
-  # ***** Unary Ops (broadcasted) *****
-  def Add(x:Tensor,y:Tensor, broadcast=None, axis=None): return x + y
-  def Sub(x:Tensor|int,y:Tensor): return x - y # some test has input as int
-  def Div(x:Tensor,y:Tensor): return x.div(y, rounding_mode='trunc' if dtypes.is_int(x.dtype) else None)
-  def Less(x:Tensor,y:Tensor): return x < y
-  def LessOrEqual(x:Tensor,y:Tensor): return x <= y
-  def Greater(x:Tensor,y:Tensor): return x > y
-  def GreaterOrEqual(x:Tensor,y:Tensor): return x >= y
-  def Equal(x:Tensor,y:Tensor): return x == y
-  def And(x:Tensor,y:Tensor): return (x==y).where(x, False)
-  def Or(x:Tensor,y:Tensor): return (x==y).where(x, True)
-  def Xor(x:Tensor,y:Tensor): return x.bool().bitwise_xor(y.bool())
-  def BitwiseAnd(x:Tensor,y:Tensor): return x & y
-  def BitwiseOr(x:Tensor,y:Tensor): return x | y
-  def BitwiseXor(x:Tensor,y:Tensor): return x ^ y
-  def BitwiseNot(x:Tensor): return ~x
-  def Mod(x:Tensor,y:Tensor,fmod=0):
-    if fmod: return x - x.div(y, rounding_mode="trunc") * y
-    return x % y
-
-  # ***** Casting Ops *****
-  # TODO: saturate
-  def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(dtype_parse(to, "Cast op"))
-  def CastLike(x:Tensor, target_type:Tensor, saturate:int=1): return x.cast(target_type.dtype)
-
-  # ***** Reduce Ops *****
-  def Max(*data_0:Tensor): return functools.reduce(Tensor.maximum, data_0)
-  def Min(*data_0:Tensor): return functools.reduce(Tensor.minimum, data_0)
-  def Sum(*data_0:Tensor): return functools.reduce(Tensor.add, data_0)
-  def Mean(*data_0:Tensor): return Sum(*data_0) / len(data_0)
-  def ReduceMax(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
-    return data.max(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceMin(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
-    return data.min(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceSum(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
-    return data.sum(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceMean(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
-    return data.mean(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceSumSquare(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
-    return ReduceSum(data.square(), axes, keepdims, noop_with_empty_axes)
-  def ReduceProd(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
-    return data.prod(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceL1(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
-    return ReduceSum(data.abs(), axes, keepdims, noop_with_empty_axes)
-  def ReduceL2(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
-    return ReduceSumSquare(data, axes, keepdims, noop_with_empty_axes).sqrt()
-  def ReduceLogSum(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
-    return ReduceSum(data, axes, keepdims, noop_with_empty_axes).log()
-  def ReduceLogSumExp(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
-    return ReduceSum(data.exp(), axes, keepdims, noop_with_empty_axes).log()
-  def ArgMax(x:Tensor, axis:int=0, keepdims:int=1, select_last_index:int=0):
-    if select_last_index: return ((x.shape[axis]-1) - x.flip(axis).argmax(axis, keepdim=keepdims)).cast(dtypes.int64)
-    return x.argmax(axis, keepdim=keepdims).cast(dtypes.int64)
-  def ArgMin(x, axis:int=0, keepdims:int=1, select_last_index:int=0):
-    return ArgMax(-x, axis=axis, keepdims=keepdims, select_last_index=select_last_index)
-
-  # ***** Movement Ops *****
-  def Reshape(data:Tensor, shape:list[int], allowzero:int=0):
-    return data.reshape([x if x != 0 else (0 if allowzero else data.shape[i]) for i,x in enumerate(shape)])
-  def Flatten(x:Tensor, axis:int=1): return x.reshape(prod(x.shape[0:axis]), -1)
-  def Expand(x:Tensor, shape:list[int]): return x.expand(_broadcast_shape(x.shape, tuple(shape)))
-  def Shrink(x:Tensor, bias:float=0.0, lambd:float=0.5): return (x < -lambd)*(x+bias) + (x > lambd)*(x-bias)
-  def Transpose(x:Tensor, perm:list[int]|None=None): return x.permute(order=perm or list(range(x.ndim)[::-1]))
-
-  def Squeeze(data:Tensor, axes:list[int]|None=None):
-    return data.squeeze() if axes is None else functools.reduce(lambda d, dim: d.squeeze(dim), sorted(axes, reverse=True), data)
-  def Unsqueeze(data:Tensor, axes:list[int]): return functools.reduce(lambda d, dim: d.unsqueeze(dim), sorted(axes), data)
-
-  def Tile(x:Tensor, repeats:list[int]): return x.repeat(repeats)
-  def Concat(*xs:Tensor, axis:int): return Tensor.cat(*xs, dim=axis)
-  def Slice(data:Tensor, starts:list[int], ends:list[int], axes:list[int]|None=None, steps:list[int]|None=None):
-    axes = axes or list(range(data.ndim))
-    steps = steps or [1]*data.ndim
-    slices = [slice(0,x,1) for x in data.shape]
-    for i, axis in enumerate(axes): slices[axis] = slice(starts[i], ends[i], steps[i])
-    return data[tuple(slices)]
-
-  def Split(data:Tensor, split:list[int]|None=None, num_outputs:int=0, axis:int=0):
-    sz = data.shape[axis]
-    if split is None: split = [sz // num_outputs + (1 if i < sz % num_outputs else 0) for i in range(num_outputs)]
-    return data.split(split, axis)
-
-  def Pad(x:Tensor, pads:list[int], constant_value:ConstType|None=None, axes:list[int]|None=None,
-          mode:Literal["constant", "reflect", "edge", "wrap"]="constant", value=0):
-    value = constant_value or value
-    axes = axes or list(range(x.ndim))
-    real_pads = [0] * (x.ndim*2)
-    for i,axis in enumerate(axes): real_pads[axis%x.ndim], real_pads[axis%x.ndim+x.ndim] = pads[i], pads[i+len(axes)]
-    return x.pad(padding=_onnx_pads_to_tiny_pads(real_pads), mode={"edge":"replicate", "wrap":"circular"}.get(mode, mode), value=value)
-
-  def CenterCropPad(t:Tensor, shape:list[int], axes:list[int]|None=None):
-    shrink_arg:list[None|tuple[int,int]] = [None] * t.ndim
-    pad_arg:list[None|tuple[int,int]] = [None] * t.ndim
-    for s, x in zip(shape, axes or range(t.ndim)):
-      tx = t.shape[x]
-      if s < tx: shrink_arg[x] = (tx//2 - (s+1)//2, tx//2 + s//2)
-      elif s > tx: pad_arg[x] = ((s-tx)//2, (s-tx+1)//2)
-    return t.shrink(tuple(shrink_arg)).pad(tuple(pad_arg))
-
-  # ***** Processing Ops *****
-  def AveragePool(X: Tensor, kernel_shape:list[int], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, count_include_pad:int=0,
-                  dilations:list[int]|int=1, pads:list[int]|int=0, strides:list[int]|int=1):
-    return X.avg_pool2d(kernel_shape, strides, dilations, _resolve_pool_pads(X, pads, kernel_shape, dilations, strides, auto_pad),
-                        ceil_mode=ceil_mode, count_include_pad=count_include_pad)
-
-  def MaxPool(X: Tensor, kernel_shape:list[int], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, dilations:list[int]|int=1, pads:list[int]|int=0,
-              storage_order:int=0, strides:list[int]|int=1):
-    pads = _resolve_pool_pads(X, pads, kernel_shape, dilations, strides, auto_pad)
-    ret, idx = X.max_pool2d(kernel_shape, strides, dilations, pads, ceil_mode=ceil_mode, return_indices=True)
-    return ret, idx.transpose(-2, -1).cast(dtypes.int64) if storage_order else idx.cast(dtypes.int64)
-
-  def Conv(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:list[int]|int=1, group:int=1,
-          kernel_shape:list[int]|None=None, pads:list[int]|int=0, strides:list[int]|int=1):
-    return X.conv2d(W, B, stride=strides, groups=group, dilation=dilations,
-                    padding=_resolve_pool_pads(X, pads, kernel_shape or W.shape[2:], dilations, strides, auto_pad))
-
-  def ConvTranspose(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:list[int]|int=1, group:int=1,
-                    kernel_shape:list[int]|None=None, pads:list[int]|None=None, output_shape:list[int]|None=None, output_padding:list[int]|int=0,
-                    strides:list[int]|int=1):
-    input_shape, kernel_shape = X.shape[2:], (kernel_shape or W.shape[2:])
-    strides, dilations, output_padding = (make_tuple(x, len(input_shape)) for x in (strides, dilations, output_padding))
-    if output_shape is not None: # we pad according to output_shape
-      pads = _auto_pad([s*(i-1) + op + ((k-1)*d+1) - os for s,i,op,k,d,os in
-                        zip(strides, input_shape, output_padding, kernel_shape, dilations, output_shape)], auto_pad)
-    if pads is None: # we generate pads
-      output_shape = output_shape or [X.shape[i+2] * strides[i] for i in range(len(strides))]
-      pads = [strides[i]*(input_shape[i]-1)+output_padding[i]+((kernel_shape[i]-1)*dilations[i]+1)-output_shape[i] for i in range(len(input_shape))]
-      pads = _auto_pad(pads, auto_pad) if auto_pad != "NOTSET" else [0] * len(input_shape) * 2
-    pads = _onnx_pads_to_tiny_pads(pads)
-    return X.conv_transpose2d(W, B, stride=strides, groups=group, dilation=dilations, padding=pads, output_padding=output_padding)
-
-  def MaxUnpool(xT: Tensor, xI: Tensor, outshape: list[int]|None=None, kernel_shape:list[int]=None, pads:list[int]|int=0, strides:list[int]|int=1):
-    return Tensor.max_unpool2d(xT, xI, kernel_shape, strides, 1, pads, outshape if outshape is None else tuple(outshape))
-
-  def GlobalAveragePool(X:Tensor): return X.mean(axis=tuple(range(2, X.ndim)), keepdim=True)
-  def GlobalMaxPool(X:Tensor): return X.max(axis=tuple(range(2, X.ndim)), keepdim=True)
-
-  def Gemm(A:Tensor, B:Tensor, C:Tensor|None=None, alpha:float=1.0, beta:float=1.0, transA:int=0, transB:int=0, broadcast=0):
-    ret = alpha * (A.transpose(transA) @ B.transpose(transB))
-    if C is not None: ret = ret + beta * (C if broadcast == 0 else C.reshape([-1 if i < len(C.shape) else 1 for i in range(ret.ndim)][::-1]))
-    return ret
-
-  def Einsum(*Inputs:list[Tensor], equation:str): return Tensor.einsum(equation, *Inputs)
-
-  def CumSum(X:Tensor, axis:int|list[int], exclusive:int=0, reverse:int=0):
-    axis = X._resolve_dim(_resolve_const(axis))
-    if reverse: X = X.flip(axis)
-    if exclusive: X = X.pad(tuple((1,0) if i == axis else None for i in range(X.ndim)))\
-                        .shrink(tuple((0,X.shape[axis]) if i == axis else None for i in range(X.ndim)))
-    return X.cumsum(axis).flip(axis) if reverse else X.cumsum(axis)
-
-  def Trilu(x:Tensor, k:int|list[int]=0, upper:int=1):
-    k_ = _resolve_const(k)
-    return x.triu(k_) if upper else x.tril(k_)
-
-  def Resize(X:Tensor, roi:list[float]|None=None, scales:list[float]|None=None, sizes:list[int]|None=None, antialias:int=0,
-            axes:list[int]|None=None, coordinate_transformation_mode:str='half_pixel', cubic_coeff_a:float=-0.75, exclude_outside:int=0,
-            extrapolation_value:float=0.0, keep_aspect_ratio_policy:str='stretch', mode:str='nearest', nearest_mode:str='round_prefer_floor'):
-    def _apply_nearest_mode(index: Tensor, input_dim, mode: str):
-      if mode == "round_prefer_floor": index = (index - 0.5).ceil()
-      elif mode == "round_prefer_ceil": index = (index + 0.5).floor()
-      elif mode in ["floor", "ceil"]: index = getattr(index, mode)()
-      else: raise ValueError(f"invalid {nearest_mode=}")
-      return index.cast(dtypes.int32).clip(0, input_dim-1)
-    def _apply_transformation(index: Tensor, input_dim, scale_dim, mode):
-      # TODO: needs more testing, not confident in this
-      # NOTE: their reference implementation differ from the implementation in their reference docs
-      # https://github.com/onnx/onnx/blob/main/onnx/reference/ops/op_resize.py
-      # https://github.com/onnx/onnx/blob/main/docs/Operators.md#Resize
-      output_dim = scale_dim * input_dim
-      if mode == "half_pixel": index = (index + 0.5) / scale_dim - 0.5
-      elif mode == "align_corners": index = index * (input_dim - 1) / (output_dim - 1) if output_dim != 1 else Tensor([0])
-      elif mode == "asymmetric": index = index / scale_dim
-      elif mode == "pytorch_half_pixel": index = (index + 0.5) / scale_dim - 0.5 if output_dim != 1 else Tensor([-0.5])
-      elif mode == "half_pixel_symmetric": index = input_dim / 2 * (1 - int(output_dim) / output_dim) + (index + 0.5) / scale_dim - 0.5
-      else: raise NotImplementedError(f"invalid {coordinate_transformation_mode=}")
-      return index.clip(0, input_dim-1)
-
-    scales, sizes = (None if scales is None else scales[2-(X.ndim-len(scales)):]), (None if sizes is None else sizes[2-(X.ndim-len(sizes)):])
-    # we pre permute the axes and permute back after resize
-    axes, input_shape, = (axes or list(range(X.ndim))), cast(tuple[int, ...], X.shape[2:]),
-    perm = [a for a in range(len(X.shape)) if a not in axes] + list(axes)
-    X = X.permute(*perm)
-
-    if sizes is not None:
-      if keep_aspect_ratio_policy in ["not_larger", "not_smaller"]:
-        scale_fxn = min if keep_aspect_ratio_policy == "not_larger" else max
-        scales = [scale_fxn([sizes[i] / input_shape[i] for i in range(len(input_shape)) if i+2 in axes])] * 2
-        sizes = [int((scales[0] * input_shape[i]) + 0.5) if i+2 in axes else input_shape[i] for i in range(X.ndim-2)]
-      else:
-        scales = [size / input_shape for size, input_shape in zip(sizes, input_shape)]
-    else:
-      sizes = [int(sc*sh) for sc, sh in zip(scales, input_shape)]
-
-    # NOTE: this transformation makes it so that we can't just call Tensor.interpolate
-    # in Tensor.interpolate, we use indexes without any transformation
-    indexes = []
-    for shape, size, scale in zip(input_shape, sizes, scales):
-      indexes.append(_apply_transformation(Tensor.arange(size), shape, scale, coordinate_transformation_mode))
-
-    if mode == "nearest":
-      indexes = [_apply_nearest_mode(index, shape, nearest_mode) for (index, shape) in zip(indexes, input_shape)]
-      X = X[(..., *Tensor.meshgrid(*indexes))]
-    if mode == "linear":
-      expand = list(X.shape)
-      for i in range(-len(sizes), 0):
-        reshape, index = [1] * X.ndim, indexes[i]
-        reshape[i] = expand[i] = sizes[i]
-        low, high, perc = [y.reshape(reshape).expand(expand) for y in (index.floor().int(), index.ceil().int(), index - index.floor())]
-        X = X.gather(i, low).lerp(X.gather(i, high), perc)
-    if mode == "cubic": raise NotImplementedError("cubic interpolation is not implemented")
-    return X.permute(*argsort(perm)) if perm else X
-  def Upsample(X, scales, mode): return Resize(X=X, scales=scales, mode=mode)  # deprecated
-
-  def TopK(X:Tensor, K:int|list[int], axis:int=-1, largest:int=1, sorted:int=1):  # noqa: A002
-    val, idx = X.topk(_resolve_const(K), axis, largest, sorted)
-    return val, idx.cast(dtypes.int64)
-
-  # ***** Neural Network Ops *****
-  def BatchNormalization(X:Tensor, scale:Tensor, B:Tensor, input_mean:Tensor, input_var:Tensor, epsilon:float=1e-05, momentum:float=0.9,
-                        training_mode:int=0, spatial=1, is_test=0):
-    if training_mode:
-      x_detached = X.detach()
-      current_mean = x_detached.mean(axis=(0,2,3))
-      y = (x_detached - current_mean.reshape(shape=[1, -1, 1, 1]))
-      current_var = (y*y).mean(axis=(0,2,3))
-      current_invstd = current_var.add(epsilon).rsqrt()
-
-      running_mean = input_mean * momentum + current_mean * (1 - momentum)
-      running_var = input_var * momentum + current_var * (1 - momentum)
-
-      return X.batchnorm(scale, B, current_mean, current_invstd), running_mean, running_var
-    return X.batchnorm(scale, B, input_mean, (input_var + epsilon).rsqrt())
-  def GroupNormalization(x:Tensor, scale:Tensor, bias:Tensor, num_groups:int, epsilon:float=1e-05):
-    x = x.reshape(x.shape[0], num_groups, -1).layernorm(eps=epsilon).reshape(x.shape)
-    return x * scale.reshape(1, -1, *[1] * (x.ndim-2)) + bias.reshape(1, -1, *[1] * (x.ndim-2))
-  def InstanceNormalization(x:Tensor, scale:Tensor, bias:Tensor, epsilon:float=1e-05):
-    return GroupNormalization(x, scale, bias, num_groups=x.shape[1], epsilon=epsilon)
-  def LayerNormalization(x:Tensor, scale:Tensor, bias:Tensor, axis:int=-1, epsilon:float=1e-05, stash_type:int=1):
-    assert stash_type == 1, "only float32 is supported"
-    axes = tuple(i for i in range(axis if axis >= 0 else x.ndim + axis, x.ndim))
-    mean = x.mean(axis=axes, keepdim=True)
-    return x.layernorm(axes, epsilon).mul(scale).add(bias), mean, (x.sub(mean)).square().mean(axis=axes, keepdim=True).add(epsilon).rsqrt()
-  def SkipLayerNormalization(x:Tensor, skip:Tensor, gamma:Tensor, beta:Tensor|None=None, bias:Tensor|None=None, epsilon:float=1e-12):
-    x = x + skip
-    if bias is not None: x = x + bias
-    ret = x.layernorm(eps=epsilon) * gamma
-    if beta is not None: ret = ret + beta
-    return ret, None, None, x
-  def EmbedLayerNormalization(input_ids: Tensor, segment_ids:Tensor, word_embedding:Tensor, position_embedding:Tensor,
-                              segment_embedding:Tensor, gamma=None, beta=None, mask:Tensor|None=None,
-                              position_ids:Tensor|None=None, epsilon=1e-12, mask_index_type=0):
-    # https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.EmbedLayerNormalization
-    assert (segment_ids is None) is (segment_embedding is None)
-    assert mask is None and not mask_index_type, "functionality not supported yet"  # TODO
-    input_shape = input_ids.shape
-    seq_length = input_shape[1]
-    compute_seg_emb = (segment_embedding is not None and segment_ids is not None)
-    vocab_size, max_position_embeddings = word_embedding.shape[0], position_embedding.shape[0]
-    type_vocab_size  = (segment_embedding.shape[0] if compute_seg_emb else None)
-
-    def embedding(x:Tensor, vocab_size, weight:Tensor) -> Tensor:
-      return x.unsqueeze(-1).expand(*x.shape, vocab_size)._one_hot_along_dim(vocab_size) @ weight
-
-    # bert embedding layer
-    if position_ids is None: position_ids = Tensor.arange(seq_length, requires_grad=False).unsqueeze(0).expand(*input_shape)
-    wrd_embedding_res = embedding(input_ids, vocab_size, word_embedding)
-    pos_embedding_res = embedding(position_ids, max_position_embeddings, position_embedding)
-    seg_embedding_res = embedding(segment_ids, type_vocab_size, segment_embedding) if compute_seg_emb else None
-
-    embedding_sum = wrd_embedding_res + pos_embedding_res
-    if seg_embedding_res is not None: embedding_sum = embedding_sum + seg_embedding_res
-    out = embedding_sum.layernorm(eps=epsilon) * gamma + beta
-    return out, None, embedding_sum
-  def MeanVarianceNormalization(x:Tensor, axis:list[int]=[0,2,3]):
-    return (x - x.mean(axis, keepdim=True)) / (x.std(axis, keepdim=True, correction=0) + 1e-9)
-
-  def OneHot(indices:Tensor, depth:float|int|list[int|float], values:Tensor, axis:int=-1):
-    # Scalar or Rank 1 tensor containing exactly one element
-    depth = int(_resolve_const(depth))
-    indices = indices.int()
-    indices = (indices < 0).where(indices+depth, indices)
-    return indices.unsqueeze(axis)._one_hot_along_dim(depth, dim=axis).where(values[1], values[0])
-
-  def DepthToSpace(X:Tensor, blocksize:int, mode:str="DCR"):
-    return X.rearrange("b (c h1 w1) h w -> b c (h h1) (w w1)" if mode=="CRD" else "b (h1 w1 c) h w -> b c (h h1) (w w1)", h1=blocksize, w1=blocksize)
-  def SpaceToDepth(X:Tensor, blocksize:int):
-    return X.rearrange("b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=blocksize, w1=blocksize)
-
-  # Reimplemented here because you need legacy RNG for passing ONNX tests.
-  def dropout_7(data:Tensor, ratio:float=0.5, training_mode:bool=False, seed:int|None=None):
-    import numpy as np
-    if not training_mode: return data, data.full_like(True, dtype=dtypes.bool)
-    if seed is not None:
-      rand = Tensor(np.random.RandomState(seed).random(cast(tuple[int,...], data.shape)), requires_grad=False, dtype=data.dtype, device=data.device)
-    else:
-      rand = data.rand_like(requires_grad=False)
-    mask = rand >= ratio
-    return data * mask / (1.0 - ratio), mask
-  # 6 with 'is_test' needed for https://github.com/MTlab/onnx2caffe/raw/refs/heads/master/model/MobileNetV2.onnx
-  def dropout_6(data:Tensor, ratio:float=0.5, is_test=0): return dropout_7(data, ratio, training_mode=not is_test)
-  Dropout = {OpSetId(Domain.ONNX, 6):dropout_6, OpSetId(Domain.ONNX, 7):dropout_7}
-
-  def LRN(x:Tensor, size:int, alpha:float=1e-4, beta:float=0.75, bias:float=1.0):
-    pooled_x = (x**2).rearrange('b c h w -> b 1 c (h w)').pad((0,0,(size-1)//2, size//2)).avg_pool2d((size, 1), 1)
-    return x / (pooled_x.reshape(x.shape) * alpha + bias).pow(beta)
-
-  def NegativeLogLikelihoodLoss(x:Tensor, target:Tensor, weight:Tensor|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean"):
-    return x.nll_loss(target, weight, ignore_index, reduction)
-  def SoftmaxCrossEntropyLoss(scores:Tensor, labels:Tensor, weights:Tensor|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean"):
-    log_probs = scores.log_softmax(1)
-    return log_probs.nll_loss(labels, weights, ignore_index, reduction), log_probs
-
-  def AffineGrid(theta:Tensor, size:list[int], align_corners:int=0):
-    N, _, *spatial_dims = size
-    def generate_grid(steps):
-      if align_corners: return Tensor.linspace(-1, 1, steps, device=theta.device)
-      return Tensor.linspace(-1+1/steps, 1-1/steps, steps, device=theta.device)
-    grids = Tensor.meshgrid(*(generate_grid(d) for d in spatial_dims))
-    base_grid = Tensor.stack(*reversed(grids), Tensor.ones_like(grids[0], device=theta.device), dim=-1)
-    base_grid = base_grid.reshape(1, prod(spatial_dims), len(grids)+1).expand(N, -1, -1)
-    return (base_grid @ theta.transpose(1, 2)).reshape(N, *spatial_dims, -1)
-
-  def attention_contrib(x:Tensor, weights:Tensor, bias:Tensor|None=None, mask_index:Tensor|None=None, past:Tensor|None=None,
-                        attention_bias:Tensor|None=None, past_sequence_length:Tensor|None=None,  do_rotary:int=0, mask_filter_value:float=-10000.0,
-                        num_heads:int|None=None, past_present_share_buffer:int|None=None, qkv_hidden_sizes:list[int]|None=None,
-                        rotary_embedding_dim:int|None=None, scale:float|None=None, unidirectional:int=0):
-    assert not do_rotary and not attention_bias, "TODO"
-    if qkv_hidden_sizes is None: qkv_hidden_sizes = [weights.shape[1] // 3] * 3
-    qkv = x.linear(weights, bias)
-    q, k, v = qkv.split(qkv_hidden_sizes, dim=2)
-
-    batch_size, seq_len, _ = x.shape
-    q_head_size, k_head_size, v_head_size = (sz // num_heads for sz in qkv_hidden_sizes)
-    q, k, v = (x.reshape(batch_size, seq_len, num_heads, hsz).transpose(1, 2) for x, hsz in zip((q, k, v), (q_head_size, k_head_size, v_head_size)))
-
-    present = None
-    if past is not None:
-      k, v = past[0].cat(k, dim=2), past[1].cat(v, dim=2)
-      present = k.stack(v)
-
-    if scale is None: scale = 1.0 / math.sqrt(q_head_size)
-    attn_scores = q @ k.transpose(-1, -2) * scale
-
-    if mask_index is not None:
-      assert 4 >= mask_index.ndim >= 1, f"{mask_index.ndim=}"
-      if mask_index.ndim != 1: mask = mask_index.bool()
-      else:
-        if mask_index.shape[0] == batch_size:
-          mask = Tensor.arange(attn_scores.shape[-1], requires_grad=False, device=mask_index.device).unsqueeze(0) < mask_index.unsqueeze(1)
-        elif mask_index.shape[0] == 2*batch_size:
-          end_positions = mask_index[:batch_size]
-          start_positions = mask_index[batch_size:]
-          arange = Tensor.arange(seq_len).unsqueeze(0)
-          mask = (arange < end_positions.unsqueeze(1)) & (arange >= start_positions.unsqueeze(1))
-        else: raise NotImplementedError("mask_index with shape (3 * batch_size + 2) is not implemented")
-      while mask.ndim < 4: mask = mask.unsqueeze(1)
-      attn_scores = mask.where(attn_scores, mask_filter_value)
-
-    if unidirectional:
-      causal_mask = Tensor.ones((seq_len, seq_len), dtype=dtypes.bool).tril()
-      attn_scores = causal_mask.where(attn_scores, mask_filter_value)
-
-    output = attn_scores.softmax(-1) @ v
-    output = output.transpose(1, 2).reshape(batch_size, seq_len, -1)
-    return output, present
-
-  def attention_onnx(Q:Tensor, K:Tensor, V:Tensor, attn_mask:Tensor|None=None, past_key:Tensor|None=None, past_value:Tensor|None=None,
-                     is_causal:int=0, kv_num_heads:int|None=None, q_num_heads:int|None=None, qk_matmul_output_mode:int=0, scale:float|None=None,
-                     softcap:float=0.0, softmax_precision:int|None=None):
-    input_shape_len = Q.ndim
-    if input_shape_len == 3:
-      assert q_num_heads is not None and kv_num_heads is not None
-      Q = Q.reshape(Q.shape[0], q_num_heads, Q.shape[1], -1)
-      K = K.reshape(K.shape[0], kv_num_heads, K.shape[1], -1)
-      V = V.reshape(V.shape[0], kv_num_heads, V.shape[1], -1)
-
-    if past_key is not None: K = past_key.cat(K, dim=2)
-    if past_value is not None: V = past_value.cat(V, dim=2)
-    present_key, present_value = K, V
-
-    _q_heads, _kv_heads = q_num_heads or Q.shape[1], kv_num_heads or K.shape[1]
-    if _q_heads != _kv_heads:
-      K = K.repeat((1, _q_heads // _kv_heads, 1, 1))
-      V = V.repeat((1, _q_heads // _kv_heads, 1, 1))
-
-    effective_scale = scale if scale is not None else 1.0 / (Q.shape[-1] ** 0.5)
-    scores = (Q @ K.transpose(-1, -2)) * effective_scale
-    qk_matmul_return_val = scores
-
-    if is_causal:
-      causal_mask = Tensor.ones(Q.shape[-2], K.shape[-2], device=Q.device, dtype=dtypes.bool, requires_grad=False).tril(0)
-      scores = scores.masked_fill(causal_mask.logical_not(), -float("inf"))
-
-    if attn_mask is not None:
-      mask_to_add = attn_mask.where(0, -float("inf")) if attn_mask.dtype == dtypes.bool else attn_mask
-      scores = scores + mask_to_add
-    if qk_matmul_output_mode == 1: qk_matmul_return_val = scores
-
-    if softcap > 0.0: scores = (scores / softcap).tanh() * softcap
-    if qk_matmul_output_mode == 2: qk_matmul_return_val = scores
-
-    if softmax_precision: scores = scores.cast({1: dtypes.float32, 10: dtypes.float16, 16: dtypes.bfloat16}[softmax_precision])
-    qk_softmax = scores.softmax(-1).cast(Q.dtype)
-    if qk_matmul_output_mode == 3: qk_matmul_return_val = qk_softmax
-
-    output = (qk_softmax @ V).cast(Q.dtype)
-    if input_shape_len == 3: output = output.permute(0, 2, 1, 3).reshape(Q.shape[0], Q.shape[2], -1)
-    return output, present_key, present_value, qk_matmul_return_val
-  Attention = {OpSetId(Domain.ONNX, 1): attention_onnx, OpSetId(Domain.MICROSOFT_CONTRIB_OPS, 1): attention_contrib}
-
-  def RMSNormalization(X:Tensor, scale:Tensor, axis:int=-1, epsilon:float=1e-5):
-    norm = X.square().mean(axis=tuple(range(axis + X.ndim if axis < 0 else axis, X.ndim)), keepdim=True).add(epsilon).rsqrt()
-    return X * norm * scale
-
-  def RotaryEmbedding(X:Tensor, cos_cache:Tensor, sin_cache:Tensor, position_ids:Tensor|None=None, interleaved:int=0, num_heads:int|None=None,
-                      rotary_embedding_dim:int=0):
-    original_input_shape = X.shape
-
-    if X.ndim == 4: X = X.permute(0, 2, 1, 3)
-    elif X.ndim == 3:
-      assert num_heads is not None, "num_heads must be provided for 3D input"
-      X = X.reshape(*X.shape[:-1], num_heads, X.shape[-1] // num_heads)
-
-    head_size = X.shape[-1]
-    rot_dim = rotary_embedding_dim or head_size
-    x_rotate, x_pass = X[..., :rot_dim], X[..., rot_dim:]
-
-    cos = cos_cache[position_ids] if position_ids is not None else cos_cache[:X.shape[1]]
-    sin = sin_cache[position_ids] if position_ids is not None else sin_cache[:X.shape[1]]
-    cos = cos[..., :rot_dim//2].unsqueeze(2)
-    sin = sin[..., :rot_dim//2].unsqueeze(2)
-
-    if interleaved:
-      x1, x2 = x_rotate[..., ::2], x_rotate[..., 1::2]
-      real = x1 * cos - x2 * sin
-      imag = x1 * sin + x2 * cos
-      x_rotated = Tensor.stack(real, imag, dim=-1).flatten(start_dim=-2)
-    else:
-      x1, x2 = x_rotate.chunk(2, dim=-1)
-      real = x1 * cos - x2 * sin
-      imag = x1 * sin + x2 * cos
-      x_rotated = real.cat(imag, dim=-1)
-
-    output = x_rotated.cat(x_pass, dim=-1)
-    return output.flatten(start_dim=2) if len(original_input_shape) == 3 else output.permute(0, 2, 1, 3)
-
-  # ***** Indexing Ops *****
-  def ArrayFeatureExtractor(x:Tensor, indices:Tensor): return x[..., indices]
-
-  def Gather(x:Tensor, indices:Tensor, axis:int=0):
-    if indices.numel() < 9: # NOTE lessor kernels for smaller indices but kernel number increases depending on size of indices
-      ret_shape = x.shape[:axis] + indices.shape + x.shape[axis+1:]
-      if indices.ndim > 1: indices = indices.flatten()
-      index_consts = [_cached_to_python_const(indices)] if indices.shape == () else _cached_to_python_const(indices)
-      index_consts = [x.shape[axis]+i if i<0 else i for i in index_consts]
-      args = [[(0,x) if j != axis else (i,i+1) for j, x in enumerate(x.shape)] for i in index_consts]
-      return x.shrink(arg=tuple(args[0])).cat(*[x.shrink(arg=tuple(arg)) for arg in args[1:]], dim=axis).reshape(ret_shape)
-    # NOTE faster gather, fixed number of kernels, but exceeds limited kernels for openpilot
-    return x[tuple([slice(None) if i != axis else indices for i in range(x.ndim)])]
-  def Scatter(*args, **kwargs): return ScatterElements(*args, **kwargs) # deprecated
-
-  def GatherND(x:Tensor, indices:Tensor, batch_dims:int=0):
-    if batch_dims == 0: return x[tuple(i.squeeze(-1) for i in indices.split(1, -1))]
-    x_shape, i_shape = x.shape, indices.shape
-    b = math.prod(x.shape[dim] for dim in range(batch_dims))
-    # NOTE: each batched dim of both input and indices are equal
-    x = x.reshape(b, *x.shape[batch_dims:])
-    indices = indices.reshape(b, *indices.shape[batch_dims:])
-    b_idx = Tensor.arange(b, device=x.device).reshape(b, *(1,)*(indices.ndim - 2)).expand(*indices.shape[:-1])
-    ret = x[(b_idx,) + tuple(i.squeeze(-1) for i in indices.split(1, -1))]
-    return ret.reshape(*x_shape[:batch_dims], *i_shape[batch_dims:-1], *ret.shape[indices.ndim-1:])
-  def ScatterND(x:Tensor, indices:Tensor, updates:Tensor, reduction:Literal["none", "add", "mul"]='none'):
-    assert updates.shape == indices.shape[:-1] + x.shape[cast(int, indices.shape[-1]):]
-    x = x.contiguous()
-    for index, u in zip(indices.split(1, 0), updates.split(1, 0)):
-      i = tuple(idx.squeeze(-1) for idx in index.squeeze(0).split(1, -1))
-      u = u.squeeze(0)
-      if reduction == "none": x[i] = u
-      elif reduction == "add": x[i] += u
-      elif reduction == "mul": x[i] *= u
-      else: raise NotImplementedError("reduction doesn't support max or min")
-    return x
-
-  def ScatterElements(x: Tensor, indices: Tensor, updates: Tensor, axis=0, reduction:Literal["none", "add", "mul", "min", "max"]="none"):
-    indices = (indices < 0).where(x.shape[axis], 0) + indices
-    if reduction == "none": return x.scatter(axis, indices, updates)
-    return x.scatter_reduce(axis, indices, updates, {"add": "sum", "mul": "prod", "min": "amin", "max": "amax"}.get(reduction))
-  def GatherElements(x:Tensor, indices:Tensor, axis:int):
-    indices = (indices < 0).where(x.shape[axis], 0) + indices
-    return x.gather(axis, indices)
-
-  def Compress(inp:Tensor, condition:list[bool], axis:int|None=None):
-    if axis is None:
-      inp = inp.flatten()
-      axis = 0
-    if axis < 0: axis += inp.ndim
-    con = Tensor([i for i,cond in enumerate(condition) if cond]) # compress in python
-    return inp[tuple(con if i == axis else slice(None) for i in range(inp.ndim))]
-
-  # ***** Quantization Ops *****
-  def QuantizeLinear(x:Tensor, y_scale:Tensor, y_zero_point:Tensor|int=0, axis:int=1, block_size:int=0, output_dtype:int=0, saturate=1):
-    if isinstance(y_zero_point, Tensor): out_dtype = y_zero_point.dtype
-    elif output_dtype != 0: out_dtype = dtype_parse(output_dtype, "QuantizeLinear op")
-    else: out_dtype = dtypes.uint8
-    y_scale, y_zero_point = _prepare_quantize(x, y_scale, y_zero_point, axis, block_size)
-    if out_dtype == dtypes.uchar:
-      # this appears to work in practice, at least for uchar out_dtype. it folds with the quantize stuff
-      ret = _clamp_cast((x / y_scale + 0.4999999 + y_zero_point).int(), out_dtype)
-    else:
-      ret = _clamp_cast(((x / y_scale).round() + y_zero_point), out_dtype)
-    return ret.contiguous()
-
-  def DynamicQuantizeLinear(x: Tensor):
-    # only support uint8
-    qmin, qmax = dtypes.min(dtypes.uint8), dtypes.max(dtypes.uint8)
-    scale = (x.max().maximum(0) + ((-x).max()).maximum(0)) / (qmax - qmin)
-    zero_point = _clamp_cast((qmin - x.min() / scale).round(), dtypes.uint8)
-    y = _clamp_cast((x / scale).round() + zero_point, dtypes.uint8)
-    return y, scale, zero_point
-
-  def DequantizeLinear(x:Tensor, x_scale:Tensor, x_zero_point:Tensor|int=0, axis:int=1, block_size:int=0):
-    x_scale, x_zero_point = _prepare_quantize(x, x_scale, x_zero_point, axis, block_size)
-    return ((x.int() - x_zero_point) * x_scale).cast(x_scale.dtype)
-
-  def QLinearConv(x:Tensor, x_scale:Tensor, x_zero_point:Tensor|int, w:Tensor, w_scale:Tensor, w_zero_point:Tensor|int, y_scale:Tensor,
-                  y_zero_point: Tensor|int, B:Tensor|None=None, **opts):
-    return _qlinearop_quantized(Conv, [x,w], [x_zero_point,w_zero_point], [x_scale,w_scale], y_scale, y_zero_point, **{"B":B, **opts})
-
-  def QLinearMatMul(a:Tensor, a_scale:Tensor, a_zero_point:Tensor|int, b:Tensor, b_scale:Tensor, b_zero_point:Tensor|int, y_scale:Tensor,
-                    y_zero_point:Tensor|int) -> Tensor:
-    return _qlinearop_quantized(Tensor.matmul, [a,b], [a_zero_point,b_zero_point], [a_scale,b_scale], y_scale, y_zero_point)
-
-  def QLinearAdd(a:Tensor, a_scale:Tensor, a_zero_point:Tensor, b:Tensor, b_scale:Tensor, b_zero_point:Tensor, c_scale:Tensor, c_zero_point:Tensor):
-    return _qlinearop_float(Tensor.add, [a,b], [a_zero_point,b_zero_point], [a_scale,b_scale], c_scale, c_zero_point)
-
-  def QLinearMul(a:Tensor, a_scale:Tensor, a_zero_point:Tensor, b:Tensor, b_scale:Tensor, b_zero_point:Tensor, c_scale:Tensor, c_zero_point:Tensor):
-    return _qlinearop_quantized(Tensor.mul, [a,b], [a_zero_point,b_zero_point], [a_scale,b_scale], c_scale, c_zero_point)
-
-  def QLinearGlobalAveragePool(X:Tensor, x_scale:Tensor, x_zero_point:Tensor, y_scale:Tensor, y_zero_point:Tensor, channels_last:int):
-    assert channels_last == 0, "TODO NHWC"
-    return _qlinearop_float(GlobalAveragePool, [X], [x_zero_point], [x_scale], y_scale, y_zero_point)
-
-  def ConvInteger(x: Tensor, w: Tensor, x_zero_point: Tensor | int = 0, w_zero_point: Tensor | int = 0, B: Tensor | None = None, **opts) -> Tensor:
-    return _op_integer(Conv, [x,w], [x_zero_point,w_zero_point], **{"B":B, **opts})
-
-  def MatMulInteger(A: Tensor, B: Tensor, a_zero_point: Tensor | int = 0, b_zero_point: Tensor | int = 0) -> Tensor:
-    return _op_integer(Tensor.matmul, [A,B], [a_zero_point,b_zero_point])
-
-  # ***** Training Ops *****
-  # NOTE: onnx training ops actually don't need the state for optim, all the ops work in a functional way, but we still can reuse optim.py code
-  @_onnx_training(3)
-  def Adagrad(R:Tensor, T:int, *inputs:Tensor, decay_factor:float=0.0, epsilon:float=0.0, norm_coefficient:float=0.0):
-    X, G, H = (i.detach() for i in inputs)
-    grad = norm_coefficient * X + G
-    H.assign(H + grad.square())
-    up = grad / (H.sqrt() + epsilon)
-    r = R / (1 + T * decay_factor)
-    X.assign(X.detach() - r * up)
-    return [X, H]
-
-  @_onnx_training(4)
-  def Adam(R:Tensor, T:int, *inputs:Tensor, alpha:float=0.9, beta:float=0.999, epsilon:float=0.0, norm_coefficient:float=0.0,
-          norm_coefficient_post:float=0.0):
-    from tinygrad.nn.optim import Adam as TinyAdam
-    X, G, V, H = inputs
-    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
-    # 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:
-      # `T-1` since it's applied again at the start of `_step`
-      opt.b1_t = Tensor([alpha**(T-1)], dtype=dtypes.float32, device=X.device, requires_grad=False)
-      opt.b2_t = Tensor([beta**(T-1)], dtype=dtypes.float32, device=X.device, requires_grad=False)
-    opt.step()
-    X = (1 - norm_coefficient_post) * X
-    return [X, V, H]
-
-  @_onnx_training(3)
-  def Momentum(R:Tensor, T:int, *inputs:Tensor, alpha:float, beta:float, mode:str, norm_coefficient:float):
-    X, G, V = (i.detach() for i in inputs)
-    grad = norm_coefficient * X + G
-    # NOTE: this beta_adjusted term makes it so we can't use SGD for nesterov
-    beta_adjusted = beta if T > 0 else 1
-    V.assign(alpha * V + grad * beta_adjusted)
-    X.assign(X - R * (V if mode == "standard" else (grad + alpha * V)))
-    return [X, V]
-
-  def Gradient(*inputs:Tensor, y:str, intermediate_tensors:dict[str, Tensor], **_):
-    intermediate_tensors[y].backward()
-    return tuple([t.grad for t in inputs])
-
-  return {
-    # Tensor ops
-    **{op: getattr(Tensor, op.lower()) for op in ("Neg", "Reciprocal", "Pow", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan",
-    "Asin", "Acos", "Atan", "Relu", "Sigmoid", "MatMul", "Floor", "Ceil", "IsNaN", "Softplus", "HardSwish", "Where", "Mul", "Sinh", "Cosh",
-    "Tanh", "Softsign", "Asinh", "Acosh", "Atanh",  "Elu", "Celu", "Selu", "Round", "Erf")},
-    # Implemented ops
-    **{name:obj for name,obj in locals().items() if isinstance(obj, types.FunctionType) and not name.startswith("_") and name[0].isupper()},
-    # Version ops
-    **{name:obj for name,obj in locals().items() if isinstance(obj, dict)},
-  }
-
-onnx_ops = get_onnx_ops()
@@ -1,7 +1,6 @@
 from tinygrad import Tensor
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner
-from extra.onnx import OnnxValue
+from tinygrad.frontend.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort

@@ -1,207 +0,0 @@
-# https://github.com/onnx/onnx/blob/main/onnx/onnx.proto3
-
-import os, pathlib, struct
-from io import BufferedReader
-from types import SimpleNamespace
-from tinygrad.nn.state import TensorIO
-from tinygrad.tensor import Tensor
-
-# Protobuf Wire Types
-WIRETYPE_VARINT = 0; WIRETYPE_FIXED64 = 1; WIRETYPE_LENGTH_DELIMITED = 2; WIRETYPE_START_GROUP = 3; WIRETYPE_END_GROUP = 4; WIRETYPE_FIXED32 = 5 # noqa: E702
-
-# TensorProto.DataType
-class TensorDataType:
-  UNDEFINED = 0; FLOAT = 1; UINT8 = 2; INT8 = 3; UINT16 = 4; INT16 = 5; INT32 = 6; INT64 = 7 # noqa: E702
-  STRING = 8; BOOL = 9; FLOAT16 = 10; DOUBLE = 11; UINT32 = 12; UINT64 = 13; COMPLEX64 = 14; COMPLEX128 = 15; BFLOAT16 = 16 # noqa: E702
-
-# AttributeProto.AttributeType
-class AttributeType:
-  UNDEFINED = 0; FLOAT = 1; INT = 2; STRING = 3; TENSOR = 4; GRAPH = 5; SPARSE_TENSOR = 11; TYPE_PROTO = 13; FLOATS = 6; INTS = 7 # noqa: E702
-  STRINGS = 8; TENSORS = 9; GRAPHS = 10; SPARSE_TENSORS = 12; TYPE_PROTOS = 14 # noqa: E702
-
-class PBType: FLOAT = 1; INT = 2; STRING = 3; FLOATS = 4; INTS = 5; STRINGS = 6; BYTES = 7; SUB = 8 # noqa: E702
-
-PB_INFOS: dict[str, dict] = {
-  "OperatorSetIdProto": {1: ("domain", PBType.STRING), 2: ("version", PBType.INT)},
-  "StringStringEntryProto": {1: ("key", PBType.STRING), 2: ("value", PBType.STRING)},
-  "TensorProto": {1: ("dims", PBType.INT, True), 2: ("data_type", PBType.INT), 4: ("float_data", PBType.FLOATS),
-    13: ("external_data", PBType.SUB, True, "StringStringEntryProto"), 14: ("data_location", PBType.INT),
-    5: ("int32_data", PBType.INTS), 7: ("int64_data", PBType.INTS), 8: ("name", PBType.STRING), 9: ("raw_data", PBType.BYTES),
-    10: ("double_data", PBType.FLOATS), 11: ("uint64_data", PBType.INTS)},
-  "TensorShapeProtoDimension": {1: ("dim_value", PBType.INT), 2: ("dim_param", PBType.STRING)},
-  "TensorShapeProto": {1: ("dim", PBType.SUB, True, "TensorShapeProtoDimension")},
-  "ModelProto": {1: ("ir_version", PBType.INT), 5: ("model_version", PBType.INT),
-    2: ("producer_name", PBType.STRING), 3: ("producer_version", PBType.STRING), 4: ("domain", PBType.STRING), 6: ("doc_string", PBType.STRING),
-    7: ("graph", PBType.SUB, False, ("GraphProto", lambda: {"node": [], "initializer": [], "input": [], "output": [], "value_info": []})),
-    8: ("opset_import",PBType.SUB, True, "OperatorSetIdProto")},
-  "GraphProto": {2: ("name", PBType.STRING), 10: ("doc_string", PBType.STRING),
-    1: ("node", PBType.SUB, True, ("NodeProto", lambda: {"input": [], "output": [], "attribute": [], "domain": None})),
-    5: ("initializer", PBType.SUB, True, ("TensorProto", lambda: {"dims": [], "float_data": None, "int32_data": None, "string_data": None,
-                                                                  "int64_data": None, "double_data": None, "uint64_data": None, "raw_data": None})),
-    11: ("input", PBType.SUB, True, "ValueInfoProto"), 12: ("output", PBType.SUB, True, "ValueInfoProto")},
-  "NodeProto": { 1: ("input", PBType.STRING, True), 2: ("output", PBType.STRING, True), 3: ("name", PBType.STRING),
-    4: ("op_type", PBType.STRING), 6: ("doc_string", PBType.STRING), 7: ("domain", PBType.STRING),
-    5: ("attribute", PBType.SUB, True, ("AttributeProto", lambda: {"floats": [], "ints": [], "strings": []}))},
-  "AttributeProto": {1: ("name", PBType.STRING), 20: ("type", PBType.INT), 3: ("i", PBType.INT), 8: ("ints", PBType.INT, True),
-    2: ("f", PBType.FLOAT), 7: ("floats", PBType.FLOAT, True), 4: ("s", PBType.BYTES), 9: ("strings", PBType.BYTES, True),
-    5:("t", PBType.SUB, False, ("TensorProto", lambda: {"dims": [], "float_data": None, "int32_data": None, "string_data": None, "int64_data": None,
-                                                        "double_data": None, "uint64_data": None, "raw_data": None}))},
-  "ValueInfoProto": {1: ("name", PBType.STRING), 2: ("type", PBType.SUB, False, "TypeProto"), 3: ("doc_string", PBType.STRING)},
-  "TypeProto": {1: ("tensor_type", PBType.SUB, False, "TypeProtoTensor"), 4: ("sequence_type", PBType.SUB, False, "TypeProtoSequence"),
-    9: ("optional_type", PBType.SUB, False, "TypeProtoOptional"), 6: ("denotation", PBType.STRING)},
-  "TypeProtoSequence": {1: ("elem_type", PBType.SUB, False, "TypeProto")},
-  "TypeProtoOptional": {1: ("elem_type", PBType.SUB, False, "TypeProto")},
-  "TypeProtoTensor": {1: ("elem_type", PBType.INT), 2: ("shape", PBType.SUB, False, ("TensorShapeProto", lambda: {"dim": []}))},
-}
-
-def onnx_load(fn: Tensor|str|pathlib.Path, load_external_data: bool=True):
-  parser = OnnxParser(fn, load_external_data)
-  onnx_model = parser.parse()
-  model = dict_to_namespace(onnx_model)
-  return model
-
-def gen_result(obj: dict, key_name, val, repeated: bool):
-  if repeated: obj.setdefault(key_name, []).append(val)
-  else: obj[key_name] = val
-
-def dict_to_namespace(d):
-  if isinstance(d, dict): return SimpleNamespace(**{k: dict_to_namespace(v) for k, v in d.items()})
-  elif isinstance(d, list): return [dict_to_namespace(i) for i in d]
-  return d
-
-class OnnxParser:
-  def __init__(self, inp: Tensor|str|pathlib.Path, load_external_data: bool=True):
-    self.file_path: pathlib.Path|None = None
-    self.load_external_data = load_external_data
-    if not isinstance(inp, Tensor):
-      self.file_path = pathlib.Path(inp)
-      self.tensor = Tensor(self.file_path)
-    else: self.tensor = inp
-    self.attr_func_dict = { PBType.BYTES: self._handle_bytes, PBType.SUB: self._handle_sub_message, PBType.FLOATS: self._handle_packed_floats,
-      PBType.INT: self._handle_int64, PBType.INTS: self._handle_packed_int64s, PBType.STRING: self._handle_string, PBType.FLOAT: self._handle_float}
-    self.registered_handles = {}
-    for pb_name in PB_INFOS:
-      res = {}
-      for fid, config in PB_INFOS[pb_name].items():
-        parser_fn, repeated = None, False
-        if len(config) == 2: name, attr = config
-        elif len(config) == 3: name, attr, repeated = config
-        elif len(config) == 4: name, attr, repeated, parser_fn = config
-        handler_fn = self.attr_func_dict[attr]
-        def _wrapper_handler(obj, reader, wt, h=handler_fn, n=name, p=parser_fn, r=repeated): return h(obj, n, reader, wt, parser_func=p, repeated=r)
-        res[fid] = _wrapper_handler
-      self.registered_handles[pb_name] = res
-
-  def parse(self):
-    reader = BufferedReader(TensorIO(self.tensor))
-    return self._parse_message(reader, "ModelProto", lambda: {"opset_import": [], "domain": None, "graph": None})
-
-  def decode_varint(self, reader: BufferedReader) -> int:
-    result = 0
-    shift = 0
-    while True:
-      data = reader.read(1)
-      if data == b"": raise EOFError("decode_varint EOF")
-      result |= (data[0] & 0x7F) << shift
-      if not (data[0] & 0x80): return result
-      shift += 7
-      if shift >= 70: raise ValueError("Varint too long")
-
-  def skip_field_value(self, reader: BufferedReader, wire_type):
-    if wire_type == WIRETYPE_VARINT: self.decode_varint(reader)
-    elif wire_type == WIRETYPE_FIXED64: reader.seek(8, os.SEEK_CUR)
-    elif wire_type == WIRETYPE_FIXED32: reader.seek(4, os.SEEK_CUR)
-    elif wire_type == WIRETYPE_LENGTH_DELIMITED: reader.seek(self.decode_varint(reader), os.SEEK_CUR)
-    else: raise ValueError(f"Unknown wire type: {wire_type}")
-
-  def _parse_message(self, reader, message_field_handlers_name, initial_obj_factory=lambda: {}):
-    message_field_handlers = self.registered_handles[message_field_handlers_name]
-    obj = initial_obj_factory()
-    while True:
-      try:
-        tag_val = self.decode_varint(reader)
-        field_number = tag_val >> 3
-        wire_type = tag_val & 0x07
-        if handler := message_field_handlers.get(field_number):
-          handler(obj, reader, wire_type)
-        else: self.skip_field_value(reader, wire_type)
-      except EOFError: break
-    if message_field_handlers_name == "TensorProto" and self.load_external_data and obj.get("data_location", 0) == 1: self._parse_external_data(obj)
-    return obj
-
-  def _handle_delimited(self, reader:BufferedReader, use_tensor=False) -> Tensor|bytes:
-    str_len = self.decode_varint(reader)
-    if not use_tensor: return reader.read(str_len)
-    raw = reader.raw
-    assert isinstance(raw, TensorIO)
-    res = raw._tensor[reader.tell():(reader.tell()+str_len)]
-    reader.seek(str_len, os.SEEK_CUR)
-    return res
-
-  def _handle_string(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for string field '{key_name}'")
-    value = self._handle_delimited(reader)
-    assert isinstance(value, bytes)
-    gen_result(obj, key_name, value.decode("utf-8"), repeated)
-
-  def _handle_bytes(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for bytes field '{key_name}'")
-    value = self._handle_delimited(reader, use_tensor=True)
-    gen_result(obj, key_name, value, repeated)
-
-  def _handle_int64(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_VARINT: raise ValueError(f"Expected varint for int64 field '{key_name}'")
-    val = self.decode_varint(reader)
-    gen_result(obj, key_name, val - 2**64 if val & (1 << 63) else val, repeated)
-
-  def _handle_float(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_FIXED32: raise ValueError(f"Expected fixed32 for float field '{key_name}'")
-    val, = struct.unpack("<f", reader.read(4))
-    gen_result(obj, key_name, val, repeated)
-
-  def _handle_packed_int64s(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError("Packed int64s expected length_delimited")
-    total_bytes_len = self.decode_varint(reader)
-    old_pos = reader.tell()
-    values = []
-    while reader.tell() < total_bytes_len + old_pos:
-      val = self.decode_varint(reader) # need copy here because packed ints are varint
-      values.append(val - 2**64 if val & (1 << 63) else val)
-    obj[key_name] = values
-
-  def _handle_packed_floats(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError("Packed floats expected length_delimited")
-    value = self._handle_delimited(reader, use_tensor=True)
-    obj[key_name] = value
-
-  def _handle_sub_message(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
-    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for sub-message field '{key_name}'")
-    value = self._handle_delimited(reader, use_tensor=True)
-    assert isinstance(value, Tensor)
-    if isinstance(parser_func, str): sub_obj = self._parse_message(BufferedReader(TensorIO(value)), parser_func)
-    elif isinstance(parser_func, tuple): sub_obj = self._parse_message(BufferedReader(TensorIO(value)), parser_func[0], parser_func[1])
-    else: sub_obj = parser_func(BufferedReader(TensorIO(value)))
-    gen_result(obj, key_name, sub_obj, repeated)
-
-  def _parse_external_data(self, obj):
-    if "external_data" not in obj: raise ValueError("no external_data")
-    location = None
-    length = None
-    offset = 0
-    for kv in obj["external_data"]:
-      if kv["key"] == "location": location = kv["value"]
-      if kv["key"] == "offset": offset = int(kv["value"])
-      if kv["key"] == "length": length = int(kv["value"])
-    if location is None: raise ValueError("no location in external_data")
-    if self.file_path is None:
-      # get onnx file path from Tensor
-      if isinstance(self.tensor.device, str) and self.tensor.device.startswith("DISK:"):
-        self.file_path = pathlib.Path(self.tensor.device[5:])
-        if not (ext_path := self.file_path.parent.joinpath(location)).exists():
-          raise Exception(f"external location not exists: {ext_path}, may caused by symbolic link, try passing onnx file path to onnx_load")
-      else: raise Exception("onnx external_data need the origin file path, try passing onnx file path to onnx_load")
-    ext_path = self.file_path.parent.joinpath(location)
-    if not ext_path.exists(): raise Exception(f"external location not exists: {ext_path}")
-    ext_tensor = Tensor(ext_path)
-    obj["raw_data"] = ext_tensor[offset:offset+length] if length is not None else ext_tensor[offset:]
-    obj["data_location"] = 0
@@ -5,9 +5,9 @@ from tinygrad.nn import Linear
 from tinygrad.tensor import Tensor
 from tinygrad.nn.optim import Adam
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.opt.search import actions
+from tinygrad.codegen.opt.search import actions
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.helpers import getenv

 # stuff needed to unpack a kernel
@@ -17,7 +17,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.uop.ops import Variable
 inf, nan = float('inf'), float('nan')
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps

 INNER = 256
 class PolicyNet:
@@ -10,11 +10,11 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.uop.ops import Variable
 inf, nan = float('inf'), float('nan')
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps

 # more stuff
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.search import actions
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.search import actions
 from extra.optimization.helpers import lin_to_feats
 from extra.optimization.pretrain_valuenet import ValueNet
 from tinygrad.nn.optim import Adam
@@ -1,8 +1,8 @@
 import random
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.opt.search import actions
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import actions
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import tqdm

 tactions = set()
@@ -1,6 +1,6 @@
 # stuff needed to unpack a kernel
 from tinygrad import Variable
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -11,7 +11,7 @@ inf, nan = float('inf'), float('nan')
 UOps = Ops

 # kernel unpacker
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 def ast_str_to_ast(ast_str:str) -> UOp: return eval(ast_str)
 def ast_str_to_lin(ast_str:str, opts=None): return Kernel(ast_str_to_ast(ast_str), opts=opts)
 def kern_str_to_lin(kern_str:str, opts=None):
@@ -103,7 +103,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
  return ret

 from tinygrad.device import Device, Buffer
-from tinygrad.opt.search import _ensure_buffer_alloc, _time_program
+from tinygrad.codegen.opt.search import _ensure_buffer_alloc, _time_program
 from tinygrad.helpers import to_function_name, CACHELEVEL, diskcache_get, diskcache_put

 def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_global_size=65536, cnt=3, disable_cache=False, clear_l2=False) -> float:  # noqa: E501
@@ -1,4 +1,4 @@
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 from tqdm import tqdm, trange
 import math
 import random
@@ -14,7 +14,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.uop.ops import Variable
 inf, nan = float('inf'), float('nan')
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps

 from extra.optimization.helpers import lin_to_feats, MAX_DIMS

@@ -3,7 +3,7 @@ import numpy as np
 import math, random
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.opt.search import actions, bufs_from_lin, get_kernel_actions
+from tinygrad.codegen.opt.search import actions, bufs_from_lin, get_kernel_actions
 from tinygrad.nn.optim import Adam
 from extra.optimization.extract_policynet import PolicyNet
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
@@ -1,6 +1,6 @@
 from typing import List, Tuple
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.search import get_kernel_actions, actions
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.search import get_kernel_actions, actions

 _net = None
 def beam_q_estimate(beam:List[Tuple[Kernel, float]]) -> List[Tuple[Kernel, float]]:
@@ -4,8 +4,8 @@ from extra.optimization.helpers import ast_str_to_lin, time_linearizer
 from tinygrad import dtypes
 from tinygrad.helpers import BEAM, getenv
 from tinygrad.device import Device, Compiled
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.search import beam_search, bufs_from_lin
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.search import beam_search, bufs_from_lin


 if __name__ == '__main__':
@@ -6,8 +6,8 @@ from copy import deepcopy
 from tinygrad.helpers import getenv, colored
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.opt.search import bufs_from_lin, actions, get_kernel_actions
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import bufs_from_lin, actions, get_kernel_actions
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
 from extra.optimization.extract_policynet import PolicyNet
 from extra.optimization.pretrain_valuenet import ValueNet
@@ -1,5 +1,5 @@
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
-from tinygrad.opt.search import bufs_from_lin, get_kernel_actions
+from tinygrad.codegen.opt.search import bufs_from_lin, get_kernel_actions

 if __name__ == "__main__":
  ast_strs = load_worlds()
@@ -1,6 +1,6 @@
 import sys, pickle, decimal, json
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent
-from tinygrad.helpers import tqdm, temp, ProfileEvent, ProfileRangeEvent
+from tinygrad.helpers import tqdm, temp, ProfileEvent, ProfileRangeEvent, TracingKey

 devices:dict[str, tuple[decimal.Decimal, decimal.Decimal, int]] = {}
 def prep_ts(device:str, ts:decimal.Decimal, is_copy): return int(decimal.Decimal(ts) + devices[device][is_copy])
@@ -11,12 +11,14 @@ def dev_ev_to_perfetto_json(ev:ProfileDeviceEvent):
          {"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 0, "args": {"name": "COMPUTE"}},
          {"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 1, "args": {"name": "COPY"}}]
 def range_ev_to_perfetto_json(ev:ProfileRangeEvent):
-  return [{"name": ev.name, "ph": "X", "ts": prep_ts(ev.device, ev.st, ev.is_copy), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device, ev.is_copy)}]
+  name = ev.name.display_name if isinstance(ev.name, TracingKey) else ev.name
+  return [{"name": name, "ph": "X", "ts": prep_ts(ev.device, ev.st, ev.is_copy), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device, ev.is_copy)}]
 def graph_ev_to_perfetto_json(ev:ProfileGraphEvent, reccnt):
  ret = []
  for i,e in enumerate(ev.ents):
    st, en = ev.sigs[e.st_id], ev.sigs[e.en_id]
-    ret += [{"name": e.name, "ph": "X", "ts": prep_ts(e.device, st, e.is_copy), "dur": float(en-st), **dev_to_pid(e.device, e.is_copy)}]
+    name = e.name.display_name if isinstance(e.name, TracingKey) else e.name
+    ret += [{"name": name, "ph": "X", "ts": prep_ts(e.device, st, e.is_copy), "dur": float(en-st), **dev_to_pid(e.device, e.is_copy)}]
    for dep in ev.deps[i]:
      d = ev.ents[dep]
      ret += [{"ph": "s", **dev_to_pid(d.device, d.is_copy), "id": reccnt+len(ret), "ts": prep_ts(d.device, ev.sigs[d.en_id], d.is_copy), "bp": "e"}]
@@ -24,6 +26,8 @@ def graph_ev_to_perfetto_json(ev:ProfileGraphEvent, reccnt):
  return ret
 def to_perfetto(profile:list[ProfileEvent]):
  # Start json with devices.
+  profile += [ProfileDeviceEvent("TINY")]
+
  prof_json = [x for ev in profile if isinstance(ev, ProfileDeviceEvent) for x in dev_ev_to_perfetto_json(ev)]
  for ev in tqdm(profile, desc="preparing profile"):
    if isinstance(ev, ProfileRangeEvent): prof_json += range_ev_to_perfetto_json(ev)
@@ -6,8 +6,8 @@ from tinygrad.helpers import getenv, BEAM
 from tinygrad.engine.jit import TinyJit
 from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item, get_program
 from tinygrad.renderer import ProgramSpec
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 import numpy as np

 def move_jit_captured_to_dev(captured, device="DSP"):
@@ -128,6 +128,12 @@ def _linalg_eigh(self, UPLO: str = 'U'):
  w, v = torch.linalg.eigh(self.cpu(), UPLO=UPLO)
  return w.tiny(), v.tiny()

+@torch.library.impl("aten::_linalg_det", "privateuseone")
+# TODO: move to tinygrad
+def _linalg_det(self: torch.Tensor):
+  result = aten._linalg_det(self.cpu())
+  return result[0].tiny(), result[1].tiny(), result[2].tiny()
+
 def upsample_backward(grad_out, output_size, input_size, *args, f=None): return f(grad_out.cpu(), output_size, input_size, *args).tiny()

 for i in [
@@ -217,15 +223,18 @@ def max_unpool2d(self:torch.Tensor, indices:torch.Tensor, output_size):

@torch.library.impl("aten::arange", "privateuseone")
 def arange(end, dtype=None, device=None, pin_memory=None):
-  return wrap(Tensor.arange(0, end, dtype=_from_torch_dtype(dtype or torch.get_default_dtype())))
+  has_float = isinstance(end, float)
+  return wrap(Tensor.arange(0, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))

@torch.library.impl("aten::arange.start", "privateuseone")
 def arange_start(start, end, dtype=None, device=None, pin_memory=None):
-  return wrap(Tensor.arange(start, end, dtype=_from_torch_dtype(dtype or torch.get_default_dtype())))
+  has_float = any(isinstance(x, float) for x in (start, end))
+  return wrap(Tensor.arange(start, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))

@torch.library.impl("aten::arange.start_step", "privateuseone")
 def arange_start_step(start, end, step, dtype=None, device=None, pin_memory=None):
-  return wrap(Tensor.arange(start, end, step, dtype=_from_torch_dtype(dtype or torch.get_default_dtype())))
+  has_float = any(isinstance(x, float) for x in (start, end, step))
+  return wrap(Tensor.arange(start, end, step, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))

@torch.library.impl("aten::convolution_overrideable", "privateuseone")
 def convolution_overrideable(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups):
@@ -362,6 +371,7 @@ from torch._decomp import get_decompositions
 decomps = [
  aten.native_batch_norm, aten.native_batch_norm_backward,
  aten.native_layer_norm_backward,
+  aten.linalg_cross,
  aten.addmm,
  aten.addcmul,
  aten.addcdiv,
@@ -371,6 +381,7 @@ decomps = [
  aten.elu,  # elu has a scale + input_scale param
  aten.elu_backward,
  aten.softplus,
+  aten.logaddexp,
  aten.threshold,
  aten.nll_loss_forward,
  aten.nll_loss_backward,
@@ -135,7 +135,7 @@ class TestTorchBackend(unittest.TestCase):
    print(c.cpu())

  def test_maxpool2d_backward(self):
-    x = torch.arange(3*3, device=device).reshape(1, 1, 3, 3).requires_grad_(True)
+    x = torch.arange(3*3, dtype=torch.float32, device=device).reshape(1, 1, 3, 3).requires_grad_(True)
    torch.nn.functional.max_pool2d(x, kernel_size=2, stride=1).sum().backward()
    np.testing.assert_equal(x.grad.squeeze().cpu().numpy(), [[0, 0, 0], [0, 1, 1], [0, 1, 1]])

@@ -198,6 +198,17 @@ class TestTorchBackend(unittest.TestCase):
    recon = (v @ torch.diag(w) @ v.T).cpu().numpy()
    np.testing.assert_allclose(recon, a.cpu().numpy(), atol=1e-6)

+  def test_linalg_det(self):
+    a = torch.diag(torch.tensor([1,2,3,4,5], dtype = torch.float32, device=device))
+    b = torch.linalg.det(a)
+    np.testing.assert_equal(b.cpu().numpy(), 120.0)
+
+  def test_linalg_cross(self):
+    a = torch.tensor([[1, 0, 0], [0, 1, 0]], dtype=torch.float32, device=device)
+    b = torch.tensor([[0, 0, 1]], dtype=torch.float32, device=device)
+    cross = torch.linalg.cross(a, b)
+    np.testing.assert_equal(cross.cpu().numpy(), np.array([[0, -1, 0], [1, 0, 0]], dtype=np.float32))
+
  def test_scalar_assign(self):
    a = torch.tensor([1, 2, 3], device=device)
    a[1] = 4
@@ -0,0 +1,75 @@
+import torch
+
+#credit to KellerJordan at https://github.com/KellerJordan/Muon/tree/master
+#some changes: classic momentum instead of weighting gradient
+#added ns_steps, ns_params, nesterov as hyperparams
+def zeropower_via_newtonschulz5(G:torch.tensor, steps:int, params:tuple[int, ...]):
+  """
+  Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
+  quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
+  of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
+  zero even beyond the point where the iteration no longer converges all the way to one everywhere
+  on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
+  where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
+  performance at all relative to UV^T, where USV^T = G is the SVD.
+  """
+  assert G.ndim >= 2 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
+
+  a, b, c = params
+  X = G
+  if G.size(-2) > G.size(-1):
+    X = X.mT
+
+  # Ensure spectral norm is at most 1
+  X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
+  # Perform the NS iterations
+  for _ in range(steps):
+    A = X @ X.mT
+    B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
+    X = a * X + B @ X
+
+  if G.size(-2) > G.size(-1):
+    X = X.mT
+
+  return X
+
+def muon_update(grad, momentum, beta=0.95, ns_steps=5, ns_params=(3.4445, -4.7750,  2.0315), nesterov=True):
+  if beta:
+    momentum.mul_(beta).add_(grad)
+    update = grad.add(momentum,alpha=beta) if nesterov else momentum
+  else: update = grad
+  if update.ndim == 4: # for the case of conv filters
+    update = update.view(len(update), -1)
+  update = zeropower_via_newtonschulz5(update, steps=ns_steps, params=ns_params)
+  return update
+
+class SingleDeviceMuon(torch.optim.Optimizer):
+  """
+  Muon variant for usage in non-distributed settings.
+  """
+  def __init__(self, params, lr=0.02, weight_decay=0.0, momentum=0.95, ns_steps=5, ns_params=(3.4445, -4.7750,  2.0315), nesterov=True):
+    defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, ns_steps=ns_steps, ns_params=ns_params, nesterov=nesterov)
+    super().__init__(params, defaults)
+
+  @torch.no_grad()
+  def step(self, closure=None):
+
+    loss = None
+    if closure is not None:
+      with torch.enable_grad():
+        loss = closure()
+
+    for group in self.param_groups:
+      for p in group["params"]:
+        if p.grad is None:
+          p.grad = torch.zeros_like(p)  # Force synchronization
+        state = self.state[p]
+        if len(state) == 0:
+          state["momentum_buffer"] = torch.zeros_like(p)
+        update = muon_update(p.grad, state["momentum_buffer"], beta=group["momentum"], ns_steps=group["ns_steps"],
+                             ns_params=group["ns_params"], nesterov=group["nesterov"])
+        p.mul_(1.0 - group["lr"] * group["weight_decay"])
+
+        p.add_(update.reshape(p.shape), alpha=-group["lr"])
+
+    return loss
@@ -0,0 +1,2 @@
+[pytest]
+norecursedirs = extra
@@ -35,6 +35,7 @@ lint.select = [
 line-length = 150

 exclude = [
+  ".git/",
  "docs/",
  "extra/",
  "tinygrad/runtime/autogen",
@@ -9,7 +9,7 @@ with open(directory / 'README.md', encoding='utf-8') as f:

 testing_minimal = [
  "numpy",
-  "torch",
+  "torch==2.7.1",
  "pytest",
  "pytest-xdist",
  "hypothesis",
@@ -18,16 +18,35 @@ testing_minimal = [
 ]

 setup(name='tinygrad',
-      version='0.10.3',
+      version='0.11.0',
      description='You like pytorch? You like micrograd? You love tinygrad! <3',
      author='George Hotz',
      license='MIT',
      long_description=long_description,
      long_description_content_type='text/markdown',
-      packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.runtime.autogen.am', 'tinygrad.codegen', 'tinygrad.nn',
-                  'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.schedule',
-                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt',
-                  'tinygrad.runtime.support.nv', 'tinygrad.apps'],
+      packages = [
+        'tinygrad',
+        'tinygrad.apps',
+        'tinygrad.codegen',
+        'tinygrad.codegen.opt',
+        'tinygrad.codegen.late',
+        'tinygrad.engine',
+        'tinygrad.frontend',
+        'tinygrad.nn',
+        'tinygrad.renderer',
+        'tinygrad.runtime',
+        'tinygrad.runtime.autogen',
+        'tinygrad.runtime.autogen.am',
+        'tinygrad.runtime.autogen.nv',
+        'tinygrad.runtime.graph',
+        'tinygrad.runtime.support',
+        'tinygrad.runtime.support.am',
+        'tinygrad.runtime.support.nv',
+        'tinygrad.schedule',
+        'tinygrad.shape',
+        'tinygrad.uop',
+        'tinygrad.viz',
+      ],
      package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
      classifiers=[
        "Programming Language :: Python :: 3",
@@ -45,6 +64,7 @@ setup(name='tinygrad',
            "pre-commit",
            "ruff",
            "numpy",
+            "typeguard",
        ],
        #'mlperf': ["mlperf-logging @ git+https://github.com/mlperf/[email protected]"],
        'testing_minimal': testing_minimal,
@@ -6,7 +6,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps
 from tinygrad import Variable

 MOCKGPU = getenv("MOCKGPU")
@@ -163,10 +163,8 @@ class TestHCQ(unittest.TestCase):
    a = Tensor.randint((3, 3, 3), dtype=dtypes.int, device=Device.DEFAULT).realize()
    b = a + 1
    si = b.schedule()[-1]
-    k = Kernel(si.ast, opts=TestHCQ.d0.renderer)
-    for i in range(3): k.apply_opt(Opt(op=OptOps.LOCAL, axis=0, arg=3))

-    runner = CompiledRunner(get_program(k.get_optimized_ast(), k.opts))
+    runner = CompiledRunner(get_program(si.ast, TestHCQ.d0.renderer, opts=[Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]))

    zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
    zt = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
@@ -338,7 +336,7 @@ class TestHCQ(unittest.TestCase):
    et = float(sig_en.timestamp - sig_st.timestamp)

    print(f"exec kernel time: {et:.2f} us")
-    assert 0.1 <= et <= (15000 if MOCKGPU or Device.DEFAULT in {"CPU", "LLVM"} else 100)
+    assert 0.1 <= et <= (100000 if MOCKGPU or Device.DEFAULT in {"CPU", "LLVM"} else 100)

  def test_speed_copy_bandwidth(self):
    if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -1,7 +1,7 @@
 import random
 from tinygrad.helpers import getenv
-from tinygrad.opt.search import beam_search, bufs_from_lin
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer

 def optimize_kernel(k):
@@ -2,7 +2,7 @@ import time
 from tinygrad import Tensor, TinyJit, Device, Context
 from tinygrad.helpers import Profiling, Timing, GlobalCounters

-# python3 test/test_speed_v_torch.py TestSpeed.test_add_a
+# python3 test/speed/external_test_speed_v_torch.py TestSpeed.test_add_a

@TinyJit
 def plus(a:Tensor, b:Tensor): return a+b
@@ -1,12 +1,8 @@
-from typing import List
 from extra.models.resnet import ResNet50
-from tinygrad import Tensor, nn
-from tinygrad.helpers import Profiling, Timing, getenv, BEAM, NOOPT, DEBUG, Context, ansilen
+from tinygrad import Tensor, nn, Device
+from tinygrad.helpers import Profiling, Timing, getenv
 from tinygrad.uop.ops import Ops
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites, rewrites_for_linearizer
-from tinygrad.opt.search import beam_search, bufs_from_lin
 from tinygrad.uop.spec import type_verify

 if __name__ == "__main__":
@@ -31,26 +27,13 @@ if __name__ == "__main__":
      if not SCHEDULE_ONLY:
        asts = list({x.ast.key:x.ast for x in sched if x.ast.op is Ops.SINK}.values())
        if (restrict_kernel := getenv("RESTRICT_KERNEL", -1)) != -1: asts = asts[restrict_kernel:restrict_kernel+1]
-        kernels: List[Kernel] = []
-        with Timing(f"***** model opts({len(asts):2d}) in  "):
-          with Profiling(PROFILE >= 3):
-            for ast in asts:
-              k = Kernel(ast)
-              if BEAM:
-                with Context(DEBUG=max(2, DEBUG.value)): k = beam_search(k, bufs_from_lin(k), BEAM.value)
-              elif NOOPT: pass
-              else: k.apply_opts(hand_coded_optimizations(k))
-              kernels.append(k)
-
-        with Timing("***** model prep in      "):
-          kernels = [(k, k.get_optimized_ast(), get_rewrites_for_renderer(k.opts, linearizer=False)) for k in kernels]

+        rewrites = get_rewrites_for_renderer(Device.default.renderer, linearizer=False)
        with Profiling(PROFILE, fn="/tmp/rewrite.prof"):
          with Timing("***** model rewrite in   "):
            rewritten_uops = []
-            for i,(k,u,rewrites) in enumerate(kernels):
-              with Timing(f"rewrite {i:2d} {k.name}{' '*(50-ansilen(k.name))}", enabled=getenv("VERBOSE", 0)):
-                rewritten_uops.append(apply_rewrites(u, rewrites))
+            for u in asts:
+              rewritten_uops.append(apply_rewrites(u, rewrites))

        if LINEARIZE:
          with Timing("***** model linearize in "):
@@ -8,7 +8,7 @@ if __name__ == "__main__":
  GlobalCounters.reset()
  t.softmax(-1, dtype="half", _single_kernel=True).realize()

-  from tinygrad.opt.kernel import Kernel, Opt, OptOps
+  from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
  from tinygrad.helpers import get_single_element
  GlobalCounters.reset()
  si = get_single_element(t.softmax(-1, dtype="half", _single_kernel=True).schedule())
@@ -1,8 +1,8 @@
 # ruff: noqa: E501
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.opt.search import bufs_from_lin
+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
@@ -10,7 +10,7 @@ if __name__ == "__main__":

  model, kv = Transformer.from_gguf(Tensor.from_url(models["1B"]), max_context=4096)

-  tok = SimpleTokenizer(kv["tokenizer.ggml.tokens"])
+  tok = SimpleTokenizer.from_gguf_kv(kv)
  bos_id: int = kv['tokenizer.ggml.bos_token_id']
  eos_id: int = kv['tokenizer.ggml.eos_token_id']

@@ -4,10 +4,10 @@ os.environ["VALIDATE_HCQ"]="1"

 import unittest, random
 import numpy as np
-from tinygrad.opt.kernel import Kernel, KernelOptError
+from tinygrad.codegen.opt.kernel import Kernel, KernelOptError
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.opt.search import Opt, OptOps
+from tinygrad.codegen.opt.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from test.external.fuzz_linearizer import compare_linearizer, compare_states, get_fuzz_rawbuf_like

@@ -3,7 +3,7 @@ from tinygrad.runtime.support.hip_comgr import compile_hip
 from tinygrad import Tensor
 from tinygrad.device import Device
 from tinygrad.engine.schedule import create_schedule
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel

 class TestHIPCompileSpeed(unittest.TestCase):
  @unittest.skipIf(Device.DEFAULT != "HIP", "only run on HIP")
@@ -2,11 +2,11 @@ import unittest, struct, array, ctypes
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.helpers import to_mv
 from tinygrad.runtime.ops_nv import NVDevice, HWQueue
-from tinygrad.opt.search import Opt, OptOps
+from tinygrad.codegen.opt.search import Opt, OptOps
 from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
 from test.external.fuzz_linearizer import get_fuzz_rawbufs

-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.uop.ops import LazyOp, Ops, ReduceOps, BufferOps, MemBuffer
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
@@ -53,6 +53,7 @@ backend_test.exclude('test_dynamicquantizelinear_cpu')
 backend_test.exclude('test_dynamicquantizelinear_expanded_cpu')

 # BUG: ORT fails these with numerical error but we match ORT numerically
+# see: https://onnx.ai/backend-scoreboard/onnxruntime_details_stable.html
 # tested in external_test_onnx_ops.py::TestMainOnnxOps.test_qlinearmatmul_2D_int8_float16
 backend_test.exclude('test_qlinearmatmul_2D_int8_float16_cpu')
 # tested in external_test_onnx_ops.py::TestMainOnnxOps.test_qlinearmatmul_3D_int8_float16
@@ -65,6 +66,10 @@ backend_test.exclude('test_qlinearmatmul_3D_int8_float32_cpu')
 backend_test.exclude('test_maxunpool_export_with_output_shape_cpu')
 # tested in external_test_onnx_ops.py::TestMainOnnxOps.test_averagepool_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_True
 backend_test.exclude('test_averagepool_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_True_cpu')
+# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_resize_downsample_scales_linear_align_corners
+backend_test.exclude('test_resize_downsample_scales_linear_align_corners_cpu')
+# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_resize_downsample_scales_cubic_align_corners
+backend_test.exclude('test_resize_downsample_scales_cubic_align_corners_cpu')

 # about different dtypes
 if not is_dtype_supported(dtypes.float64):
@@ -165,10 +170,6 @@ backend_test.exclude('test_deform_conv_*')
 backend_test.exclude('test_lppool_*')
 backend_test.exclude('test_scan_*')
 backend_test.exclude('test_split_to_sequence_*')
-backend_test.exclude('test_resize_downsample_scales_cubic_*') # unsure how to implement cubic
-backend_test.exclude('test_resize_downsample_sizes_cubic_*') # unsure how to implement cubic
-backend_test.exclude('test_resize_upsample_scales_cubic_*') # unsure how to implement cubic
-backend_test.exclude('test_resize_upsample_sizes_cubic_*') # unsure how to implement cubic
 backend_test.exclude('test_ai_onnx_ml_tree_ensemble_*') # https://github.com/onnx/onnx/blob/main/onnx/reference/ops/aionnxml/op_tree_ensemble.py#L121

 # rest of the failing tests
@@ -178,6 +179,8 @@ backend_test.exclude('test_resize_tf_crop_and_resize_axes_3_2_cpu') # tf_crop_an
 backend_test.exclude('test_resize_tf_crop_and_resize_extrapolation_value_cpu') # tf_crop_and_resize value not implemented
 backend_test.exclude('test_resize_downsample_scales_linear_antialias_cpu') # antialias not implemented
 backend_test.exclude('test_resize_downsample_sizes_linear_antialias_cpu') # antialias not implemented
+backend_test.exclude('test_resize_downsample_scales_cubic_antialias_cpu') # antialias not implemented
+backend_test.exclude('test_resize_downsample_sizes_cubic_antialias_cpu') # antialias not implemented
 backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_value_only_mapping_cpu') # bad data type string
 backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad data type string

@@ -75,6 +75,49 @@ class TestMainOnnxOps(TestOnnxOps):
    outputs = ["y"]
    self.helper_test_single_op("Gather", inputs, attributes, outputs)

+  # NOTE: resize OP is sensitive to numerical errors
+  def _test_resize_scales(self, scale_values, **kwargs):
+    for sc in scale_values:
+      for ct_mode in ["half_pixel", "align_corners", "asymmetric", "pytorch_half_pixel", "half_pixel_symmetric"]:
+        with self.subTest(coordinate_transformation_mode=ct_mode, scale=sc, **kwargs):
+          X = np.array([[[[1, 2, 3, 4],
+                          [5, 6, 7, 8],
+                          [9,10,11,12]]]], dtype=np.float32)
+          scales = np.array([1.0, 1.0, sc, sc], dtype=np.float32)
+          inputs = {"X": X, "roi": np.array([], dtype=np.float32), "scales": scales}
+          attributes = {"coordinate_transformation_mode": ct_mode, **kwargs}
+          outputs = ["out"]
+          self.helper_test_single_op("Resize", inputs, attributes, outputs)
+
+  def test_resize_linear_mode(self):
+    self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="linear")
+
+  def test_resize_nearest_mode(self):
+    # excluded 3.5 because some values divide into slight numerical differences, which when rounded gives wrong results
+    self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 20.0], mode="nearest")
+
+  def test_resize_cubic_mode(self):
+    self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="cubic", exclude_outside=1)
+    self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="cubic", exclude_outside=0)
+
+  def test_resize_downsample_scales_linear_align_corners(self):
+    # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-131
+    X = np.array([[[[1, 2, 3, 4], [5, 6, 7, 8]]]], dtype=np.float32)
+    scales = np.array([1.0, 1.0, 0.6, 0.6], dtype=np.float32)
+    inputs = {"X": X, "roi": np.array([], dtype=np.float32), "scales": scales}
+    attributes = {"mode": "linear", "coordinate_transformation_mode": "align_corners"}
+    outputs = ["out"]
+    self.helper_test_single_op("Resize", inputs, attributes, outputs)
+
+  def test_resize_downsample_scales_cubic_align_corners(self):
+    # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-131
+    X = np.array([[[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]]], dtype=np.float32)
+    scales = np.array([1.0, 1.0, 0.8, 0.8], dtype=np.float32)
+    inputs = {"X": X, "roi": np.array([], dtype=np.float32), "scales": scales}
+    attributes = {"mode": "cubic", "coordinate_transformation_mode": "align_corners"}
+    outputs = ["out"]
+    self.helper_test_single_op("Resize", inputs, attributes, outputs)
+
  def test_maxunpool_export_with_output_shape(self):
    # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-91
    xT = np.array([[[[5, 6], [7, 8]]]], dtype=np.float32)
@@ -251,7 +294,6 @@ class TestTrainingOnnxOps(TestOnnxOps):
        outputs = ["X_out", "V_out"]
        self._validate_training("Momentum", onnx_fxn, inputs, attributes, outputs)

-  @unittest.expectedFailure  # TODO: regression from removing StrEnum in Domain
  def test_adam_t_greater_than_zero(self):
    from onnx.backend.test.case.node.adam import apply_adam
    for t in [1, 3, 100]:
@@ -3,8 +3,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.uop.ops import Ops
 from tinygrad.device import is_dtype_supported
-from extra.onnx import data_types
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner, OnnxDataType
 from hypothesis import given, strategies as st

 # copied from test_const_folding.py
@@ -86,8 +85,8 @@ class TestOnnxRunner(unittest.TestCase):
      output = runner({'inp': Tensor([1])})['output']
      np.testing.assert_equal(output.numpy(), weights + 1)

-all_dtypes = list(data_types.keys())
-device_supported_dtypes = {odt for odt, dtype in data_types.items() if is_dtype_supported(dtype)}
+all_dtypes = list(OnnxDataType)
+device_supported_dtypes = {odt for odt in OnnxDataType if is_dtype_supported(odt.to_dtype())}

 class TestOnnxRunnerDtypes(unittest.TestCase):
  """
@@ -95,7 +94,7 @@ class TestOnnxRunnerDtypes(unittest.TestCase):
  External tensors (inputs) preserve their original dtype - user must ensure compatibility with device.
  """
  def _get_expected_dtype(self, onnx_dtype: int, is_input: bool):
-    true_dtype = data_types[onnx_dtype]
+    true_dtype = OnnxDataType(onnx_dtype).to_dtype()
    # inputs always preserve their true dtype.
    if is_input:
      return true_dtype
@@ -1,17 +1,19 @@
 from transformers import AutoTokenizer
 from datasets import load_dataset
-from tinygrad.apps.llm import SimpleTokenizer
-from tinygrad.helpers import tqdm, getenv
+from tinygrad.apps.llm import SimpleTokenizer, gpt2_decode_vocab, get_llama_re
+from tinygrad.helpers import tqdm, getenv, partition

 # use ALLOW_FAILED=-1 to go over the entire dataset without printing.
 if __name__ == "__main__":
  base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
-  vocab_words = [ word for word, _ in sorted(base_tokenizer.get_vocab().items(), key=lambda t: t[1]) ]
+  special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()),
+    lambda e: e[1] in base_tokenizer.all_special_ids)
  inv_vocab = { tid: word for word, tid in base_tokenizer.get_vocab().items() }
-  simple_tokenizer = SimpleTokenizer(vocab_words)
+  simple_tokenizer = SimpleTokenizer(get_llama_re(), gpt2_decode_vocab(dict(normal_tokens)), dict(special_tokens))

  color_codes = [ 91, 92, 94, 93, 95 ]
-  def color_tokens(tids): return "".join(f"\033[{color_codes[i%len(color_codes)]}m{inv_vocab[t]}" for i, t in enumerate(tids)) + "\033[0m"
+  def color_tokens(tids):
+    return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([t])}" for i, t in enumerate(tids)) + "\033[0m"

  ds = load_dataset("OpenAssistant/oasst1")
  allow_failed = getenv("ALLOW_FAILED", 10)
@@ -2,11 +2,11 @@
 import unittest

 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.opt.search import Opt, OptOps
+from tinygrad.codegen.opt.search import Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel

 from test.external.fuzz_linearizer import run_linearizer

@@ -1,33 +1,25 @@
 import random
-from z3 import Int, Solver, sat
-from tinygrad import dtypes, Device
-from tinygrad.uop.ops import UOp, Ops, UPat, graph_rewrite, PatternMatcher
-from tinygrad.codegen.optional import fast_idiv
+import z3
+from tinygrad import dtypes
+from tinygrad.uop.spec import uops_to_z3, z3_cdiv
+from tinygrad.uop.ops import UOp
+from tinygrad.uop.decompositions import fast_idiv
 random.seed(42)

-z3_renderer = PatternMatcher([
-  (UPat((Ops.DEFINE_VAR, Ops.SPECIAL), name="x"), lambda x: UOp(Ops.NOOP, arg=x.arg[0])),
-  # Because fast_idiv only works for non-negative integers we can emulate machine arithmetic with modulo operations.
-  (UPat(Ops.SHR, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"(({x.src[0].arg}/(2**{x.src[1].arg}))%{dtypes.max(x.dtype)+1})")),
-  (UPat(Ops.MUL, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"(({x.src[0].arg}*{x.src[1].arg})%{dtypes.max(x.dtype)+1})")),
-  (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: UOp(Ops.NOOP, arg=str(x.arg))),
-  (UPat(Ops.CAST, src=UPat(Ops.NOOP), name="x"), lambda x:  UOp(Ops.NOOP, arg=f"{x.src[0].arg}")),
-])
-
-def render(self) -> str:
-  ret = graph_rewrite(self.simplify(), z3_renderer)
-  return ret.arg if ret.op is Ops.NOOP else str(ret)
-
+powers_of_two = [2**i for i in range(64)]
 if __name__ == "__main__":
-  x = Int('x')
-  for _ in range(10_000):
+  for i in range(10_000):
+    if i % 1000 == 0:
+      print(f"Progress: {i}")
    dt = random.choice(dtypes.ints)
-    u = UOp(Ops.DEFINE_VAR, dt, arg=('x', 0, random.randint(1, dtypes.max(dt))), src=())
+    u = UOp.variable('x', random.randint(dt.min, 0), random.randint(1, dt.max), dtype=dt)
    d = random.randint(1, max(1, u.arg[2]))
-
-    expr = fast_idiv(Device[Device.DEFAULT].renderer, u, d)
+    if d in powers_of_two: continue
+    expr = fast_idiv(None, u, d)
    if expr is None: continue
-    solver = Solver()
-    solver.add(x>=u.arg[1], x<=u.arg[2])
-    if solver.check(eval(render(expr)) != x/d) == sat:
-      assert False, f"Failed: {render(expr)} != x//{d} at x={solver.model()[x]}\nx={u}\nd={d}"
+
+    solver = z3.Solver()
+    z3_expr, x =uops_to_z3(solver, expr, u)
+
+    if solver.check(z3_expr != z3_cdiv(x, d)) == z3.sat:
+      assert False, f"Failed: {expr.render()} != x//{d} at x={solver.model()}\nx={u}\nd={d}\n{z3_expr=}\n{x/d=}"
@@ -21,9 +21,9 @@ if os.getenv("VALIDATE_HCQ", 0) != 0:

 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.kernel import Opt, OptOps
-from tinygrad.opt.search import get_kernel_actions, bufs_from_lin
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.search import get_kernel_actions, bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing
 from tinygrad.uop.ops import UOp, Ops
@@ -1,8 +1,8 @@
 import random, operator
 import z3
 from tinygrad import Variable, dtypes
-from tinygrad.uop.ops import UOp, graph_rewrite
-from tinygrad.uop.spec import z3_renderer
+from tinygrad.uop.ops import UOp
+from tinygrad.uop.spec import uops_to_z3
 from tinygrad.helpers import DEBUG, Context

 seed = random.randint(0, 100)
@@ -57,8 +57,7 @@ if __name__ == "__main__":

    solver = z3.Solver()
    solver.set(timeout=5000)  # some expressions take very long verify, but its very unlikely they actually return sat
-    z3_sink = graph_rewrite(expr.sink(simplified_expr, u1, u2, u3), z3_renderer, ctx=(solver, {}))
-    z3_expr, z3_simplified_expr = z3_sink.src[0].arg, z3_sink.src[1].arg
+    z3_expr, z3_simplified_expr, v1, v2, v3 = uops_to_z3(solver, expr, simplified_expr, u1, u2, u3)
    check = solver.check(z3_simplified_expr != z3_expr)
    if check == z3.unknown and DEBUG>=1:
      skipped += 1
@@ -69,7 +68,6 @@ if __name__ == "__main__":
            f"expr = {expr.render(simplify=False)}\n")
    elif check == z3.sat:
      m = solver.model()
-      v1, v2, v3 = z3_sink.src[2].arg, z3_sink.src[3].arg, z3_sink.src[4].arg
      n1, n2, n3 = m[v1], m[v2], m[v3]
      u1_val, u2_val, u3_val = u1.const_like(n1.as_long()), u2.const_like(n2.as_long()), u3.const_like(n3.as_long())
      with Context(CORRECT_DIVMOD_FOLDING=1):
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 # compare kernels created by HEAD against master
 import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools, functools, base64, codecs
+from dataclasses import replace
 from typing import Callable, Any

 ASSERT_DIFF = int((flag:="[pr]") in os.getenv("COMMIT_MESSAGE", flag) or flag in os.getenv("PR_TITLE", flag))
@@ -11,7 +12,9 @@ try:
  from tinygrad.renderer import Renderer, ProgramSpec
  from tinygrad.engine.realize import get_program
  from tinygrad.uop.ops import UOp, Ops, KernelInfo
+  from tinygrad.codegen.opt.kernel import Opt
  from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
+  from tinygrad.device import Device
 except ImportError as e:
  print(repr(e))
  exit(int(ASSERT_DIFF))
@@ -47,9 +50,13 @@ def replay_kernelize(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[
    return "\n".join([f"{len(asts)} kernels", *asts])
  return to_str(new_sink), to_str(ret[big_sink]), (big_sink,)

-def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
-  input_ast = ast.replace(arg=KernelInfo(opts_to_apply=p.applied_opts, name=p.name)) if ast.arg is None else ast
-  p2 = get_program(input_ast, renderer)
+def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
+  # NOTE: this always uses the opts_to_apply path
+  sink_arg = ast.arg or KernelInfo(opts_to_apply=p.applied_opts)
+  input_ast = ast.replace(arg=replace(sink_arg, name=p.name))
+  # if no renderer was provided, open the device to get it
+  if renderer is None: renderer = Device[p.device].renderer
+  p2 = get_program(input_ast, renderer=renderer)
  def to_str(ret:ProgramSpec) -> str:
    # PYTHON renderer pickles UOps, first unpickle and decode here
    if p.device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(ret.src))])
@@ -74,6 +81,7 @@ def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
      warnings.warn(f"detected changes in over {MAX_DIFF_PCT}%. skipping further diff generation.", ProcessReplayWarning)
      early_stop.set()
      break
+    name, loc = "", ""
    try:
      name, args, kwargs, ctx_vals, loc, ret = pickle.loads(row[0])
      ctx_vars = {k:v.value for k,v in ctx_vals.items() if k != "DEBUG" and (var:=ContextVar._cache.get(k)) is not None and var.value != v.value}
@@ -90,8 +98,7 @@ def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
        warnings.warn("PROCESS REPLAY DETECTED CHANGE", ProcessReplayWarning)
    except Exception as e:
      changed += 1
-      warnings.warn(e, ProcessReplayWarning)
-  conn.commit()
+      warnings.warn(f"{name=} {loc=} {e=}", ProcessReplayWarning)
  cur.close()

 # *** generic runner to map rows of a table to a function in parallel
@@ -103,12 +110,11 @@ def _pmap(fxns:dict[str, Callable]) -> None:
  except sqlite3.OperationalError:
    raise RuntimeError(f"{TABLE_NAME} isn't accessible in master, did DB_VERSION change?")
  finally:
-    conn.commit()
    cur.close()

  with multiprocessing.get_context("spawn").Pool(multiprocessing.cpu_count()) as pool:
-    inputs = list(range(0, row_count, PAGE_SIZE))
-    list(tqdm(pool.imap_unordered(functools.partial(diff, fxns=fxns), inputs), total=len(inputs)))
+    bar = tqdm(total=row_count)
+    for _ in pool.imap_unordered(functools.partial(diff, fxns=fxns), range(0, row_count, PAGE_SIZE)): bar.update(PAGE_SIZE)
    pool.close()
    pool.join()
    pool.terminate()
@@ -123,5 +129,5 @@ if __name__ == "__main__":
  logging.info(f"running process replay with {ASSERT_DIFF=}")
  try: _pmap(replayers)
  except Exception as e:
-    logging.info("process replay err", e)
+    logging.info(f"process replay err: {e}")
    exit(int(ASSERT_DIFF))
@@ -1,7 +1,7 @@
 from tinygrad import Device
 from tinygrad.helpers import getenv, DEBUG, BEAM
-from tinygrad.opt.search import beam_search, bufs_from_lin
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer

 if __name__ == "__main__":
@@ -2,8 +2,8 @@ from tinygrad import Device, dtypes
 from tinygrad.helpers import getenv, colorize_float, DEBUG
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
-from tinygrad.opt.heuristic import hand_coded_optimizations
-from tinygrad.opt.search import bufs_from_lin
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.runtime.ops_amd import AMDDevice
@@ -2,8 +2,8 @@ from tinygrad import Device, dtypes
 from tinygrad.helpers import getenv, colorize_float
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
-from tinygrad.opt.heuristic import hand_coded_optimizations
-from tinygrad.opt.search import bufs_from_lin
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.tensor import _to_np_dtype
 import numpy as np
@@ -1,10 +1,10 @@
 import itertools
 from tinygrad import Device
 from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import getenv, colorize_float
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.opt.search import bufs_from_lin
+from tinygrad.codegen.opt.search import bufs_from_lin
 from tinygrad.runtime.ops_cuda import PTXCompiler, PTXRenderer, CUDACompiler

 if __name__ == "__main__":
@@ -3,7 +3,7 @@ from collections import defaultdict
 from extra.optimization.helpers import kern_str_to_lin, time_linearizer
 from test.external.fuzz_linearizer import compare_linearizer
 from tinygrad.helpers import colored
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel

 # Use this with the LOGKERNS options to verify that all executed kernels are valid and evaluate to the same ground truth results

@@ -87,16 +87,19 @@ class AMDDriver(VirtDriver):
        functools.partial(TextFileDesc, text=gpu_props.format(drm_render_minor=gpu_id))),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0',
               functools.partial(DirFileDesc, child_names=[str(am.GC_HWID), str(am.SDMA0_HWID), str(am.NBIF_HWID)])),
+      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/major', functools.partial(TextFileDesc, text='11')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/minor', functools.partial(TextFileDesc, text='0')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/revision', functools.partial(TextFileDesc, text='0')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/base_addr',
               functools.partial(TextFileDesc, text='0x00001260\n0x0000A000\n0x0001C000\n0x02402C00')),
+      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/major', functools.partial(TextFileDesc, text='6')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/minor', functools.partial(TextFileDesc, text='0')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/revision', functools.partial(TextFileDesc, text='0')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/base_addr',
               functools.partial(TextFileDesc, text='0x00001260\n0x0000A000\n0x0001C000\n0x02402C00')),
+      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/major', functools.partial(TextFileDesc, text='4')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/minor', functools.partial(TextFileDesc, text='3')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/revision', functools.partial(TextFileDesc, text='0')),
@@ -1,7 +1,7 @@
 import ctypes, time
 from test.mockgpu.gpu import VirtGPU
 from tinygrad.helpers import getbits, to_mv, init_c_struct_t
-import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4, tinygrad.runtime.autogen.am.soc21 as soc21
+import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4

 SDMA_MAX_COPY_SIZE = 0x400000

@@ -14,6 +14,9 @@ regSQ_THREAD_TRACE_BUF0_SIZE = 0x39e9 + amd_gpu.GC_BASE__INST0_SEG1
 regSQ_THREAD_TRACE_WPTR = 0x39ef + amd_gpu.GC_BASE__INST0_SEG1
 regSQ_THREAD_TRACE_STATUS = 0x39f4 + amd_gpu.GC_BASE__INST0_SEG1

+class SQTT_EVENTS:
+  THREAD_TRACE_FINISH = 0x00000037
+
 CACHE_FLUSH_AND_INV_TS_EVENT = 0x14

 WAIT_REG_MEM_FUNCTION_ALWAYS = 0
@@ -208,7 +211,7 @@ class PM4Executor(AMDQueue):
    assert n == 0
    event_dw = self._next_dword()
    match (event_dw & 0xFF): # event type
-      case soc21.THREAD_TRACE_FINISH:
+      case SQTT_EVENTS.THREAD_TRACE_FINISH:
        old_idx = self.gpu.regs.grbm_index
        for se in range(self.gpu.regs.n_se):
          self.gpu.regs.grbm_index = 0b011 << 29 | se << 16 # select se, broadcast sa and instance
--- a/Show More
+++ b/Show More