fix unit test dtypes

.github/actions/setup-tinygrad/action.yml

@@ -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 }}-${{ env.APT_CACHE_VERSION }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}
 
     - 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'
@@ -225,20 +225,16 @@ runs:
     - name: Install gpuocelot dependencies (MacOS)
       if: inputs.ocelot == 'true' && runner.os == 'macOS'
       shell: bash
-      run: |
-        pkgs=(cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses)
-        for f in "${pkgs[@]}"; do
-          brew ls --versions "$f" >/dev/null 2>&1 || brew install --quiet "$f"
-        done
+      run: brew install --quiet cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
     - name: Cache gpuocelot
       if: inputs.ocelot == 'true'
       id: cache-build
       uses: actions/cache@v4
       env:
-        cache-name: cache-gpuocelot-build-1
+        cache-name: cache-gpuocelot-build
       with:
         path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.BUILD_CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-0
     - name: Clone/compile gpuocelot
       if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
       shell: bash

.github/workflows/benchmark.yml

@@ -62,8 +62,10 @@ 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/speed/external_test_speed_v_torch.py | tee torch_speed.txt
+      run: BIG=2 MPS=1 python3.11 test/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
@@ -187,7 +189,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/speed/external_test_speed_v_torch.py | tee torch_speed.txt
+      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/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
@@ -389,7 +391,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/speed/external_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/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
@@ -603,12 +605,12 @@ 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: 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: 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.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

.github/workflows/test.yml

@@ -1,10 +1,8 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '3'
-  APT_CACHE_VERSION: '1'
-  BUILD_CACHE_VERSION: '1'
+  DOWNLOAD_CACHE_VERSION: '10'
+  PYTHON_CACHE_VERSION: '2'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
@@ -32,9 +30,9 @@ jobs:
     - name: External Benchmark Schedule
       run: PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
     - name: Speed Test
-      run: LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: LLVM=1 python3 test/test_speed_v_torch.py
     - name: Speed Test (BEAM=2)
-      run: BEAM=2 LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: BEAM=2 LLVM=1 python3 test/test_speed_v_torch.py
 
   docs:
     name: Docs
@@ -48,11 +46,6 @@ 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
@@ -336,6 +329,7 @@ jobs:
       run: |
         pip3 install --upgrade --force-reinstall ruff==0.11.0
         python3 -m ruff check .
+        python3 -m ruff check extra/onnx.py
         python3 -m ruff check examples/mlperf/ --ignore E501
     - name: Lint tinygrad with pylint
       run: python -m pylint tinygrad/
@@ -343,8 +337,7 @@ jobs:
       run: |
         python -m mypy --strict-equality --lineprecision-report .
         cat lineprecision.txt
-    - name: Run TYPED=1
-      run: TYPED=1 python -c "import tinygrad"
+        python -m mypy --strict-equality extra/onnx.py
 
   unittest:
     name: Unit Tests
@@ -382,8 +375,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 < 17500 lines
-      run: MAX_LINE_COUNT=17500 python sz.py
+    - name: Repo line count < 16000 lines
+      run: MAX_LINE_COUNT=16000 python sz.py
 
   fuzzing:
     name: Fuzzing
@@ -462,7 +455,7 @@ jobs:
   testopenpilot:
     name: 'openpilot Compile Tests'
     runs-on: ubuntu-22.04
-    timeout-minutes: 15
+    timeout-minutes: 10
     env:
       IGNORE_OOB: 0
     steps:
@@ -593,56 +586,6 @@ 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
@@ -678,6 +621,12 @@ 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 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
 
   testwebgpu:
     name: Linux (WebGPU)
@@ -737,9 +686,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/device/test_amd_llvm.py
+        run: python test/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/device/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/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
@@ -864,14 +813,14 @@ jobs:
         AMD: 1
         FORWARD_ONLY: 1
       run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
+        python3 -m pytest -n=auto test/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/device/test_hcq.py test/test_tiny.py test/device/test_amd_llvm.py --durations=20
+        python -m pytest -n=auto test/test_hcq.py test/test_tiny.py test/test_amd_llvm.py --durations=20
     - name: Run pytest (ptx)
       env:
         MOCKGPU: 1
@@ -879,7 +828,7 @@ jobs:
         NV: 1
         FORWARD_ONLY: 1
       run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
+        python3 -m pytest -n=auto test/test_hcq.py test/test_tiny.py --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -990,18 +939,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 --durations 20
+          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
       - 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 --durations 20
+          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
           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 --durations 20
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
       - name: Show remote server logs
         if: always()
         run: |

autogen_stubs.sh

@@ -198,7 +198,11 @@ 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
@@ -372,6 +376,26 @@ 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 \

docs/tensor/elementwise.md

@@ -78,7 +78,6 @@ 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
 

docs/tinybox.md

@@ -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, 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.
+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.
 
 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.
 

examples/beautiful_mnist.py

@@ -29,7 +29,8 @@ if __name__ == "__main__":
     opt.zero_grad()
     samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
     loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
-    return loss.realize(*opt.schedule_step())
+    opt.step()
+    return loss
 
   @TinyJit
   def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100

examples/hlb_cifar10.py

@@ -118,7 +118,7 @@ class SpeedyResNet:
 # hyper-parameters were exactly the same as the original repo
 bias_scaler = 58
 hyp = {
-  'seed' : 201,
+  'seed' : 200,
   'opt': {
     'bias_lr':            1.76 * bias_scaler/512,
     'non_bias_lr':        1.76 / 512,

examples/mlperf/model_train.py

@@ -1297,9 +1297,6 @@ def train_llama3():
   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
@@ -1321,10 +1318,6 @@ def train_llama3():
   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):
@@ -1346,8 +1339,6 @@ def train_llama3():
       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)
@@ -1378,7 +1369,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)).clamp(max_=1.0)
+        p.grad = p.grad * opt_gradient_clip_norm / (total_norm + 1e-6)
 
     optim.step()
     scheduler.step()
@@ -1387,40 +1378,16 @@ def train_llama3():
     loss.realize(lr)
     return loss, lr
 
-  @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()
+  if getenv("FAKEDATA", 0):
+    def fake_data():
+      for _ in range(SAMPLES // GBS):
+        yield Tensor.randint(GBS, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
+    iter = fake_data()
+  else:
+    from examples.mlperf.dataloader import batch_load_llama3
+    iter = batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))
 
-  # ** 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
+  i = 0
   for tokens in tqdm(iter, total=SAMPLES//GBS):
     t = time.perf_counter()
     GlobalCounters.reset()
@@ -1435,33 +1402,9 @@ def train_llama3():
     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"
+      fn = f"{ckpt_dir}/{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')

extra/amdpci/proclogs.py

@@ -1,16 +1,6 @@
-import re, ctypes, sys, importlib
+import re, ctypes, sys
 
-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
+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
 
 def parse_amdgpu_logs(log_content, register_names=None):
   register_map = register_names
@@ -33,11 +23,16 @@ 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 = {}
-  dev = PCIIface(None, 0)
-  for x, y in dev.dev_impl.__dict__.items():
-    if isinstance(y, AMRegister):
-      for inst, addr in y.addr.items(): reg_names[addr] = f"{x}, xcc={inst}"
+  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)])
 
   with open(sys.argv[1], 'r') as f:
     log_content = log_content_them = f.read()

extra/gemm/amd_uop_matmul.py

@@ -5,7 +5,7 @@ from tinygrad.dtype import AddrSpace
 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.codegen.opt.kernel import axis_colors, Opt, OptOps
+from tinygrad.codegen.opt.kernel import axis_colors
 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,28 +44,13 @@ 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)
   c = a@b
   sink = c.schedule()[-1].ast
   L = 16
-  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(N//BM, 0), 2:UOp.range(N//BN, 1)})
+  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(dtypes.int, N//BM, 0), 2:UOp.range(dtypes.int, N//BN, 1)})
   sink = graph_rewrite(sink, view_left+pm)
   axis_types = (AxisType.GLOBAL, AxisType.LOCAL, AxisType.GLOBAL, AxisType.LOCAL, AxisType.REDUCE)
   return sink.replace(arg=KernelInfo(name="top_"+to_colored(sink.full_shape, axis_types), axis_types=axis_types))
@@ -186,7 +171,7 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
 
   c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), arg=2)
 
-  i = UOp.range(c_regs.dtype.size, 16)
+  i = UOp.range(dtypes.int, c_regs.dtype.size, 16)
   init_store = c_regs[i].store(UOp.const(dtypes.float, 0.0), i)
 
   if kernel4:
@@ -197,53 +182,53 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
     kId = 0
 
     # load from globals into locals
-    i = UOp.range(nbReadsB, 0)
+    i = UOp.range(dtypes.int, nbReadsB, 0)
     index_x = BN * blockIdx_x + rBIdx
     index_y = rBIdy + i * strideReadB + kId
     Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
 
-    i = UOp.range(nbReadsA, 1)
+    i = UOp.range(dtypes.int, nbReadsA, 1)
     index_x = rAIdx + kId
     index_y = BM * blockIdx_y + rAIdy + i * strideReadA
     As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
 
     # iterate over the middle chunk
-    kId_range = UOp.range(N//BK-1, 2)
+    kId_range = UOp.range(dtypes.int, N//BK-1, 2)
     kId = kId_range*BK
 
     barrier = UOp.barrier(As_store, Bs_store)
 
     # load from globals into registers (next round)
-    i = UOp.range(nbReadsB, 3)
+    i = UOp.range(dtypes.int, nbReadsB, 3)
     index_x = BN * blockIdx_x + rBIdx
     index_y = rBIdy + i * strideReadB + kId + BK
     regB_store = regB[i].store(b[N * index_y + index_x].load(), i)
 
-    i = UOp.range(nbReadsA, 4)
+    i = UOp.range(dtypes.int, nbReadsA, 4)
     index_x = rAIdx + kId + BK
     index_y = BM * blockIdx_y + rAIdy + i * strideReadA
     regA_store = regA[i].store(a[N * index_y + index_x].load(), i)
 
     def inner_loop(first_range, inp_dep=()):
       # inner unroll
-      k = UOp.range(BK, first_range+0)
+      k = UOp.range(dtypes.int, BK, first_range+0)
 
       # load from locals into registers
-      iterWave = UOp.range(nbIterWaveN, first_range+1)
-      i = UOp.range(TN, first_range+2)
+      iterWave = UOp.range(dtypes.int, nbIterWaveN, first_range+1)
+      i = UOp.range(dtypes.int, TN, first_range+2)
       index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
       B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(*inp_dep), iterWave, i)
 
-      iterWave = UOp.range(nbIterWaveM, first_range+3)
-      i = UOp.range(TM, first_range+4)
+      iterWave = UOp.range(dtypes.int, nbIterWaveM, first_range+3)
+      i = UOp.range(dtypes.int, TM, first_range+4)
       index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
       A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(*inp_dep), iterWave, i)
 
       # do the GEMM math
-      iterWaveM = UOp.range(nbIterWaveM, first_range+5)
-      yt = UOp.range(TM, first_range+6)
-      iterWaveN = UOp.range(nbIterWaveN, first_range+7)
-      xt = UOp.range(TN, first_range+8)
+      iterWaveM = UOp.range(dtypes.int, nbIterWaveM, first_range+5)
+      yt = UOp.range(dtypes.int, TM, first_range+6)
+      iterWaveN = UOp.range(dtypes.int, nbIterWaveN, first_range+7)
+      xt = UOp.range(dtypes.int, TN, first_range+8)
       x = iterWaveN * TN + xt
       y = iterWaveM * TM + yt
       c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
@@ -256,12 +241,12 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
     sink = inner_loop(5, (barrier, regB_store, regA_store)).barrier()
 
     # load from registers into locals
-    i = UOp.range(nbReadsB, 14)
+    i = UOp.range(dtypes.int, nbReadsB, 14)
     index_x = BN * blockIdx_x + rBIdx
     index_y = rBIdy + i * strideReadB + kId + BK
     Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(regB[i].load(sink), i, kId_range)
 
-    i = UOp.range(nbReadsA, 15)
+    i = UOp.range(dtypes.int, nbReadsA, 15)
     index_x = rAIdx + kId + BK
     index_y = BM * blockIdx_y + rAIdy + i * strideReadA
     As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(regA[i].load(sink), i, kId_range)
@@ -269,40 +254,40 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
     # final iteration without the copy
     sink = inner_loop(16, (UOp.barrier(Bs_store, As_store),))
   else:
-    kId_range = UOp.range(N//BK, 0)
+    kId_range = UOp.range(dtypes.int, N//BK, 0)
     kId = kId_range*BK
 
     # load from globals into locals
-    i = UOp.range(nbReadsB, 1)
+    i = UOp.range(dtypes.int, nbReadsB, 1)
     index_x = BN * blockIdx_x + rBIdx
     index_y = rBIdy + i * strideReadB + kId
     Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
 
-    i = UOp.range(nbReadsA, 2)
+    i = UOp.range(dtypes.int, nbReadsA, 2)
     index_x = rAIdx + kId
     index_y = BM * blockIdx_y + rAIdy + i * strideReadA
     As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
 
     barrier = UOp.barrier(As_store, Bs_store)
 
-    k = UOp.range(BK, 3)
+    k = UOp.range(dtypes.int, BK, 3)
 
     # load from locals into registers
-    iterWave = UOp.range(nbIterWaveN, 4)
-    i = UOp.range(TN, 5)
+    iterWave = UOp.range(dtypes.int, nbIterWaveN, 4)
+    i = UOp.range(dtypes.int, TN, 5)
     index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
     B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(barrier), iterWave, i)
 
-    iterWave = UOp.range(nbIterWaveM, 6)
-    i = UOp.range(TM, 7)
+    iterWave = UOp.range(dtypes.int, nbIterWaveM, 6)
+    i = UOp.range(dtypes.int, TM, 7)
     index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
     A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(barrier), iterWave, i)
 
     # do the GEMM math
-    iterWaveM = UOp.range(nbIterWaveM, 8)
-    yt = UOp.range(TM, 9)
-    iterWaveN = UOp.range(nbIterWaveN, 10)
-    xt = UOp.range(TN, 12)
+    iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 8)
+    yt = UOp.range(dtypes.int, TM, 9)
+    iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 10)
+    xt = UOp.range(dtypes.int, TN, 12)
     x = iterWaveN * TN + xt
     y = iterWaveM * TM + yt
     c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
@@ -310,10 +295,10 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
                             iterWaveM, iterWaveN, yt, xt, k, kId_range)
 
   # store c_regs into c
-  iterWaveM = UOp.range(nbIterWaveM, 1000)
-  yt = UOp.range(TM, 1001)
-  iterWaveN = UOp.range(nbIterWaveN, 1002)
-  xt = UOp.range(TN, 1003)
+  iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 1000)
+  yt = UOp.range(dtypes.int, TM, 1001)
+  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 1002)
+  xt = UOp.range(dtypes.int, TN, 1003)
   xOut = blockIdx_x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave
   yOut = blockIdx_y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave
   indexC = N * (yOut + yt) + xOut + xt
@@ -324,15 +309,10 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
 
 if __name__ == "__main__":
   HL = getenv("HL")
-  if HL == 3: hprg = rangeify_kernel3()
-  elif HL == 2: hprg = top_spec_kernel3()
+  if HL == 2: hprg = top_spec_kernel3()
   elif HL == 1: hprg = hl_spec_kernel3()
   else: hprg = hand_spec_kernel3()
-  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)
+  prg = get_program(hprg, Device.default.renderer)
   print(prg.src)
   if getenv("SRC"): exit(0)
   hrunner = CompiledRunner(prg)

extra/gemm/max_matmul.py

@@ -56,7 +56,7 @@ def randoms():
 def ast_to_cuda_prog(compiler, ast, opts):
   k = Kernel(ast)
   k.apply_opts(opts)
-  p = get_program(k.ast, k.opts, k.applied_opts)
+  p = get_program(k.get_optimized_ast(), k.opts)
   return CUDAProgram(device, p.function_name, compiler.compile(p.src))
 
 if __name__ == "__main__":

extra/gemm/tinygrad_nv_matmul.py

@@ -29,7 +29,7 @@ if __name__ == "__main__":
       Opt(op=OptOps.LOCAL, axis=0, amt=2),
     ]
   k.apply_opts(opts)
-  prg = get_program(k.ast, k.opts, k.applied_opts)
+  prg = get_program(k.get_optimized_ast(), k.opts)
   new_src = prg.src
   # can mod source here
   prg = replace(prg, src=new_src)

extra/huggingface_onnx/README.md

@@ -1,61 +0,0 @@
-# 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
-```

extra/huggingface_onnx/collect_metadata.py

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

extra/huggingface_onnx/download_models.py

@@ -0,0 +1,29 @@
+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))

extra/huggingface_onnx/huggingface_manager.py

@@ -1,230 +0,0 @@
-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
-  )

extra/huggingface_onnx/run_models.py

@@ -1,11 +1,10 @@
-import onnx, yaml, tempfile, time, argparse, json
+import onnx, yaml, tempfile, time, collections, pprint, 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) -> dict[str, Any]:
+def get_config(root_path: Path):
   ret = {}
   for path in root_path.rglob("*config.json"):
     config = json.load(path.open())
@@ -13,19 +12,19 @@ def get_config(root_path: Path) -> dict[str, Any]:
       ret.update(config)
   return ret
 
-def get_tolerances(file_name: str) -> tuple[float, float]:
+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
   # 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(models)} models **")
+  print(f"** Validating {len(model_paths)} models **")
   for model_id, (root_path, relative_path) in models.items():
     print(f"validating model {model_id}")
     model_path = root_path / relative_path
@@ -37,6 +36,25 @@ 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)
@@ -53,9 +71,12 @@ 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")
-  parser.add_argument("--validate", type=str, default="",
-                      help="Validate correctness of models from the specified YAML configuration file")
+  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.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
@@ -64,13 +85,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.validate or args.debug):
-    parser.error("Please provide either --validate <yaml_file> or --debug <repo_id>.")
+  if not (args.check_ops or args.validate or args.debug):
+    parser.error("Please provide either --validate, --check_ops, or --debug.")
   if args.truncate != -1 and not args.debug:
     parser.error("--truncate and --debug should be used together for debugging")
 
-  if args.validate:
-    with open(args.validate, 'r') as f:
+  if args.check_ops or args.validate:
+    with open(args.input, '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 = {
@@ -80,16 +101,22 @@ if __name__ == "__main__":
         if model["file"].endswith(".onnx")
       }
 
-    validate_repos(model_paths)
+    if args.check_ops:
+      pprint.pprint(retrieve_op_stats(model_paths))
+
+    if args.validate:
+      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 = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=DOWNLOADS_DIR)
-      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
+      root_path = 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)
       config = get_config(root_path)
       for onnx_model in root_path.rglob("*.onnx"):
         rtol, atol = get_tolerances(onnx_model.name)
@@ -101,8 +128,8 @@ if __name__ == "__main__":
       onnx_model = path[-1]
       assert path[-1].endswith(".onnx")
       repo_id, relative_path = "/".join(path[:2]), "/".join(path[2:])
-      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=DOWNLOADS_DIR)
-      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
+      root_path = 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)
       config = get_config(root_path)
       rtol, atol = get_tolerances(onnx_model)
       print(f"validating {relative_path} with truncate={args.truncate}, {rtol=}, {atol=}")

extra/onnx_helpers.py

@@ -1,6 +1,7 @@
 from tinygrad import Tensor
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner, OnnxValue
+from tinygrad.frontend.onnx import OnnxRunner
+from extra.onnx import OnnxValue
 import numpy as np
 import onnxruntime as ort
 

extra/optimization/test_beam_search.py

@@ -16,9 +16,9 @@ class TestBeamSearch(unittest.TestCase):
     BEAM.value = self.old_beam
 
   def test_variable_ast_beam(self):
-    vi = Variable("a", 1, 10).bind(3)
-    a = rand(10, 3)[:vi]
-    a = (a+1).realize()
+    with Context(IGNORE_OOB=1):
+      a = rand(3, 3).reshape((Variable("a", 1, 10).bind(3), 3))
+      a = (a+1).realize()
 
   def test_big_prime_number(self):
     a = rand(367, 367)
@@ -42,16 +42,18 @@ class TestBeamSearch(unittest.TestCase):
 
   def test_variable_big_prime_number(self):
     v = Variable("v", 1, 400).bind(367)
-    a = rand(367, 400)
-    b = rand(400, 367)
-    c = (a[:, :v] @ b[:v, :]).realize()
-    np.testing.assert_allclose(c.numpy(), a[:, :367].numpy() @ b[:367, :].numpy(), atol=1e-4, rtol=1e-4)
+    a = rand(367, 367)
+    b = rand(367, 367)
+    with Context(IGNORE_OOB=1):
+      c = (a.reshape(367, v) @ b.reshape(v, 367)).realize()
+      np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), atol=1e-4, rtol=1e-4)
 
   def test_variable_shrink_prime_number(self):
     v = Variable("v", 1, 400).bind(367)
     a = rand(400, 367)
-    b = (a.shrink(((0,v), None))+1).reshape(367,367).realize()
-    np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)
+    with Context(IGNORE_OOB=1):
+      b = (a.shrink(((0,v), None))+1).reshape(367,367).realize()
+      np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)
 
   def test_no_mutate_rawbuffers(self):
     a = rand(3, 3).realize()

extra/replay_pkl.py

@@ -58,7 +58,7 @@ if __name__ == "__main__":
             GlobalCounters.kernel_count -= 1
 
         if not getenv("NOOPT"): k.apply_opts(hand_coded_optimizations(k))
-        p2 = get_program(k.ast, k.opts, k.applied_opts)
+        p2 = get_program(k.get_optimized_ast(), k.opts)
         new_ei = replace(ei, prg=CompiledRunner(p2))
         new_ei.run()
         new_jit.append(new_ei)

extra/torch_backend/backend.py

@@ -381,7 +381,6 @@ 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,

pytest.ini

@@ -1,2 +0,0 @@
-[pytest]
-norecursedirs = extra

ruff.toml

@@ -35,7 +35,6 @@ lint.select = [
 line-length = 150
 
 exclude = [
-  ".git/",
   "docs/",
   "extra/",
   "tinygrad/runtime/autogen",

setup.py

@@ -18,35 +18,16 @@ testing_minimal = [
 ]
 
 setup(name='tinygrad',
-      version='0.11.0',
+      version='0.10.3',
       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.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',
-      ],
+      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.codegen.opt',
+                  'tinygrad.runtime.support.nv', 'tinygrad.apps'],
       package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
       classifiers=[
         "Programming Language :: Python :: 3",
@@ -64,7 +45,6 @@ setup(name='tinygrad',
             "pre-commit",
             "ruff",
             "numpy",
-            "typeguard",
         ],
         #'mlperf': ["mlperf-logging @ git+https://github.com/mlperf/[email protected]"],
         'testing_minimal': testing_minimal,

test/external/external_benchmark_kernel_launch.py

@@ -2,7 +2,7 @@ import time
 from tinygrad import Tensor, TinyJit, Device, Context
 from tinygrad.helpers import Profiling, Timing, GlobalCounters
 
-# python3 test/speed/external_test_speed_v_torch.py TestSpeed.test_add_a
+# python3 test/test_speed_v_torch.py TestSpeed.test_add_a
 
 @TinyJit
 def plus(a:Tensor, b:Tensor): return a+b

test/external/external_benchmark_sdxl_softmax.py

@@ -24,5 +24,5 @@ if __name__ == "__main__":
   #k.apply_opt(Opt(OptOps.GROUP, 1, 32))
   #k.apply_opt(Opt(OptOps.GROUP, 0, 32))
   from tinygrad.engine.realize import CompiledRunner, ExecItem
-  run = CompiledRunner(prg:=get_program(k.ast, k.opts, k.applied_opts))
+  run = CompiledRunner(prg:=get_program(k.get_optimized_ast(), k.opts))
   ExecItem(run, si.bufs).run()

test/external/external_debug_metal_sd_conv.py

@@ -35,7 +35,7 @@ k = Kernel(ast)
 k.apply_opts(opts)
 bufs = bufs_from_lin(k)
 
-prg = CompiledRunner(get_program(k.ast, k.opts, k.applied_opts))
+prg = CompiledRunner(get_program(k.get_optimized_ast(), k.opts))
 
 for i in range(10):
   speed = prg(bufs, var_vals={}, wait=True)

test/external/external_test_onnx_backend.py

@@ -134,6 +134,7 @@ backend_test.exclude('test_simple_rnn_*')
 
 # no control flow
 # control flow uses AttributeProto.GRAPH
+backend_test.exclude('test_if_*')
 backend_test.exclude('test_loop*')
 backend_test.exclude('test_range_float_type_positive_delta_expanded_cpu') # requires loop
 backend_test.exclude('test_affine_grid_2d_align_corners_expanded_cpu')
@@ -182,8 +183,6 @@ backend_test.exclude('test_resize_downsample_scales_cubic_antialias_cpu') # anti
 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
-backend_test.exclude('test_if_opt_cpu') # ValueError: 13 is not a valid AttributeType
-backend_test.exclude('test_if_seq_cpu') # NotImplementedError: op='SequenceConstruct' is not supported
 
 backend_test.exclude('test_scatternd_min_cpu') # min not yet supported
 backend_test.exclude('test_scatternd_max_cpu') # max not yet supported

test/external/external_test_onnx_ops.py

@@ -100,25 +100,6 @@ class TestMainOnnxOps(TestOnnxOps):
     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_if(self, then_value, else_value):
-    then_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, then_value.shape)
-    else_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, else_value.shape)
-
-    then_const_node = onnx.helper.make_node("Constant", inputs=[], outputs=["res"], value=onnx.numpy_helper.from_array(then_value))
-    else_const_node = onnx.helper.make_node("Constant", inputs=[], outputs=["res"], value=onnx.numpy_helper.from_array(else_value))
-
-    then_body = onnx.helper.make_graph([then_const_node], "then_body", [], [then_out])
-    else_body = onnx.helper.make_graph([else_const_node], "else_body", [], [else_out])
-
-    self.helper_test_single_op("If", {"cond": np.array(False).astype(bool)}, {"then_branch": then_body, "else_branch": else_body}, ["res"])
-    self.helper_test_single_op("If", {"cond": np.array(True).astype(bool)}, {"then_branch": then_body, "else_branch": else_body}, ["res"])
-
-  def test_if_different_shapes_broadcastable(self):
-    self._test_if(np.array([[1], [2]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
-
-  def test_if_different_shapes_not_broadcastable(self):
-    self._test_if(np.array([[1, 2, 3], [4, 5, 6]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
-
   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)

test/external/external_test_onnx_runner.py

@@ -3,7 +3,8 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.uop.ops import Ops
 from tinygrad.device import is_dtype_supported
-from tinygrad.frontend.onnx import OnnxRunner, OnnxDataType
+from extra.onnx import OnnxDataType
+from tinygrad.frontend.onnx import OnnxRunner
 from hypothesis import given, strategies as st
 
 # copied from test_const_folding.py

test/external/fuzz_fast_idiv.py

@@ -1,8 +1,8 @@
 import random
 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.spec import z3_renderer, z3_cdiv
+from tinygrad.uop.ops import UOp, graph_rewrite
 from tinygrad.uop.decompositions import fast_idiv
 random.seed(42)
 
@@ -19,7 +19,8 @@ if __name__ == "__main__":
     if expr is None: continue
 
     solver = z3.Solver()
-    z3_expr, x =uops_to_z3(solver, expr, u)
+    z3_sink = graph_rewrite(expr.sink(u), z3_renderer, ctx=(solver, {}))
+    z3_expr, x = z3_sink.src[0].arg, z3_sink.src[1].arg
 
     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=}"

test/external/fuzz_symbolic.py

@@ -1,8 +1,8 @@
 import random, operator
 import z3
 from tinygrad import Variable, dtypes
-from tinygrad.uop.ops import UOp
-from tinygrad.uop.spec import uops_to_z3
+from tinygrad.uop.ops import UOp, graph_rewrite
+from tinygrad.uop.spec import z3_renderer
 from tinygrad.helpers import DEBUG, Context
 
 seed = random.randint(0, 100)
@@ -57,7 +57,8 @@ 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_expr, z3_simplified_expr, v1, v2, v3 = uops_to_z3(solver, expr, simplified_expr, u1, u2, u3)
+    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
     check = solver.check(z3_simplified_expr != z3_expr)
     if check == z3.unknown and DEBUG>=1:
       skipped += 1
@@ -68,6 +69,7 @@ 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):

test/external/process_replay/process_replay.py

@@ -99,6 +99,7 @@ def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
     except Exception as e:
       changed += 1
       warnings.warn(f"{name=} {loc=} {e=}", ProcessReplayWarning)
+  conn.commit()
   cur.close()
 
 # *** generic runner to map rows of a table to a function in parallel
@@ -110,11 +111,12 @@ 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:
-    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)
+    inputs = list(range(0, row_count, PAGE_SIZE))
+    list(tqdm(pool.imap_unordered(functools.partial(diff, fxns=fxns), inputs), total=len(inputs)))
     pool.close()
     pool.join()
     pool.terminate()

test/mockgpu/amd/amddriver.py

@@ -87,19 +87,16 @@ 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')),

test/mockgpu/amd/amdgpu.py

@@ -1,8 +1,7 @@
 import ctypes, time
 from test.mockgpu.gpu import VirtGPU
-from test.mockgpu.helpers import _try_dlopen_remu
 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
+import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4, tinygrad.runtime.autogen.am.soc21 as soc21
 
 SDMA_MAX_COPY_SIZE = 0x400000
 
@@ -15,9 +14,6 @@ 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
@@ -25,6 +21,19 @@ WAIT_REG_MEM_FUNCTION_EQ  = 3 # ==
 WAIT_REG_MEM_FUNCTION_NEQ = 4 # !=
 WAIT_REG_MEM_FUNCTION_GEQ = 5 # >=
 
+REMU_PATHS = ["extra/remu/target/release/libremu.so", "libremu.so", "/usr/local/lib/libremu.so",
+              "extra/remu/target/release/libremu.dylib", "libremu.dylib", "/usr/local/lib/libremu.dylib", "/opt/homebrew/lib/libremu.dylib"]
+def _try_dlopen_remu():
+  for path in REMU_PATHS:
+    try:
+      remu = ctypes.CDLL(path)
+      remu.run_asm.restype = ctypes.c_int32
+      remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
+        ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
+    except OSError: pass
+    else: return remu
+  print("Could not find libremu.so")
+  return None
 remu = _try_dlopen_remu()
 
 def create_sdma_packets():
@@ -199,7 +208,7 @@ class PM4Executor(AMDQueue):
     assert n == 0
     event_dw = self._next_dword()
     match (event_dw & 0xFF): # event type
-      case SQTT_EVENTS.THREAD_TRACE_FINISH:
+      case soc21.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

test/mockgpu/cuda/cuda.py

@@ -2,14 +2,16 @@ from __future__ import annotations
 from typing import Any
 import ctypes, time
 from tinygrad.runtime.autogen import cuda as orig_cuda
-from test.mockgpu.helpers import _try_dlopen_gpuocelot
 from tinygrad.helpers import mv_address
 
 for attr in dir(orig_cuda):
   if not attr.startswith('__'):
     globals()[attr] = getattr(orig_cuda, attr)
 
-gpuocelot_lib = _try_dlopen_gpuocelot()
+try:
+  gpuocelot_lib = ctypes.CDLL(ctypes.util.find_library("gpuocelot"))
+  gpuocelot_lib.ptx_run.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.POINTER(ctypes.c_void_p), ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int]  # noqa: E501
+except Exception: pass
 
 # Global state
 class CUDAState:

test/mockgpu/helpers.py

@@ -1,29 +0,0 @@
-import ctypes, ctypes.util
-
-def _try_dlopen_gpuocelot():
-  GPUOCELOT_PATHS = [ctypes.util.find_library("gpuocelot")] if ctypes.util.find_library("gpuocelot") is not None else []
-  GPUOCELOT_PATHS += ["libgpuocelot.so", "/usr/local/lib/libgpuocelot.so",
-                      "libgpuocelot.dylib", "/usr/local/lib/libgpuocelot.dylib", "/opt/homebrew/lib/libgpuocelot.dylib"]
-  for path in GPUOCELOT_PATHS:
-    try:
-      gpuocelot_lib = ctypes.CDLL(path)
-      gpuocelot_lib.ptx_run.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.POINTER(ctypes.c_void_p), ctypes.c_int, ctypes.c_int,
-        ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int]
-    except OSError: pass
-    else: return gpuocelot_lib
-  print("Could not find libgpuocelot.so")
-  return None
-
-def _try_dlopen_remu():
-  REMU_PATHS = ["extra/remu/target/release/libremu.so", "libremu.so", "/usr/local/lib/libremu.so",
-               "extra/remu/target/release/libremu.dylib", "libremu.dylib", "/usr/local/lib/libremu.dylib", "/opt/homebrew/lib/libremu.dylib"]
-  for path in REMU_PATHS:
-    try:
-      remu = ctypes.CDLL(path)
-      remu.run_asm.restype = ctypes.c_int32
-      remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
-        ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
-    except OSError: pass
-    else: return remu
-  print("Could not find libremu.so")
-  return None

test/mockgpu/nv/nvgpu.py

@@ -2,7 +2,6 @@ import ctypes, ctypes.util, time
 import tinygrad.runtime.autogen.nv_gpu as nv_gpu
 from enum import Enum, auto
 from test.mockgpu.gpu import VirtGPU
-from test.mockgpu.helpers import _try_dlopen_gpuocelot
 from tinygrad.helpers import to_mv, init_c_struct_t
 
 def make_qmd_struct_type():
@@ -17,7 +16,10 @@ def make_qmd_struct_type():
 qmd_struct_t = make_qmd_struct_type()
 assert ctypes.sizeof(qmd_struct_t) == 0x40 * 4
 
-gpuocelot_lib = _try_dlopen_gpuocelot()
+try:
+  gpuocelot_lib = ctypes.CDLL(ctypes.util.find_library("gpuocelot"))
+  gpuocelot_lib.ptx_run.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.POINTER(ctypes.c_void_p), ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int]  # noqa: E501
+except Exception: pass
 
 class SchedResult(Enum): CONT = auto(); YIELD = auto() # noqa: E702
 

test/models/test_onnx.py

@@ -9,15 +9,7 @@ except ModuleNotFoundError:
   raise unittest.SkipTest("onnx not installed, skipping onnx test")
 from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.tensor import Tensor
-from tinygrad.device import Device
-from tinygrad.helpers import CI, fetch, temp, Context
-
-try:
-  from extra.onnx_helpers import validate
-  from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry
-  HUGGINGFACE_AVAILABLE = True
-except ModuleNotFoundError:
-  HUGGINGFACE_AVAILABLE = False
+from tinygrad.helpers import CI, fetch, temp
 
 def run_onnx_torch(onnx_model, inputs):
   import torch
@@ -32,7 +24,6 @@ OPENPILOT_MODEL = "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/mod
 np.random.seed(1337)
 
 class TestOnnxModel(unittest.TestCase):
-  @unittest.skip("this isn't a test, it can't fail")
   def test_benchmark_openpilot_model(self):
     onnx_model = fetch(OPENPILOT_MODEL)
     run_onnx = OnnxRunner(onnx_model)
@@ -146,36 +137,5 @@ class TestOnnxModel(unittest.TestCase):
     print(cls, _LABELS[cls])
     assert "car" in _LABELS[cls] or _LABELS[cls] == "convertible"
 
-@unittest.skipUnless(HUGGINGFACE_AVAILABLE and Device.DEFAULT == "METAL", "only run on METAL")
-class TestHuggingFaceOnnxModels(unittest.TestCase):
-  @classmethod
-  def setUpClass(cls):
-    cls._ctx = Context(MAX_BUFFER_SIZE=0)
-    cls._ctx.__enter__()
-
-  @classmethod
-  def tearDownClass(cls):
-    cls._ctx.__exit__()
-
-  def _validate(self, repo_id, model_file, custom_inputs, rtol=1e-4, atol=1e-4):
-    onnx_model_path = snapshot_download_with_retry(
-      repo_id=repo_id,
-      allow_patterns=["*.onnx", "*.onnx_data"],
-      cache_dir=str(DOWNLOADS_DIR)
-    )
-    onnx_model_path = onnx_model_path / model_file
-    file_size = onnx_model_path.stat().st_size
-    print(f"Validating model: {repo_id}/{model_file} ({file_size/1e6:.2f}M)")
-    validate(onnx_model_path, custom_inputs, rtol=rtol, atol=atol)
-
-  def test_xlm_roberta_large(self):
-    repo_id = "FacebookAI/xlm-roberta-large"
-    model_file = "onnx/model.onnx"
-    custom_inputs = {
-      "input_ids": np.random.randint(0, 250002, (1, 11), dtype=np.int64),
-      "attention_mask": np.ones((1, 11), dtype=np.int64),
-    }
-    self._validate(repo_id, model_file, custom_inputs)
-
 if __name__ == "__main__":
   unittest.main()

test/models/test_whisper.py

@@ -16,7 +16,7 @@ TRANSCRIPTION_2 = "a slightly longer audio file so that we can test batch transc
 TEST_FILE_3_URL = 'https://homepage.ntu.edu.tw/~karchung/miniconversations/mc45.mp3'
 TRANSCRIPTION_3 = "Just lie back and relax. Is the level of pressure about right? Yes, it's fine, and I'd like conditioner please. Sure. I'm going to start the second lathering now. Would you like some Q-tips? How'd you like it cut? I'd like my bangs and the back trimmed, and I'd like the rest thinned out a bit and layered. Where would you like the part? On the left, right about here. Here, have a look. What do you think? It's fine. Here's a thousand anti-dollars. It's 30-ant extra for the rants. Here's your change and receipt. Thank you, and please come again. So how do you like it? It could have been worse, but you'll notice that I didn't ask her for her card. Hmm, yeah. Maybe you can try that place over there next time."   # noqa: E501
 
-@unittest.skipIf(Device.DEFAULT in ["CPU", "LLVM"], "slow")
+@unittest.skipIf(CI and Device.DEFAULT in ["CPU"], "slow")
 @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16 support")
 class TestWhisper(unittest.TestCase):
   @classmethod

test/test_compile_failures.py

@@ -1,10 +1,7 @@
-import unittest, io
-from contextlib import redirect_stdout
+import unittest
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.helpers import OSX
 from tinygrad.engine.realize import lower_schedule
 from tinygrad.device import is_dtype_supported
-from tinygrad.engine.realize import get_program
 
 class TestCompileFailures(unittest.TestCase):
   def compile(self, out:Tensor):
@@ -17,17 +14,5 @@ class TestCompileFailures(unittest.TestCase):
   def test_add_max_uchar(self):
     self.compile((Tensor.empty(1024, dtype='uint8') + Tensor.empty(1024, dtype='uint8')).max())
 
-class TestDisassembly(unittest.TestCase):
-  # TODO: fails on llvm. llvm.LLVMGetHostCPUName() returns "generic"
-  @unittest.skipUnless(Device.DEFAULT in ("CPU",) and OSX, "m series cpus support fp16 arithmetic")
-  def test_float16_alu(self):
-    c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
-    s = c.schedule()[-1]
-    p = get_program(s.ast, Device[Device.DEFAULT].renderer)
-    lib = Device[Device.DEFAULT].compiler.compile(p.src)
-    out = io.StringIO()
-    with redirect_stdout(out): Device[Device.DEFAULT].compiler.disassemble(lib)
-    assert "fcvt" not in out.getvalue()
-
 if __name__ == '__main__':
   unittest.main()

test/test_const_folding.py

@@ -1,10 +1,11 @@
 import unittest, itertools, math
+from typing import Any
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.dtype import DType, ConstType
+from tinygrad.dtype import DType
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen import full_rewrite_to_sink
-from tinygrad.device import is_dtype_supported
 import numpy as np
+from tinygrad.device import is_dtype_supported
 from test.helpers import not_support_multi_device
 
 def _check_ast_count(desired_count:int, t:Tensor):
@@ -24,7 +25,7 @@ class TestUnaryOpsConstFolding(unittest.TestCase):
     _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
     _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))
 
-  @unittest.expectedFailure  # no two level fold
+  @unittest.expectedFailure  # no two level fold at lazybuffer
   def test_neg_folding(self):
     _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
     _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
@@ -103,7 +104,7 @@ class TestBinaryOpsConstFolding(unittest.TestCase):
 
 class TestBitcastConstFolding(unittest.TestCase):
   def test_scalar_bitcast(self):
-    def t(cases: dict[DType, ConstType]):
+    def t(cases: dict[DType, Any]):
       for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
         if not math.isnan(from_v):
           r = full_rewrite_to_sink(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0]
@@ -142,12 +143,13 @@ class TestIndexingConstFolding(unittest.TestCase):
     _check_ast_count(1, t[:,:,Tensor(1)+2,:])
     _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
 
+  @unittest.expectedFailure
   def test_const_tensor_index(self):
-    # TODO: these can be 0, implement const tensor folded indexing
+    # TODO: implement const tensor folded indexing
     t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
-    _check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
-    _check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
-    _check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])
+    _check_ast_count(0, t[:,:,Tensor.ones(2,1),:])
+    _check_ast_count(0, t[:,:,Tensor.ones(1,2)+2,:])
+    _check_ast_count(0, t[:,:,Tensor.ones(1,1),Tensor.zeros(2,1,2)])
 
 class TestMovedConstFolding(unittest.TestCase):
   def test_add_shrunk_zero(self):
@@ -164,6 +166,7 @@ class TestMovedConstFolding(unittest.TestCase):
     _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
 
   def test_cast_padded(self):
+    # NOTE: this is folded due to CAST_BEFORE_VIEW
     if is_dtype_supported(dtypes.int16):
       _check_ast_count(0, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
       np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])

test/test_copy_speed.py

@@ -3,7 +3,7 @@ from tinygrad import Tensor, Device, TinyJit
 from tinygrad.helpers import Timing, CI, OSX
 import multiprocessing.shared_memory as shared_memory
 
-N = 256
+N = 256 if CI else 4096
 class TestCopySpeed(unittest.TestCase):
   @classmethod
   def setUpClass(cls): Device[Device.DEFAULT].synchronize()

test/test_define_reg.py

@@ -0,0 +1,32 @@
+import unittest
+from tinygrad import dtypes, Device, Tensor, Context
+from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.engine.realize import get_program, ExecItem, CompiledRunner
+
+class TestDefineReg(unittest.TestCase):
+  def test_simple(self, at=AxisType.UPCAST):
+    N = 16
+    bout = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
+    a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
+    a_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(N, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((N,N), (0,1)))
+
+    out = a_col.load(a_col.store(a.load()))
+    sink = bout.store(out).sink(arg=KernelInfo(name="regcopy", axis_types=(AxisType.LOOP, at)))
+    prg = get_program(sink, Device.default.renderer)
+
+    with Context(DEBUG=0):
+      a = Tensor.randn(N, N).realize()
+      b = Tensor.empty(N, N).realize()
+    hrunner = CompiledRunner(prg)
+    ExecItem(hrunner, [b.uop.buffer, a.uop.buffer]).run(wait=True)
+    with Context(DEBUG=0):
+      self.assertEqual((b-a).mean().item(), 0.0)
+
+  @unittest.skipIf(getenv("PTX"), "ptx needs regs to be unrolled")
+  def test_simple_loop(self): self.test_simple(AxisType.LOOP)
+
+if __name__ == '__main__':
+  unittest.main()

test/test_disassembly.py

@@ -0,0 +1,21 @@
+import unittest, io
+from tinygrad import Tensor, dtypes
+from contextlib import redirect_stdout
+from tinygrad.device import Device
+from tinygrad.helpers import OSX
+from tinygrad.engine.realize import get_program
+
+class TestDisassembly(unittest.TestCase):
+  # TODO: fails on llvm. llvm.LLVMGetHostCPUName() returns "generic"
+  @unittest.skipUnless(Device.DEFAULT in ("CPU",) and OSX, "m series cpus support fp16 arithmetic")
+  def test_float16_alu(self):
+    c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
+    s = c.schedule()[-1]
+    p = get_program(s.ast, Device[Device.DEFAULT].renderer)
+    lib = Device[Device.DEFAULT].compiler.compile(p.src)
+    out = io.StringIO()
+    with redirect_stdout(out): Device[Device.DEFAULT].compiler.disassemble(lib)
+    assert "fcvt" not in out.getvalue()
+
+if __name__ == "__main__":
+  unittest.main()

test/test_dtype.py

@@ -4,8 +4,9 @@ import torch
 from typing import Any, List
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, DEBUG, CI
-from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype
+from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8
 from tinygrad import Device, Tensor, dtypes
+from tinygrad.tensor import _to_np_dtype
 from hypothesis import assume, given, settings, strategies as strat
 from test.helpers import rand_for_dtype
 from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
@@ -23,10 +24,6 @@ def get_available_cast_dtypes(dtype: DType) -> List[DType]:
   # dont cast internal dtypes
   return [v for k, v in DTYPES_DICT.items() if v != dtype and is_dtype_supported(v) and not k.startswith("_")]
 
-def _to_torch_storage_type(dtype:DType):
-  if dtype == dtypes.bfloat16: return torch.float32
-  return _to_torch_dtype(dtype)
-
 def _test_to_np(a:Tensor, np_dtype, target):
   if DEBUG >= 2: print(a)
   na = a.numpy()
@@ -49,10 +46,10 @@ def _test_cast(a:Tensor, target_dtype:DType):
 
   _test_op(lambda: a.cast(target_dtype), target_dtype, list(a.numpy().astype(_to_np_dtype(target_dtype))))
 def _test_bitcast(a:Tensor, target_dtype:DType, target=None):
+  if target_dtype == dtypes.bfloat16: raise unittest.SkipTest("no test for bf16 bitcast yet")
   if getenv("PTX") and a.dtype == dtypes.int8 and target_dtype.itemsize != a.dtype.itemsize:
     raise unittest.SkipTest("shape changing bitcast of int8 broken on PTX")
-  expected = torch.tensor(a.tolist(), dtype=_to_torch_storage_type(a.dtype)).view(_to_torch_dtype(target_dtype))
-  _test_op(lambda: a.bitcast(target_dtype), target_dtype, target or expected.tolist())
+  _test_op(lambda: a.bitcast(target_dtype), target_dtype, target or a.numpy().view(_to_np_dtype(target_dtype)).tolist())
 
 class TestDType(unittest.TestCase):
   DTYPE: Any = None
@@ -129,7 +126,7 @@ class TestDType(unittest.TestCase):
 
   def test_finfo(self):
     if self.DTYPE not in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]: return
-    info = ml_dtypes.finfo(ml_dtypes.bfloat16 if self.DTYPE is dtypes.bfloat16 else _to_np_dtype(self.DTYPE))
+    info = np.finfo(_to_np_dtype(self.DTYPE))
     assert info.bits == self.DTYPE.itemsize*8
     assert info.nexp == dtypes.finfo(self.DTYPE)[0]
     assert info.nmant == dtypes.finfo(self.DTYPE)[1]
@@ -302,10 +299,10 @@ class TestBitCast(unittest.TestCase):
   @given(strat.sampled_from(dtype_ints + dtype_floats), strat.sampled_from(dtype_ints + dtype_floats))
   def test_shape_change_bitcast(self, dt1, dt2):
     # NOTE: this has to be assume to prevent hypothesis from skipping all samples
+    assume(dt2 != dtypes.bfloat16 and dt1 != dtypes.bfloat16) # no test for bf16 bitcast yet
     assume(not (getenv("PTX") and dt1 == dtypes.int8)) # TODO: bitcasting int8 fails in PTX
     data = rand_for_dtype(dt1, 32).reshape(2, 2, 8)
-    expected = torch.tensor(data.tolist(), dtype=_to_torch_storage_type(dt1)).view(_to_torch_dtype(dt2))
-    _test_op(lambda: Tensor(data, dtype=dt1).bitcast(dt2), dt2, expected.tolist())
+    _test_op(lambda: Tensor(data, dtype=dt1).bitcast(dt2), dt2, data.view(_to_np_dtype(dt2)).tolist())
 
   def test_shape_change_bitcast_exceptions(self):
     with self.assertRaises(RuntimeError):
@@ -345,9 +342,6 @@ class TestUint64DType(TestDType):
 
 class TestBoolDType(TestDType): DTYPE = dtypes.bool
 
-@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
-class TestBFloat16Type(TestDType): DTYPE = dtypes.bfloat16
-
 class TestPtrDType(unittest.TestCase):
   def test_vec_double(self):
     dt1 = dtypes.float.vec(4).ptr().vec(4)
@@ -424,7 +418,7 @@ class TestDtypeUsage(unittest.TestCase):
 class TestOpsBFloat16(unittest.TestCase):
   def test_cast(self):
     # TODO: helper_test_op breaks in unrelated part
-    # TODO: wrong output with GPU=1 on mac
+    # TODO: wrong output with GPU=1 / PYTHON=1 on mac
     data = [60000.0, 70000.0, 80000.0]
     np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy())
 

test/test_dtype_alu.py

@@ -1,13 +1,16 @@
-import unittest, operator, math
+import unittest
+
 from tinygrad import Tensor, dtypes, Device
+import operator
+import numpy as np
+from hypothesis import given, strategies as strat, settings, HealthCheck
 from tinygrad.dtype import DType
 from tinygrad.helpers import CI, getenv
+from tinygrad.engine.realize import run_schedule
+from tinygrad.uop.ops import GroupOp
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
-import numpy as np
-import pytest
-from hypothesis import given, strategies as strat, settings, HealthCheck
-
+import pytest, math
 pytestmark = pytest.mark.filterwarnings("ignore")
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
@@ -60,21 +63,25 @@ def universal_test(a, b, dtype, op):
   ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype)
   tensor_value = (op[0](ta, tb)).numpy()
   numpy_value = op[1](ta.numpy(), tb.numpy())
-  if dtype in dtypes.floats:
-    atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2)}.get(dtype, (1e-10, 1e-7))
-    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
+  if dtype == dtypes.bfloat16: np.testing.assert_allclose(tensor_value, numpy_value, atol=1e-3, rtol=1e-2)
+  elif dtype in dtypes_float: np.testing.assert_allclose(tensor_value, numpy_value, atol=1e-10)
   else: np.testing.assert_equal(tensor_value, numpy_value)
 
 def universal_test_unary(a, dtype, op):
   if not isinstance(op, tuple): op = (op, op)
   ta = Tensor([a], dtype=dtype)
   out: Tensor = op[0](ta)
+  sched = out.schedule()
+  ast = sched[-1].ast
+  run_schedule(sched)
   tensor_value = out.numpy()
   numpy_value = op[1](ta.numpy())
-  if dtype in dtypes.floats:
-    atol, rtol = {dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 1e-2)}.get(dtype, (1e-6, 1e-5))
-    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
+  if dtype in (dtypes.float16, dtypes.bfloat16): np.testing.assert_allclose(tensor_value, numpy_value, atol=1e-3, rtol=1e-2)
+  elif dtype in dtypes_float: np.testing.assert_allclose(tensor_value, numpy_value, atol=1e-6, rtol=1e-5)
   else: np.testing.assert_equal(tensor_value, numpy_value)
+  if op[0] != Tensor.reciprocal: # reciprocal is not supported in most backends
+    op = [x for x in ast.toposort() if x.op in GroupOp.Unary][0]
+    assert op.dtype == dtype
 
 def universal_test_cast(a, in_dtype, dtype):
   tensor_value = Tensor([a], dtype=in_dtype).cast(dtype)
@@ -92,44 +99,45 @@ def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType):
   np.testing.assert_allclose(tensor_value, numpy_value, rtol=1e-6 if getenv("PTX") else 1e-7)
 
 class TestDTypeALU(unittest.TestCase):
-  @unittest.skipUnless(is_dtype_supported(dtypes.float64), f"no float64 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float64, Device.DEFAULT), f"no float64 on {Device.DEFAULT}")
   @given(ht.float64, ht.float64, strat.sampled_from(binary_operations))
   def test_float64(self, a, b, op): universal_test(a, b, dtypes.float64, op)
 
   @given(ht.float32, ht.float32, strat.sampled_from(binary_operations))
   def test_float32(self, a, b, op): universal_test(a, b, dtypes.float32, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.float16), f"no float16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16, Device.DEFAULT), f"no float16 on {Device.DEFAULT}")
   @given(ht.float16, ht.float16, strat.sampled_from(binary_operations))
   def test_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16, Device.DEFAULT), f"no bfloat16 on {Device.DEFAULT}")
   @given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
   def test_bfloat16(self, a, b, op): universal_test(a, b, dtypes.bfloat16, op)
 
   @given(ht.float32, strat.sampled_from(unary_operations))
   def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.float16), f"no float16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16, Device.DEFAULT), f"no float16 on {Device.DEFAULT}")
   @given(ht.float16, strat.sampled_from(unary_operations))
   def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16, Device.DEFAULT), f"no bfloat16 on {Device.DEFAULT}")
   @given(ht.bfloat16, strat.sampled_from(unary_operations))
+  @unittest.skipIf(Device.DEFAULT in ["AMD"], "broken on AMD?")
   def test_bfloat16_unary(self, a, op): universal_test_unary(a, dtypes.bfloat16, op)
 
   @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
   def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.uint16), f"no uint16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint16, Device.DEFAULT), f"no uint16 on {Device.DEFAULT}")
   @given(ht.uint16, ht.uint16, strat.sampled_from(integer_binary_operations))
   def test_uint16(self, a, b, op): universal_test(a, b, dtypes.uint16, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.uint32), f"no uint32 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint32, Device.DEFAULT), f"no uint32 on {Device.DEFAULT}")
   @given(ht.uint32, ht.uint32, strat.sampled_from(integer_binary_operations))
   def test_uint32(self, a, b, op): universal_test(a, b, dtypes.uint32, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.uint64), f"no uint64 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint64, Device.DEFAULT), f"no uint64 on {Device.DEFAULT}")
   @given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
   def test_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
 
@@ -142,7 +150,7 @@ class TestDTypeALU(unittest.TestCase):
   @given(ht.int32, ht.int32, strat.sampled_from(integer_binary_operations))
   def test_int32(self, a, b, op): universal_test(a, b, dtypes.int32, op)
 
-  @unittest.skipUnless(is_dtype_supported(dtypes.int64), f"no int64 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.int64, Device.DEFAULT), f"no int64 on {Device.DEFAULT}")
   @given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
   def test_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
 
@@ -172,7 +180,7 @@ class TestDTypeALU(unittest.TestCase):
   @settings(suppress_health_check=[HealthCheck.filter_too_much])
   @given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned(self, a, float_dtype, unsigned_dtype):
-    if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
+    if not is_dtype_supported(float_dtype, Device.DEFAULT): float_dtype = dtypes.float32
     float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
     float_strat = float_strat.filter(lambda x: 0 < x < dtypes.max(unsigned_dtype))
     universal_test_cast(a.draw(float_strat), float_dtype, unsigned_dtype)
@@ -180,7 +188,7 @@ class TestDTypeALU(unittest.TestCase):
   @settings(suppress_health_check=[HealthCheck.filter_too_much])
   @given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned_overflow(self, a, float_dtype, unsigned_dtype):
-    if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
+    if not is_dtype_supported(float_dtype, Device.DEFAULT): float_dtype = dtypes.float32
     float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
     overflow_strat = float_strat.filter(lambda x: x > dtypes.max(unsigned_dtype) and x <= dtypes.max(dtypes.int32))
     universal_test_cast(a.draw(overflow_strat), float_dtype, unsigned_dtype)
@@ -188,7 +196,7 @@ class TestDTypeALU(unittest.TestCase):
   @settings(suppress_health_check=[HealthCheck.filter_too_much])
   @given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned_underflow(self, a, float_dtype, unsigned_dtype):
-    if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
+    if not is_dtype_supported(float_dtype, Device.DEFAULT): float_dtype = dtypes.float32
     float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
     underflow_strat = float_strat.filter(lambda x: x < 0 and x >= dtypes.min(dtypes.int32))
     universal_test_cast(a.draw(underflow_strat), float_dtype, unsigned_dtype)

test/test_linearizer.py

@@ -12,7 +12,7 @@ from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program
 from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import prod, Context, getenv, CI, flatten, dedup, AMX, AMD_LLVM
-from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace
+from tinygrad.dtype import DType, dtypes, AddrSpace
 from tinygrad.codegen import apply_rewrites, rewrites_for_views
 
 def push_views(ast): return apply_rewrites(ast, rewrites_for_views)
@@ -33,10 +33,11 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
   r = a.matmul(b, dtype=dtype_out)
   if dtype_in == dtypes.bfloat16: r = r.float()
   realized_ast, bufs = helper_realized_ast(r)
-  opts = [Opt(op=OptOps.TC, axis=axis, arg=(tc_select, tc_opt, use_tensor_cores))]
-  prg = CompiledRunner(replace(get_program(realized_ast, opts=opts), device=Device.DEFAULT))
+  k = Kernel(realized_ast)
+  k.apply_tensor_cores(use_tensor_cores, axis=axis, tc_select=tc_select, tc_opt=tc_opt)
+  prg = CompiledRunner(replace(get_program(k.get_optimized_ast(), k.opts), device=Device.DEFAULT))
   if use_tensor_cores == 1: assert len([uop for uop in prg.p.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
-  assert len([x for x in prg.p.uops[-1].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
+  assert len([x for x in k.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
   prg.exec(bufs)
   if dtype_in == dtypes.half: tc_atol, tc_rtol = 1e-2, 1e-3
   elif dtype_in == dtypes.bfloat16: tc_atol, tc_rtol = 1e-2, 1e-2
@@ -116,7 +117,6 @@ class TestLinearizer(unittest.TestCase):
         if skip and i in skip: continue
         assert ranges[i-1] != u, f"multireduce nested the ranges! {ranges[i-1], {u}}"
 
-  @unittest.skip("broken. should not depends on push_views and implementation details of getitem")
   @unittest.skipIf(CI and Device.DEFAULT in {"PTX", "AMD", "NV"}, "very slow")
   def test_indexing_multireduce(self):
     dataset = Tensor.rand(16384, 256).realize()
@@ -133,7 +133,7 @@ class TestLinearizer(unittest.TestCase):
     uops = get_program(lin.get_optimized_ast(), lin.opts).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
-    # RANGE -> LOAD -> RANGE -> STORE
+    # RANGE -> LOAD -> RANGE -> ASSIGN
     #assert any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]])
 
   def test_three_nested_range(self):
@@ -143,7 +143,7 @@ class TestLinearizer(unittest.TestCase):
     uops = get_program(lin.get_optimized_ast(), lin.opts).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
-    # RANGE -> RANGE -> LOAD -> RANGE -> STORE
+    # RANGE -> RANGE -> LOAD -> RANGE -> ASSIGN
     # NOTE: nothing should toposort between the first two ranges
     #assert ranges[0]+1 == ranges[1]
     #assert any(x.op is Ops.LOAD for x in uops[ranges[1]:ranges[2]])
@@ -154,7 +154,7 @@ class TestLinearizer(unittest.TestCase):
     lin = helper_linearizer_opt(out, wanna_output=[24])[0]
     uops = get_program(lin.get_optimized_ast(), lin.opts).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
-    # RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
+    # RANGE -> ALU -> RANGE -> ALU + LOAD -> ASSIGN
     assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
     assert not any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]])
     assert any(x.op in {*GroupOp.ALU, Ops.LOAD} for x in uops[ranges[1]:])
@@ -166,7 +166,7 @@ class TestLinearizer(unittest.TestCase):
     lin = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])[0]
     uops = get_program(lin.get_optimized_ast(), lin.opts).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
-    # LOAD -> RANGE -> LOAD -> STORE
+    # LOAD -> RANGE -> LOAD -> ASSIGN
     assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
 
   def test_range_outer_op_before_phi_nested_range(self):
@@ -178,11 +178,11 @@ class TestLinearizer(unittest.TestCase):
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
     #if getenv("PTX"):
-    # LOAD -> RANGE -> CAST -> ALU -> ALU -> LOAD -> ALU -> RANGE -> ALU -> STORE
+    # LOAD -> RANGE -> CAST -> ALU -> ALU -> LOAD -> ALU -> RANGE -> ALU -> ASSIGN
     #  assert uops[ranges[0]-2].op is Ops.LOAD
     #  assert ranges[1] == ranges[0]+6
     #  assert [x.op for x in uops[ranges[1]-2:ranges[1]]] == [Ops.LOAD, Ops.ALU]
-    # LOAD -> RANGE -> LOAD -> ALU -> RANGE -> STORE
+    # LOAD -> RANGE -> LOAD -> ALU -> RANGE -> ASSIGN
     #else:
     #  assert uops[ranges[0]-2].op is Ops.LOAD
     #  assert ranges[1] == ranges[0]+3
@@ -194,7 +194,7 @@ class TestLinearizer(unittest.TestCase):
     out = a.sum() * a.sum()
     lin = helper_linearizer_opt(out, wanna_output=[a.numpy().sum()*a.numpy().sum()])[0]
     uops = get_program(lin.get_optimized_ast(), lin.opts).uops
-    # RANGE -> LOAD -> STORE -> ALU
+    # RANGE -> LOAD -> ASSIGN -> ALU
     end = max(i for i,u in enumerate(uops) if u.op is Ops.ENDRANGE)
     # the INDEX can be first
     assert uops[end+1].op in GroupOp.ALU or uops[end+2].op in GroupOp.ALU
@@ -205,7 +205,7 @@ class TestLinearizer(unittest.TestCase):
     out = a.reshape(2, 1).expand(2, 3).sum() + a.reshape(2, 1).expand(2, 3).sum()
     lin = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3))).sum()*2])[0]
     uops = get_program(lin.get_optimized_ast(), lin.opts).uops
-    # RANGE -> LOAD -> STORE -> ALU
+    # RANGE -> LOAD -> ASSIGN -> ALU
     end = max(i for i,u in enumerate(uops) if u.op is Ops.ENDRANGE)
     # the INDEX can be first
     assert uops[end+1].op in GroupOp.ALU or uops[end+2].op in GroupOp.ALU
@@ -327,7 +327,11 @@ class TestLinearizer(unittest.TestCase):
       n, m, k = tc.dims[0], tc.dims[1], 2 if AMX else tc.dims[2]
       a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in)
       r = a.matmul(b, dtype=tc.dtype_out)
-      prg = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))])
+      sched = r.schedule()
+      realized_ast = push_views(sched[-1].ast)
+      kernel = Kernel(realized_ast)
+      kernel.apply_tensor_cores(1, axis=0, tc_select=-1, tc_opt=2)
+      prg = get_program(kernel.get_optimized_ast(), kernel.opts)
       if Device.DEFAULT == "LLVM":
         assert "0x201000" in prg.src
       elif Device.DEFAULT == "AMD" and AMD_LLVM:
@@ -348,7 +352,7 @@ class TestLinearizer(unittest.TestCase):
   # Internal bug: zero-stride dimensions combined with a mask may produce wrong index/valid for pad == 1 on AMD
   @unittest.skipUnless((Device.DEFAULT == "AMD") or (Device.DEFAULT == "PYTHON" and getenv("EMULATE_AMD")), "test for AMD's tc")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
-  @unittest.skip("warp elements not duplicated properly across lanes")
+  @unittest.expectedFailure
   def test_tensor_cores_padded_amd(self):
     for tc in Device[Device.DEFAULT].renderer.tensor_cores:
       if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue
@@ -418,9 +422,9 @@ class TestLinearizer(unittest.TestCase):
     x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
     r = x.matmul(y, dtype=tc.dtype_out)
     k = helper_linearizer_opt(r, [[Opt(OptOps.UNROLL, 0, 4)]], apply_tc=True, atol=3e-2, rtol=1e-3)[-1]
-    for u in get_program(k.ast, k.opts, k.applied_opts).uops:
+    for u in get_program(k.get_optimized_ast(), k.opts).uops:
       if u.op is Ops.WMMA:
-        assert u.src[-1].src[0].op != Ops.STORE
+        assert u.src[-1].src[0].op != Ops.ASSIGN
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
   @unittest.skipIf(Device.DEFAULT in {"CPU", "LLVM"}, "CPU does not support using a different type for accumulation")
@@ -429,43 +433,37 @@ class TestLinearizer(unittest.TestCase):
     x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
     r = x.matmul(y, dtype=tc.dtype_out)
     k = helper_linearizer_opt(r, [[Opt(OptOps.UNROLL, 0, 4)]], apply_tc=True, atol=3e-2, rtol=1e-3)[-1]
-    for u in get_program(k.ast, k.opts, k.applied_opts).uops:
+    for u in get_program(k.get_optimized_ast(), k.opts).uops:
       if u.op is Ops.WMMA:
         #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
-        assert u.src[-1].src[0].op != Ops.STORE
+        assert u.src[-1].src[0].op != Ops.ASSIGN
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
   @unittest.skipIf(Device.DEFAULT in {"CPU", "LLVM"}, "CPU does not support using a different type for accumulation")
   def test_tensor_cores_unroll_casted_phi_with_children(self):
-    # all STORE children are outside the loop
+    # all ASSIGN children are outside the loop
     tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out][0]
     x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
     r = x.matmul(y, dtype=tc.dtype_out).relu()
     k = helper_linearizer_opt(r, [[Opt(OptOps.UNROLL, 0, 4)]], apply_tc=True, atol=3e-2, rtol=1e-3)[-1]
-    for u in get_program(k.ast, k.opts, k.applied_opts).uops:
+    for u in get_program(k.get_optimized_ast(), k.opts).uops:
       if u.op is Ops.WMMA:
         #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
-        assert u.src[-1].src[0].op != Ops.STORE
+        assert u.src[-1].src[0].op != Ops.ASSIGN
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
   def test_simple_unroll_no_between_phi_dependencies(self):
     x, y = Tensor.rand(128, 128), Tensor.rand(128, 128)
     r = (x@y).relu()
     k = helper_linearizer_opt(r, [[Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]])[-1]
-    # the uops graph is DEFINE_REG -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
-    uops = get_program(k.ast, k.opts, k.applied_opts).uops
-    begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
-    end_range = [i for i, x in enumerate(uops) if x.op is Ops.ENDRANGE][0]
-    for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
+    # the uops graph is RANGE -> DEFINE_ACC -> 4x ALU -> 4x ASSIGN -> ENDRANGE
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
     for u in uops:
-      if u.op is Ops.STORE and isinstance(dt:=u.src[0].dtype, PtrDType) and dt.addrspace is AddrSpace.REG:
-        if uops.index(u) < begin_range:
-          assert u.src[1].op is Ops.CONST
-        else:
-          assert u.src[1].op in GroupOp.ALU
-          assert begin_range < uops.index(u) < end_range
-      # children of STORE are placed after ENDRANGE
-      if any(x.op is Ops.STORE and x.src[1].op in GroupOp.ALU for x in u.src):
+      if u.op is Ops.ASSIGN:
+        assert u.src[1].op in GroupOp.ALU
+      # children of ASSIGN are placed after ENDRANGE
+      if any(x.op is Ops.ASSIGN for x in u.src):
+        end_range = [i for i, x in enumerate(uops) if x.op is Ops.ENDRANGE][0]
         assert end_range < uops.index(u)
 
   def test_grouped_dims(self):
@@ -544,7 +542,7 @@ class TestLinearizer(unittest.TestCase):
     # shrink so that the dims do not collapse
     t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
     k = helper_linearizer_opt(t+1)[0]
-    uops = get_program(k.ast, k.opts, k.applied_opts).uops
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
     idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
     idxs = sorted(idxs, key=lambda uop: uop.arg[0])
     assert idxs[0].arg == ('gidx0', 6), idxs[0].arg
@@ -584,13 +582,12 @@ class TestLinearizer(unittest.TestCase):
   def test_phi_simplification(self):
     def helper(t, max_ops=0):
       k = helper_linearizer_opt(t)[-1]
-      uops = get_program(k.ast, k.opts, k.applied_opts).uops
+      uops = get_program(k.get_optimized_ast(), k.opts).uops
       # ignore kernel optimized IF statements for now
       if if_op:=next((u for u in uops if u.op is Ops.IF), None):
         uops = uops[:uops.index(if_op)]
       assert len(set([u.op for u in uops if u.op in {Ops.RANGE, Ops.SPECIAL}])) == 1, "has either specials or ranges, not both"
-      reg_stores = [u for u in uops if u.op is Ops.STORE and isinstance(dt:=u.src[0].dtype, PtrDType) and dt.addrspace == AddrSpace.REG]
-      assert len(reg_stores) == 0, "STORE to reg should have been simplified"
+      assert len([u for u in uops if u.op is Ops.ASSIGN]) == 0, "ASSIGN should have been simplified"
       # TODO: once uops track min/max this will be fixed
       #assert len([u for u in uops if u.op is Ops.MAX]) <= max_ops, "no unnecessary MAX ops"
 
@@ -616,7 +613,7 @@ class TestLinearizer(unittest.TestCase):
     x, y = Tensor.randn(64,64), Tensor.randn(64,64)
     out = x.matmul(y)
     k = helper_linearizer_opt(out)[-1]
-    uops = get_program(k.ast, k.opts, k.applied_opts).uops
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
     # check that the float4 cast collapses
     store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
     for val in store_vals:
@@ -641,7 +638,7 @@ class TestLinearizer(unittest.TestCase):
     x = Tensor.randn((4,3,6,6)).realize()
     out = x.flip((0,1)).contiguous()
     k = helper_linearizer_opt(out)[-1]
-    store_val = [u.src[1] for u in get_program(k.ast, k.opts, k.applied_opts).uops if u.op is Ops.STORE][0]
+    store_val = [u.src[1] for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
     assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.VECTORIZE
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -654,7 +651,7 @@ class TestLinearizer(unittest.TestCase):
             Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
     k = helper_linearizer_opt(out, opts=[opt])[-1]
     def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
-    uops = get_program(k.ast, k.opts, k.applied_opts).uops
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
     local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
     global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_GLOBAL for x in get_recursive(u.src[0]))]
     barrier = [u for u in uops if u.op is Ops.BARRIER][0]
@@ -674,7 +671,7 @@ class TestLinearizer(unittest.TestCase):
     x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
     r = (x@y).relu()
     k = helper_linearizer_opt(r)[-1]
-    uops = get_program(k.ast, k.opts, k.applied_opts).uops
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
     stores = [u for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
 
     # the float4 value stores directly in lds and we skip upcast
@@ -700,7 +697,7 @@ class TestLinearizer(unittest.TestCase):
         Opt(op=OptOps.LOCAL, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=3, arg=2)
     ]
     k = helper_linearizer_ast(ast, [Tensor.randn(240*40).realize()], opts=[opt])[-1]
-    out = [u for u in get_program(k.ast, k.opts, k.applied_opts).uops if u.op is Ops.STORE][0]
+    out = [u for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
     assert out.src[1].op is Ops.VECTORIZE and out.src[1].dtype == dtypes.float.vec(4)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -718,7 +715,7 @@ class TestLinearizer(unittest.TestCase):
             Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8),
             Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=2)]
     k = helper_linearizer_ast(ast, [Tensor.randn(8*32).realize()], opts=[opt])[-1]
-    out = [u for u in get_program(k.ast, k.opts, k.applied_opts).uops if u.op is Ops.STORE][0]
+    out = [u for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
     assert out.src[1].op is Ops.VECTORIZE and out.src[1].dtype.count != 1
 
 @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need backends that support float4")
@@ -1049,7 +1046,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
   outbufs = [real_bufs[x.src[0].base.arg] for x in realized_ast.src]
   device = real_bufs[0].device
 
-  def get_prg(k:Kernel): return CompiledRunner(replace(get_program(k.ast, k.opts, k.applied_opts), device=device))
+  def get_prg(k:Kernel): return CompiledRunner(replace(get_program(k.get_optimized_ast(), k.opts), device=device))
 
   def check_opt(opts, create_k, expected_color_size):
     k = create_k()

test/test_linearizer_dumb.py

@@ -9,6 +9,7 @@ from tinygrad.uop.ops import UOp, Ops
 from tinygrad.helpers import getenv
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.codegen.opt.search import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.engine.realize import get_program
 
 class TestLinearizerDumb(unittest.TestCase):
@@ -35,7 +36,9 @@ class TestLinearizerDumb(unittest.TestCase):
           UOp(Ops.CONST, dtypes.half, arg=0.0, src=(
              x16,)),)),)),))
     opts = [Opt(op=OptOps.TC, axis=2, arg=(-1, 2, 1)), Opt(op=OptOps.UPCAST, axis=2, arg=0), Opt(op=OptOps.UNROLL, axis=1, arg=0)]
-    prg = get_program(ast, Device["METAL"].renderer, opts)
+    k = Kernel(ast, opts=Device["METAL"].renderer)
+    k.apply_opts(opts)
+    prg = get_program(k.get_optimized_ast(), k.opts)
     print(prg.src)
     Device[Device.DEFAULT].compiler.compile_cached(prg.src)
     gate_count = len([x for x in prg.src.splitlines() if "if" in x])
@@ -72,7 +75,9 @@ class TestLinearizerDumb(unittest.TestCase):
             UOp(Ops.CONST, dtypes.int, arg=1000, src=(
                x14,)),)),)),)),))
     opts = [Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
+    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
+    k.apply_opts(opts)
+    prg = get_program(k.get_optimized_ast(), k.opts)
     print(prg.src)
     assert prg.uops is not None and not any(uop.op is Ops.MAX for uop in prg.uops), "leftover MAX"
 
@@ -88,7 +93,9 @@ class TestLinearizerDumb(unittest.TestCase):
             UOp(Ops.VIEW, dtypes.float.ptr(25), arg=ShapeTracker(views=(View(shape=(26, 49), strides=(0, -1), offset=48, mask=((0, 26), (24, 49)), contiguous=False), View(shape=(25, 25), strides=(1, 50), offset=0, mask=None, contiguous=False))), src=(
               UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(25), arg=1, src=()),)),)),)),)),))
     opts = [Opt(op=OptOps.GROUP, axis=0, arg=0), Opt(op=OptOps.PADTO, axis=0, arg=32), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=0)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
+    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
+    k.apply_opts(opts)
+    prg = get_program(k.get_optimized_ast(), k.opts)
     print(prg.src)
     if_uops = [u for u in prg.uops if u.op is Ops.IF]
     self.assertIn(len(if_uops), {1,2,3})
@@ -128,7 +135,8 @@ class TestLinearizerDumb(unittest.TestCase):
                 UOp(Ops.LOAD, dtypes.half, arg=None, src=(
                   UOp(Ops.VIEW, dtypes.half.ptr(131072000), arg=ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(1, 4096, 0), offset=0, mask=None, contiguous=False),)), src=(
                     UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(131072000), arg=2, src=()),)),)),)),)),)),)),)),))
-    prg = get_program(ast, Device[Device.DEFAULT].renderer)
+    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
+    prg = get_program(k.get_optimized_ast(), k.opts)
     print(prg.src)
 
   @unittest.expectedFailure
@@ -155,9 +163,11 @@ class TestLinearizerDumb(unittest.TestCase):
               UOp(Ops.VIEW, dtypes.float.ptr(18), arg=ShapeTracker(views=(View(shape=(3, 6), strides=(6, 1), offset=0, mask=None, contiguous=True),)), src=(
                 UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(18), arg=2, src=()),)),)),)),)),)),))
     opts = [Opt(op=OptOps.UNROLL, axis=0, arg=0)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
+    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
+    k.apply_opts(opts)
+    prg = get_program(k.get_optimized_ast(), k.opts)
     print(prg.src)
-    load_idxs = [x.src[1] for x in prg.uops if x.op is Ops.LOAD and x.src[0].arg == 2]
+    load_idxs = [x.src[1] for x in k.uops if x.op is Ops.LOAD and x.src[0].arg == 2]
     assert load_idxs[0] < load_idxs[1], f"first loaded idx {load_idxs[0].arg} then {load_idxs[1].arg}!"
 
   @unittest.expectedFailure
@@ -177,9 +187,11 @@ class TestLinearizerDumb(unittest.TestCase):
               UOp(Ops.VIEW, dtypes.float.ptr(1040), arg=ShapeTracker(views=(View(shape=(4, 5, 13, 1, 1, 1, 4, 1, 4, 3, 3), strides=(260, 13, 1, 0, 0, 0, 65, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=(
                 UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1040), arg=2, src=()),)),)),)),)),)),))
     opts = [Opt(op=OptOps.UPCAST, axis=3, arg=0), Opt(op=OptOps.UPCAST, axis=2, arg=0)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
+    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
+    k.apply_opts(opts)
+    prg = get_program(k.get_optimized_ast(), k.opts)
     print(prg.src)
-    store_idxs = [x.src[1] for x in prg.uops if x.op is Ops.STORE]
+    store_idxs = [x.src[1] for x in k.uops if x.op is Ops.STORE]
     for i in range(len(store_idxs) - 1):
       first_bounds = store_idxs[i].vmin+store_idxs[i].vmax
       next_bounds = store_idxs[i+1].vmin+store_idxs[i+1].vmax

test/test_memory_planner.py

@@ -120,19 +120,5 @@ class TestMemoryPlanner(unittest.TestCase):
     ]
     check_assign(bs)
 
-  def test_very_small_buffers(self):
-    bs = [
-      [b(0, pin=True), b(1, size=32)],
-      [b(3, size=4), b(4, size=6)],
-    ]
-    check_assign(bs)
-
-  def test_very_big_buffers(self):
-    bs = [
-      [b(0, pin=True), b(1, size=34359738368000)],
-      [b(3, size=1 << 128), b(4, size=1 << 64)],
-    ]
-    check_assign(bs)
-
 if __name__ == "__main__":
   unittest.main()

test/test_multitensor.py

@@ -373,11 +373,11 @@ class TestMultiTensor(unittest.TestCase):
     np.testing.assert_allclose(y.numpy(), y_shard.numpy(), atol=1e-6, rtol=1e-6)
 
   # NOTE: this is failing on LLVM CI, no idea why. Works locally.
-  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "LLVM", "CPU", "AMD"), "slow, and flaky on LLVM/CPU")
+  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "LLVM", "CPU"), "slow, and flaky on LLVM/CPU")
   def test_data_parallel_resnet(self):
     from extra.models.resnet import ResNet18
 
-    fake_image = Tensor.rand((2, 3, 224//16, 224//16))
+    fake_image = Tensor.rand((2, 3, 224//8, 224//8))
     fake_image_sharded = fake_image.shard(devices_2, axis=0)
     m = ResNet18()
     m.load_from_pretrained()
@@ -409,10 +409,10 @@ class TestMultiTensor(unittest.TestCase):
     # sometimes there is zeros in these grads... why?
     np.testing.assert_allclose(grad, shard_grad, atol=1e-5, rtol=1e-5)
 
-  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "LLVM", "CPU", "AMD"), "slow, and flaky on LLVM/CPU")
+  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "LLVM", "CPU"), "slow, and flaky on LLVM/CPU")
   def test_data_parallel_resnet_train_step(self):
     from extra.models.resnet import ResNet18
-    fake_image = Tensor.rand((2, 3, 224//16, 224//16))
+    fake_image = Tensor.rand((2, 3, 224//8, 224//8))
     labels = Tensor.randint(2, low=0, high=1000)
     m = ResNet18()
     self._test_model_train_step(m, fake_image, labels)
@@ -1128,7 +1128,6 @@ class TestMultiRamUsage(unittest.TestCase):
     self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
   def test_zeros_shard_self(self): self.test_zeros_shard((d0, d1))
 
-  @unittest.skip("flaky")
   def test_zeros_contiguous_shard(self):
     _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).contiguous().realize()
     self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage

test/test_net_speed.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python
+import time
+import unittest
+import torch
+from tinygrad import Tensor, Device
+from tinygrad.helpers import Profiling, CI
+
+@unittest.skipIf(CI and Device.DEFAULT in {"CUDA", "NV"}, "slow")
+class TestConvSpeed(unittest.TestCase):
+
+  def test_mnist(self):
+    # https://keras.io/examples/vision/mnist_convnet/
+    conv = 3
+    inter_chan, out_chan = 32, 64
+
+    # ****** torch baseline *******
+
+    torch.backends.mkldnn.enabled = False
+
+    conv = 3
+    inter_chan, out_chan = 32, 64
+    c1 = torch.randn(inter_chan,1,conv,conv, requires_grad=True)
+    c2 = torch.randn(out_chan,inter_chan,conv,conv, requires_grad=True)
+    l1 = torch.randn(out_chan*5*5, 10, requires_grad=True)
+
+    c2d = torch.nn.functional.conv2d
+    mp = torch.nn.MaxPool2d((2,2))
+    lsm = torch.nn.LogSoftmax(dim=1)
+
+    cnt = 5
+    fpt, bpt = 0.0, 0.0
+    for i in range(cnt):
+      et0 = time.time()
+      x = torch.randn(128, 1, 28, 28, requires_grad=True)
+      x = mp(c2d(x,c1).relu())
+      x = mp(c2d(x,c2).relu())
+      x = x.reshape(x.shape[0], -1)
+      out = lsm(x.matmul(l1))
+      out = out.mean()
+      et1 = time.time()
+      out.backward()
+      et2 = time.time()
+      fpt += (et1-et0)
+      bpt += (et2-et1)
+
+    fpt_baseline = (fpt*1000/cnt)
+    bpt_baseline = (bpt*1000/cnt)
+    print("torch forward pass:  %.3f ms" % fpt_baseline)
+    print("torch backward pass: %.3f ms" % bpt_baseline)
+
+    # ****** tinygrad compare *******
+
+    c1 = Tensor(c1.detach().numpy(), requires_grad=True)
+    c2 = Tensor(c2.detach().numpy(), requires_grad=True)
+    l1 = Tensor(l1.detach().numpy(), requires_grad=True)
+
+    cnt = 5
+    fpt, bpt = 0.0, 0.0
+    for i in range(1+cnt):
+      et0 = time.time()
+      x = Tensor.randn(128, 1, 28, 28)
+      x = x.conv2d(c1).relu().avg_pool2d()
+      x = x.conv2d(c2).relu().max_pool2d()
+      x = x.reshape(shape=(x.shape[0], -1))
+      out = x.dot(l1).log_softmax()
+      out = out.mean()
+      out.backward()  # NOTE: we have to now compute this here, but it doesn't realize
+      out.realize()
+      et1 = time.time()
+      [x.grad.realize() for x in [c1, c2, l1]]
+      et2 = time.time()
+      if i == 0:
+        pr = Profiling(sort='time', frac=0.2)
+        pr.__enter__()
+      else:
+        fpt += (et1-et0)
+        bpt += (et2-et1)
+
+    pr.__exit__()
+    fpt = (fpt*1000/cnt)
+    bpt = (bpt*1000/cnt)
+    print("forward pass:  %.3f ms, %.2fx off baseline %.3f ms" % (fpt, fpt/fpt_baseline, fpt_baseline))
+    print("backward pass: %.3f ms, %.2fx off baseline %.3f ms" % (bpt, bpt/bpt_baseline, bpt_baseline))
+
+
+if __name__ == '__main__':
+  unittest.main()

test/test_nn.py

@@ -210,27 +210,6 @@ class TestNN(unittest.TestCase):
       np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
       np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
 
-  def test_layernorm_forward(self):
-    N, C, H, W = 20, 5, 10, 10
-
-    # create in torch
-    torch_layer = torch.nn.LayerNorm([H, W]).eval()
-
-    # create in tinygrad
-    layer = LayerNorm([H, W])
-    layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True)
-    layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True)
-
-    x = Tensor.empty(N, C, H, W, requires_grad=True)
-    z = layer(x)
-    z.realize()
-
-    torch_x = torch.tensor(x.numpy(), requires_grad=True)
-    torch_z = torch_layer(torch_x)
-    torch_z.sum().backward()
-
-    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
-
   def test_layernorm(self):
     N, C, H, W = 20, 5, 10, 10
 

test/test_ops.py

@@ -393,7 +393,7 @@ class TestOps(unittest.TestCase):
   def test_trunc(self):
     helper_test_op([()], lambda x: x.trunc(), forward_only=True)
     helper_test_op([(45,35)], lambda x: x.trunc(), forward_only=True)
-    helper_test_op(None, lambda x: x.trunc(), vals=[[1.499, 1.5, 1.501, 1.0, 2.1, 0.0, -5.0, -2.499, -2.5, -2.501, 1e12, -1e12]], forward_only=True)
+    helper_test_op(None, lambda x: x.trunc(), vals=[[1.499, 1.5, 1.501, 1.0, 2.1, 0.0, -5.0, -2.499, -2.5, -2.501]], forward_only=True)
   def test_floor(self):
     helper_test_op([()], lambda x: x.floor(), forward_only=True)
     helper_test_op([(45,35)], lambda x: x.floor(), forward_only=True)
@@ -928,12 +928,6 @@ class TestOps(unittest.TestCase):
       for j in [-1., 0., 1.]:
         helper_test_op(None, torch.copysign, Tensor.copysign, vals=[[i], [j]])
 
-  def test_logaddexp(self):
-    helper_test_op([(45,65), (45,65)], torch.logaddexp, Tensor.logaddexp)
-    helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[-1.], [-1.0, 2, 3]])
-    helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[-100.0, -200, -300], [-1.0, 2, 3]])
-    helper_test_op(None, torch.logaddexp, Tensor.logaddexp, vals=[[1.0, 2000, 30000], [-1.0, 2, 3]])
-
   def test_softsign(self):
     helper_test_op([(45,65)], torch.nn.functional.softsign, Tensor.softsign)
     helper_test_op([()], torch.nn.functional.softsign, Tensor.softsign)
@@ -971,6 +965,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
     helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3), lambda t: Tensor.softplus(t, beta=3), grad_atol=1e-6)
     helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=1/3), lambda t: Tensor.softplus(t, beta=1/3), grad_atol=1e-6)
+    helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3, threshold=0.5),
+                              lambda t: Tensor.softplus(t, beta=3, threshold=0.5), grad_atol=1e-6)
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=300, high=400)
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=-400, high=-300)
     helper_test_op([()], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
@@ -2465,20 +2461,6 @@ class TestOps(unittest.TestCase):
       lambda x: Tensor.max_unpool2d(*Tensor.max_pool2d(x, kernel_size=(2,2), return_indices=True),
                                     kernel_size=(2,2), output_size=(99,99,7,6)), forward_only=True)
 
-  def test_max_unpool2d_inf(self):
-    data = [[[[math.inf, -math.inf, math.nan], [1.0, 2.0, 3.0]]]]
-    ksz = (2,2)
-    helper_test_op((),
-      lambda: torch.nn.functional.max_unpool2d(
-        *torch.nn.functional.max_pool2d(torch.tensor(data), kernel_size=ksz, return_indices=True),
-        kernel_size=ksz
-      ),
-      lambda: Tensor.max_unpool2d(
-        *Tensor.max_pool2d(Tensor(data), kernel_size=ksz, return_indices=True),
-        kernel_size=ksz
-      ),
-      forward_only=True)
-
   def test_avg_pool2d(self):
     shape = (32,2,111,28)
     for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]:
@@ -2712,10 +2694,6 @@ class TestOps(unittest.TestCase):
     i, j, k, o, p = [Tensor(tor.detach().cpu().numpy().astype(np.int32), requires_grad=False) for tor in [a,b,c,d,e]]
     return a,b,c,d,e,i,j,k,o,p
 
-  def test_fancy_indexing_inf(self):
-    data = [math.inf, -math.inf, math.nan]
-    helper_test_op((), lambda: torch.tensor(data)[torch.tensor([0, 1, 2])], lambda: Tensor(data)[Tensor([0, 1, 2])])
-
   def test_slice_fancy_indexing_no_dim_collapse(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # no dim collapse from int or dim injection from None
@@ -2826,7 +2804,11 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: x.gather(dim=0, index=torch.tensor([2, 1, 0, 1, 2], requires_grad=False)),
                          lambda x: x.gather(dim=0, index=Tensor([2, 1, 0, 1, 2])),
                          vals=[[1., 2., 3.]])
-    # gather with inf values
+
+  @unittest.expectedFailure
+  @unittest.skipIf(torch._C._get_privateuse1_backend_name() == "tiny", 'results in a success instead of a failure')
+  def test_gather_failure(self):
+    # gather with inf values do not work, other values results in nan
     helper_test_op(None, lambda x: x.gather(dim=0, index=torch.tensor([2, 1, 0, 1, 2], requires_grad=False)),
                          lambda x: x.gather(dim=0, index=Tensor([2, 1, 0, 1, 2])),
                          vals=[[-float("inf"), 2., 3.]])

test/test_outerworld_range.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor, nn, Variable, UOp
+from tinygrad import Tensor, nn, Variable, UOp, dtypes
 
 # outerworld range should support three things
 # 1. full optimizer steps (test_model_bound_range)
@@ -136,7 +136,7 @@ class TestOuterworldRange(unittest.TestCase):
   def test_model_bound_range(self):
     m, opt = get_model_and_opt()
     # TODO: should ranges be unique so you don't have to pass in the -1?
-    rng = UOp.range(self.STEPS, -1)
+    rng = UOp.range(dtypes.int, self.STEPS, -1)
     vib = Variable('i', 0, self.STEPS-1).bind(rng)
     loss = (m(self.X[vib]) - self.Y[vib]).square().mean()
     loss.backward()

test/test_pickle.py

@@ -20,7 +20,7 @@ class TestPickle(unittest.TestCase):
     self.assertEqual(pm2.rewrite(sink).key, tt.key)
 
   def test_pickle_main_pattern_matcher(self):
-    from tinygrad.codegen.late.devectorizer import sym
+    from tinygrad.codegen.devectorizer import sym
     ssym = pickle.dumps(sym)
     dsym = pickle.loads(ssym)
     self.assertEqual(dsym.patterns[0][0].location, sym.patterns[0][0].location)

test/test_profiler.py

@@ -1,6 +1,6 @@
 import unittest, struct, contextlib, statistics, time, gc
 from tinygrad import Device, Tensor, dtypes, TinyJit
-from tinygrad.helpers import CI, getenv, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
+from tinygrad.helpers import CI, getenv, Context, ProfileRangeEvent, cpu_profile, cpu_events
 from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.runtime.support.hcq import HCQCompiled
 from tinygrad.engine.realize import get_runner
@@ -209,18 +209,5 @@ class TestProfiler(unittest.TestCase):
     for ge in graphs:
       self.assertEqual(len(ge.ents), len(graphs))
 
-  def test_trace_metadata(self):
-    with Context(TRACEMETA=1):
-      a = Tensor.empty(1)+2
-      b = Tensor.empty(1)+2
-      with helper_collect_profile(TestProfiler.d0) as profile:
-        Tensor.realize(a, b)
-    profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
-    exec_points = [e for e in profile if isinstance(e, ProfilePointEvent) and e.name == "exec"]
-    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent)]
-    self.assertEqual(len(exec_points), len(range_events), 2)
-    self.assertEqual(len(dedup(e.key for e in exec_points)), 1)
-    self.assertEqual(len(dedup(e.arg['metadata'] for e in exec_points)), 1)
-
 if __name__ == "__main__":
   unittest.main()

test/test_rangeify.py

@@ -1,196 +0,0 @@
-import unittest
-from tinygrad import Tensor
-from tinygrad.helpers import RANGEIFY, Context, GlobalCounters
-from tinygrad.uop.ops import UOp
-
-N = 256
-
-@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
-class TestRangeify(unittest.TestCase):
-  def test_expand_children(self):
-    A = Tensor.empty(N, N).sum(axis=1)
-    ba = A.expand(N, N)
-    ((ba+1).sum(axis=1) + (ba+2).sum(axis=0)).realize()
-
-  def test_partial_contig(self):
-    A = Tensor.empty(64, 64, 64)
-    ret = A.sum(axis=2).contiguous(arg=(1,)).sum(axis=1)
-    ret.realize()
-
-  def test_double_gemm_real(self):
-    def go():
-      with Context(DEBUG=0):
-        Tensor.manual_seed(1337)
-        A,B,C = [Tensor.randn(N, N) for _ in range(3)]
-        Tensor.realize(A, B, C)
-      GlobalCounters.reset()
-      return (A@B@C).realize()
-    rng = go()
-    with Context(RANGEIFY=0, DEBUG=2):
-      ref = go()
-      mse = ((rng-ref)**2).sum().item()
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-2)
-
-  def test_double_gemm(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (A@B@C).realize()
-
-  def test_double_gemm_exp(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (((A@B).exp()@C).exp()).realize()
-
-  def test_double_gemm_relu(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (((A@B).relu()@C).relu()).realize()
-
-  def test_double_gemm_relu_half_contig(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (((A@B).relu().contiguous(arg=(1,))@C).relu()).realize()
-
-  def test_double_gemm_half_contig(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    ((A@B).contiguous(arg=(1,))@C).realize()
-
-  def test_double_gemm_contig(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    ((A@B).contiguous()@C).realize()
-
-  def test_many_gemm(self):
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    D = Tensor.empty(N, N)
-    E = Tensor.empty(N, N)
-    F = Tensor.empty(N, N)
-    (A@B@C@D@E@F).realize()
-
-  def test_conv2d(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    x.conv2d(w1).realize()
-
-  def test_conv2d_t(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    (x*2).conv2d(w1).realize()
-
-  def test_double_conv2d(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    w2 = Tensor.empty(12, 8, 3, 3)
-    x.conv2d(w1).conv2d(w2).realize()
-
-  def test_double_conv2d_half_contig(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    w2 = Tensor.empty(12, 8, 3, 3)
-    # NOTE: this contiguous doesn't help
-    x.conv2d(w1).contiguous(arg=(1,)).conv2d(w2).permute(0,2,3,1).contiguous().realize()
-
-  def test_double_conv2d_contig(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    w2 = Tensor.empty(12, 8, 3, 3)
-    x.conv2d(w1).contiguous().conv2d(w2).realize()
-
-  def test_transformer_ffn(self):
-    from tinygrad.apps.llm import TransformerBlock
-    from tinygrad import nn
-    blk = TransformerBlock(1024, 4096, 1, 1, 1e-5)
-    for p in nn.state.get_parameters(blk): p.replace(Tensor.empty(p.shape))
-
-    x = Tensor.empty(128, 1024)
-    out = blk._feed_forward(x)
-    out.realize()
-
-  def test_flash_attention(self):
-    BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
-
-    # bigger
-    #BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
-
-    # llama 8B
-    #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
-
-    def fa():
-      Tensor.manual_seed(1337)
-      with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
-      return q.scaled_dot_product_attention(k, v).realize()
-
-    with Context(DEBUG=4):
-      GlobalCounters.reset()
-      ret = fa()
-    with Context(RANGEIFY=0):
-      with Context(DEBUG=2):
-        GlobalCounters.reset()
-        cmp = fa()
-      with Context(DEBUG=0):
-        mse = ((cmp-ret)**2).sum().item()
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-6)
-
-# contiguous + reduce can support ranges?
-
-@unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
-class TestOuterworld(unittest.TestCase):
-  def test_passthrough_range(self):
-    t = Tensor.rand(10, 10).realize()
-
-    # passthrough ranges
-    a = UOp.range(10, -1)
-    sel = t[a]
-    cpy = sel.contiguous(a).realize()
-
-    self.assertTrue((t==cpy).all().item())
-
-  def test_flip_range(self):
-    t = Tensor.rand(10, 10).realize()
-
-    # passthrough ranges
-    a = UOp.range(10, -1)
-    sel = t[9-a]
-    cpy = sel.contiguous(a).realize()
-
-    self.assertTrue((t.flip(0)==cpy).all().item())
-
-  def test_vmap(self):
-    def f(x): return x.sum(axis=0)*2
-
-    x = Tensor.ones(3, 10, 2).contiguous()
-
-    # vmap across axis 0
-    a = UOp.range(3, -1)
-    out = f(x[a])
-    out = out.contiguous(a)
-
-    # 3x2 grid of 20
-    out.realize()
-    print(out.numpy())
-
-  def test_triple_gemm(self):
-    x = Tensor.rand(1, 16).realize()
-    W = Tensor.rand(3, 16, 16).realize()
-
-    manual = (x @ W[0] @ W[1] @ W[2]).contiguous().realize()
-
-    a = UOp.range(3, -1)
-    x = x.assign(x @ W[a])
-    out = x.contiguous(a)[-1].contiguous().realize()
-
-    self.assertTrue((manual==out).all().item())
-
-if __name__ == '__main__':
-  unittest.main()

test/test_renderer_failures.py

@@ -25,8 +25,7 @@ def _test_uop_result(inputs:List[Tensor], stores:List[UOp], local_size=None):
       initial_value=np.zeros(sz, dtype=_to_np_dtype(dtype)).data) for u in uops if u.op is Ops.STORE]
   inbufs = [cast(UOp,x.uop).base.buffer for x in inputs]
   src = Device[Device.DEFAULT].renderer.render(uops)
-  ei = CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test",
-                                  src, Device.DEFAULT, uops[-1], uops=uops, local_size=local_size))
+  ei = CompiledRunner(ProgramSpec("test", src, Device.DEFAULT, uops[-1], uops=uops, local_size=local_size))
   ei.exec(outbufs+inbufs)
   return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]
 

test/test_schedule.py

@@ -1050,14 +1050,6 @@ class TestSchedule(unittest.TestCase):
       compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
       np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
 
-    with Context(FUSE_ATTENTION=1):
-      out = Tensor.scaled_dot_product_attention(q,k,v)
-      run_schedule(check_schedule(out, 1))
-      if getenv("CHECK", 1):
-        import torch
-        compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
-        np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
-
   def test_ugly_reduceop_pairing(self):
     Tensor.manual_seed(0)
     a = Tensor.randn(4, 32).realize()

test/test_softmax_fusion.py

@@ -117,14 +117,6 @@ class TestFuse(unittest.TestCase):
     c = (a.sum(axis=1) + b.sum(axis=1)).fuse()
     self.assertListEqual(c.tolist(), [30]*16)
 
-  @unittest.skipUnless(Device.DEFAULT == "METAL", "METAL TC")
-  def test_fuse_and_tc_opt(self):
-    A = Tensor.randn(8, 8).realize()
-    B = Tensor.randn(8, 8).realize()
-    C = Tensor.ones(1, 8, 8).pad(((1,1), None, None),).sum(0)
-    out = (C + (A @ B)).fuse()
-    out.realize()
-
 class TestSoftmaxFusion(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
@@ -163,7 +155,7 @@ class TestSoftmaxFusion(unittest.TestCase):
       out = single_kernel_softmax(self.test)
       out.realize()
 
-    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
+    np.testing.assert_allclose(sout.numpy(), out.numpy())
 
   def test_auto_softmax(self):
     print("*** softmax ***")
@@ -176,7 +168,7 @@ class TestSoftmaxFusion(unittest.TestCase):
       out = self.test.contiguous().softmax(-1).fuse()
       run_one_schedule_item(out)
 
-    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
+    np.testing.assert_allclose(sout.numpy(), out.numpy())
 
   @unittest.skip("recursion error no longer raised")
   def test_softmax_bw(self):

test/test_symbolic_jit.py

@@ -2,41 +2,50 @@ import unittest
 
 from test.helpers import assert_jit_cache_len
 from tinygrad import Variable, Tensor, TinyJit
+from tinygrad.helpers import Context
 import numpy as np
 
 class TestSymbolicJit(unittest.TestCase):
+  def setUp(self):
+    # A lot of these test are out of bounds, so we ignore the bounds check
+    self.context = Context(IGNORE_OOB=1)
+    self.context.__enter__()
+
+  def tearDown(self):
+    self.context.__exit__(None, None, None)
+
   def test_plus1(self):
     def f(a): return (a+1).realize()
     jf = TinyJit(f)
-    a = Tensor.rand(3, 10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(a[:, :vi]).reshape(3, i).numpy()
-      expected = f(a[:, :i]).numpy()
+      a = Tensor.rand(3, i)
+      symbolic = jf(a.reshape(3, vi)).reshape(3, i).numpy()
+      expected = f(a).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_add(self):
     def f(a, b): return (a+b).realize()
     jf = TinyJit(f)
-    a = Tensor.rand(3, 10)
-    b = Tensor.rand(3, 10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(a[:, :vi], b[:, :vi]).reshape(3, i).numpy()
-      expected = f(a[:, :i], b[:, :i]).numpy()
+      a = Tensor.rand(3, i)
+      b = Tensor.rand(3, i)
+      symbolic = jf(a.reshape(3, vi), b.reshape(3, vi)).reshape(3, i).numpy()
+      expected = f(a, b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_matmul(self):
     def f(a, b): return (a@b).realize()
     jf = TinyJit(f)
-    a = Tensor.rand(3, 10)
-    b = Tensor.rand(10, 5)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(a[:, :vi], b[:vi, :]).numpy()
-      expected = f(a[:, :i], b[:i, :]).numpy()
+      a = Tensor.rand(3, i)
+      b = Tensor.rand(i, 5)
+      symbolic = jf(a.reshape(3, vi), b.reshape(vi, 5)).numpy()
+      expected = f(a, b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
@@ -46,119 +55,119 @@ class TestSymbolicJit(unittest.TestCase):
       s = (s+s).realize() # this one does not have symbols in input
       return s
     jf = TinyJit(f)
-    a = Tensor.rand(3, 10)
-    b = Tensor.rand(10, 5)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(a[:, :vi], b[:vi, :]).numpy()
-      expected = f(a[:, :i], b[:i, :]).numpy()
+      a = Tensor.rand(3, i)
+      b = Tensor.rand(i, 5)
+      symbolic = jf(a.reshape(3, vi), b.reshape(vi, 5)).numpy()
+      expected = f(a, b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 2)
 
   def test_attention(self):
     def f(q, k, v): return Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).realize()
     jf = TinyJit(f)
-    q = Tensor.rand(2, 1, 4, 8)
-    k = Tensor.rand(2, 10, 4, 8)
-    v = Tensor.rand(2, 10, 4, 8)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(q, k[:, :vi], v[:, :vi]).reshape(2, 4, 1, 8).numpy()
-      expected = f(q, k[:, :i], v[:, :i]).numpy()
+      q = Tensor.rand(2, 1, 4, 8)
+      k = Tensor.rand(2, i, 4, 8)
+      v = Tensor.rand(2, i, 4, 8)
+      symbolic = jf(q, k.reshape(2, vi, 4, 8), v.reshape(2, vi, 4, 8)).reshape(2, 4, 1, 8).numpy()
+      expected = f(q, k, v).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 5)
 
   def test_cat_dim0(self):
     def f(a, b): return a.cat(b, dim=0).realize()
     jf = TinyJit(f)
-    a = Tensor.rand(10, 3)
-    b = Tensor.rand(2, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(a[:vi], b).reshape(i+2, 3).numpy()
-      expected = f(a[:i], b).numpy()
+      a = Tensor.rand(i, 3)
+      b = Tensor.rand(2, 3)
+      symbolic = jf(a.reshape(vi, 3), b).reshape(i+2, 3).numpy()
+      expected = f(a, b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_cat_dim1(self):
     def f(a, b): return a.cat(b, dim=1).realize()
     jf = TinyJit(f)
-    a = Tensor.rand(3, 10)
-    b = Tensor.rand(3, 2)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(a[:, :vi], b).reshape(3, i+2).numpy()
-      expected = f(a[:, :i], b).numpy()
+      a = Tensor.rand(3, i)
+      b = Tensor.rand(3, 2)
+      symbolic = jf(a.reshape(3, vi), b).reshape(3, i+2).numpy()
+      expected = f(a, b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_cat_dim0_two_vars(self):
     def f(a, b): return a.cat(b, dim=0).realize()
     jf = TinyJit(f)
-    a = Tensor.rand(10, 3)
-    b = Tensor.rand(10, 3)
     for i in range(1, 5):
       for j in range(1, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        symbolic = jf(a[:vi], b[:vj]).reshape(i+j, 3).numpy()
-        expected = f(a[:i], b[:j]).numpy()
+        a = Tensor.rand(i, 3)
+        b = Tensor.rand(j, 3)
+        symbolic = jf(a.reshape(vi, 3), b.reshape(vj, 3)).reshape(i+j, 3).numpy()
+        expected = f(a, b).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_cat_dim1_two_vars(self):
     def f(a, b): return a.cat(b, dim=1).realize()
     jf = TinyJit(f)
-    a = Tensor.rand(3, 10)
-    b = Tensor.rand(3, 10)
     for i in range(1, 5):
       for j in range(1, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        symbolic = jf(a[:, :vi], b[:, :vj]).reshape(3, i+j).numpy()
-        expected = f(a[:, :i], b[:, :j]).numpy()
+        a = Tensor.rand(3, i)
+        b = Tensor.rand(3, j)
+        symbolic = jf(a.reshape(3, vi), b.reshape(3, vj)).reshape(3, i+j).numpy()
+        expected = f(a, b).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_two_vars_plus1_ij(self):
     def f(a, b): return (a@b+1).realize()
     jf = TinyJit(f)
-    a = Tensor.rand(10, 3)
-    b = Tensor.rand(3, 10)
     for i in range(1, 5):
       for j in range(1, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        symbolic = jf(a[:vi, :], b[:, :vj]).reshape(i, j).numpy()
-        expected = f(a[:i, :], b[:, :j]).numpy()
+        a = Tensor.rand(i, 3)
+        b = Tensor.rand(3, j)
+        symbolic = jf(a.reshape(vi, 3), b.reshape(3, vj)).reshape(i, j).numpy()
+        expected = f(a, b).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_two_vars_plus1_ji(self):
     def f(a, b): return (a@b+1).realize()
     jf = TinyJit(f)
-    a = Tensor.rand(10, 3)
-    b = Tensor.rand(3, 10)
     for i in range(1, 5):
       for j in range(1, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        symbolic = jf(a[:vj, :], b[:, :vi]).reshape(j, i).numpy()
-        expected = f(a[:j, :], b[:, :i]).numpy()
+        a = Tensor.rand(j, 3)
+        b = Tensor.rand(3, i)
+        symbolic = jf(a.reshape(vj, 3), b.reshape(3, vi)).reshape(j, i).numpy()
+        expected = f(a, b).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
   def test_jit_symbolic_shape_mismatch(self):
     @TinyJit
     def add(a, b): return (a+b).realize()
-    a = Tensor.rand(3, 10)
-    b = Tensor.rand(3, 10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      add(a[:, :vi], b[:, :vi])
+      a = Tensor.rand(3, i).reshape(3, vi)
+      b = Tensor.rand(3, i).reshape(3, vi)
+      add(a, b)
     vi2 = Variable("i", 1, 10).bind(7)
-    a = Tensor.rand(3, 7)[:, :vi2]
-    bad = Tensor.rand(4, 7)[:, :vi2]
+    a = Tensor.rand(3, 7).reshape(3, vi2)
+    bad = Tensor.rand(4, 7).reshape(4, vi2)
     with self.assertRaises(AssertionError):
       add(a, bad)
 
@@ -166,9 +175,9 @@ class TestSymbolicJit(unittest.TestCase):
     # shrink is a movement, so we pair it with a simple function to test the JIT interaction
     def f(a): return (a+1).realize()
     jf = TinyJit(f)
-    a = Tensor.rand(7, 11)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
+      a = Tensor.rand(7, 11)
       symbolic = a.shrink(((3,5),(vi,vi+2)))
       symbolic = jf(symbolic).numpy()
       expected = f(a.shrink(((3,5),(i,i+2)))).numpy()
@@ -179,9 +188,9 @@ class TestSymbolicJit(unittest.TestCase):
     # slice is a movement, so we pair it with a simple function to test the JIT interaction
     def f(a): return (a+1).realize()
     jf = TinyJit(f)
-    a = Tensor.rand(7, 11)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
+      a = Tensor.rand(7, 11)
       symbolic = a[3:5, vi:vi+2]
       symbolic = jf(symbolic).numpy()
       expected = f(a[3:5, i:i+2]).numpy()
@@ -203,11 +212,11 @@ class TestSymbolicJit(unittest.TestCase):
   def test_ones_sum(self):
     def f(a): return a.sum().realize()
     jf = TinyJit(f)
-    t = Tensor.ones(10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = jf(t[:vi]).item()
-      expected = f(t[:i]).item()
+      t = Tensor.ones(i)
+      symbolic = jf(t.reshape(vi)).item()
+      expected = f(t).item()
       np.testing.assert_equal(symbolic, expected)
 
   def test_mean(self):
@@ -217,22 +226,22 @@ class TestSymbolicJit(unittest.TestCase):
     jf = TinyJit(f)
     jf0 = TinyJit(f0)
     jf1 = TinyJit(f1)
-    a = Tensor.rand(10, 3)
-    b = Tensor.rand(10, 3)
-    c = Tensor.rand(10, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      # axis = None
-      symbolic = jf(a[:vi]).numpy()
-      expected = a[:i].mean().numpy()
+      # aixs = None
+      a = Tensor.rand(i, 3)
+      symbolic = jf(a.reshape(vi, 3)).numpy()
+      expected = a.mean().numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-      # axis = 0
-      symbolic = jf0(b[:vi]).numpy()
-      expected = b[:i].mean(0).numpy()
+      # aixs = 0
+      a = Tensor.rand(i, 3)
+      symbolic = jf0(a.reshape(vi, 3)).numpy()
+      expected = a.mean(0).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-      # axis = 1
-      symbolic = jf1(c[:vi]).reshape(i).numpy()
-      expected = c[:i].mean(1).numpy()
+      # aixs = 1
+      a = Tensor.rand(i, 3)
+      symbolic = jf1(a.reshape(vi, 3)).reshape(i).numpy()
+      expected = a.mean(1).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_mean_2d(self):
@@ -242,24 +251,24 @@ class TestSymbolicJit(unittest.TestCase):
     jf = TinyJit(f)
     jf0 = TinyJit(f0)
     jf1 = TinyJit(f1)
-    a = Tensor.rand(10, 10)
-    b = Tensor.rand(10, 10)
-    c = Tensor.rand(10, 10)
     for i in range(1, 5):
       for j in range(1, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        # axis = None
-        symbolic = jf(a[:vi, :vj]).numpy()
-        expected = a[:i, :j].mean().numpy()
+        # aixs = None
+        a = Tensor.rand(i, j)
+        symbolic = jf(a.reshape(vi, vj)).numpy()
+        expected = a.mean().numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-        # axis = 0
-        symbolic = jf0(b[:vi, :vj]).reshape(j).numpy()
-        expected = b[:i, :j].mean(0).numpy()
+        # aixs = 0
+        a = Tensor.rand(i, j)
+        symbolic = jf0(a.reshape(vi, vj)).reshape(j).numpy()
+        expected = a.mean(0).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-        # axis = 1
-        symbolic = jf1(c[:vi, :vj]).reshape(i).numpy()
-        expected = c[:i, :j].mean(1).numpy()
+        # aixs = 1
+        a = Tensor.rand(i, j)
+        symbolic = jf1(a.reshape(vi, vj)).reshape(i).numpy()
+        expected = a.mean(1).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_var(self):
@@ -269,22 +278,22 @@ class TestSymbolicJit(unittest.TestCase):
     jf = TinyJit(f)
     jf0 = TinyJit(f0)
     jf1 = TinyJit(f1)
-    a = Tensor.rand(10, 3)
-    b = Tensor.rand(10, 3)
-    c = Tensor.rand(10, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      # axis = None
-      symbolic = jf(a[:vi]).numpy()
-      expected = a[:i].var().numpy()
+      # aixs = None
+      a = Tensor.rand(i, 3)
+      symbolic = jf(a.reshape(vi, 3)).numpy()
+      expected = a.var().numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-      # axis = 0
-      symbolic = jf0(b[:vi]).numpy()
-      expected = b[:i].var(0).numpy()
+      # aixs = 0
+      a = Tensor.rand(i, 3)
+      symbolic = jf0(a.reshape(vi, 3)).numpy()
+      expected = a.var(0).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-      # axis = 1
-      symbolic = jf1(c[:vi]).reshape(i).numpy()
-      expected = c[:i].var(1).numpy()
+      # aixs = 1
+      a = Tensor.rand(i, 3)
+      symbolic = jf1(a.reshape(vi, 3)).reshape(i).numpy()
+      expected = a.var(1).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_var_2d(self):
@@ -294,24 +303,24 @@ class TestSymbolicJit(unittest.TestCase):
     jf = TinyJit(f)
     jf0 = TinyJit(f0)
     jf1 = TinyJit(f1)
-    a = Tensor.rand(10, 10)
-    b = Tensor.rand(10, 10)
-    c = Tensor.rand(10, 10)
     for i in range(1, 5):
       for j in range(1, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        # axis = None
-        symbolic = jf(a[:vi, :vj]).numpy()
-        expected = a[:i, :j].var().numpy()
+        # aixs = None
+        a = Tensor.rand(i, j)
+        symbolic = jf(a.reshape(vi, vj)).numpy()
+        expected = a.var().numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-        # axis = 0
-        symbolic = jf0(b[:vi, :vj]).reshape(j).numpy()
-        expected = b[:i, :j].var(0).numpy()
+        # aixs = 0
+        a = Tensor.rand(i, j)
+        symbolic = jf0(a.reshape(vi, vj)).reshape(j).numpy()
+        expected = a.var(0).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-        # axis = 1
-        symbolic = jf1(c[:vi, :vj]).reshape(i).numpy()
-        expected = c[:i, :j].var(1).numpy()
+        # aixs = 1
+        a = Tensor.rand(i, j)
+        symbolic = jf1(a.reshape(vi, vj)).reshape(i).numpy()
+        expected = a.var(1).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
 if __name__ == '__main__':

test/test_symbolic_ops.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Variable
 from tinygrad.shape.shapetracker import View
-from tinygrad.helpers import GlobalCounters
+from tinygrad.helpers import Context, GlobalCounters
 from tinygrad.uop.ops import sym_infer
 from tinygrad.dtype import dtypes
 from tinygrad.device import Device
@@ -9,46 +9,54 @@ from examples.gpt2 import Attention
 import numpy as np
 
 class TestSymbolicOps(unittest.TestCase):
+  def setUp(self):
+    # A lot of these test are out of bounds, so we ignore the bounds check
+    self.context = Context(IGNORE_OOB=1)
+    self.context.__enter__()
+
+  def tearDown(self):
+    self.context.__exit__(None, None, None)
+
   def test_plus1(self):
     def f(a): return (a+1).realize()
-    a = Tensor.rand(3, 10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = f(a[:, :vi]).reshape(3, i).numpy()
-      expected = f(a[:, :i]).numpy()
+      a = Tensor.rand(3, i)
+      symbolic = f(a.reshape(3, vi)).reshape(3, i).numpy()
+      expected = f(a).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_add(self):
     def f(a, b): return (a+b).realize()
-    a = Tensor.rand(3, 10)
-    b = Tensor.rand(3, 10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = f(a[:, :vi], b[:, :vi]).reshape(3, i).numpy()
-      expected = f(a[:, :i], b[:, :i]).numpy()
+      a = Tensor.rand(3, i)
+      b = Tensor.rand(3, i)
+      symbolic = f(a.reshape(3, vi), b.reshape(3, vi)).reshape(3, i).numpy()
+      expected = f(a, b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_matmul(self):
     def f(a, b): return (a@b).realize()
-    a = Tensor.rand(3, 10)
-    b = Tensor.rand(10, 5)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = f(a[:, :vi], b[:vi, :]).numpy()
-      expected = f(a[:, :i], b[:i, :]).numpy()
+      a = Tensor.rand(3, i)
+      b = Tensor.rand(i, 5)
+      symbolic = f(a.reshape(3, vi), b.reshape(vi, 5)).numpy()
+      expected = f(a, b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_attention(self, dropout_p=0.0, imin=1, imax=5, use_symbolic=True):
     def f(q, k, v): return Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=dropout_p).realize()
-    q = Tensor.rand(2, 1, 4, 8)
-    k = Tensor.rand(2, 10, 4, 8)
-    v = Tensor.rand(2, 10, 4, 8)
     for i in range(imin, imax):
       vi = Variable("i", 1, 10).bind(i) if use_symbolic else i
+      q = Tensor.rand(2, 1, 4, 8)
+      k = Tensor.rand(2, i, 4, 8)
+      v = Tensor.rand(2, i, 4, 8)
       Tensor.realize(q, k, v)
       GlobalCounters.reset()
-      symbolic = f(q, k[:, :vi, :, :], v[:, :vi, :, :]).reshape(2, 4, 1, 8).numpy()
-      expected = f(q, k[:, :i, :, :], v[:, :i, :, :]).numpy()
+      symbolic = f(q, k.reshape(2, vi, 4, 8), v.reshape(2, vi, 4, 8)).reshape(2, 4, 1, 8).numpy()
+      expected = f(q, k, v).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_attention_cmp_symbolic(self):
@@ -82,89 +90,73 @@ class TestSymbolicOps(unittest.TestCase):
 
   def test_cat_dim0(self):
     def f(a, b): return a.cat(b, dim=0).realize()
-    a = Tensor.rand(10, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
+      a = Tensor.rand(i, 3)
       b = Tensor.rand(2, 3)
-      symbolic = f(a[:vi, :], b).reshape(i+2, 3).numpy()
-      expected = f(a[:i, :], b).numpy()
+      symbolic = f(a.reshape(vi, 3), b).reshape(i+2, 3).numpy()
+      expected = f(a, b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_cat_dim1(self):
     def f(a, b): return a.cat(b, dim=1).realize()
-    a = Tensor.rand(3, 10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
+      a = Tensor.rand(3, i)
       b = Tensor.rand(3, 2)
-      symbolic = f(a[:, :vi], b).reshape(3, i+2).numpy()
-      expected = f(a[:, :i], b).numpy()
+      symbolic = f(a.reshape(3, vi), b).reshape(3, i+2).numpy()
+      expected = f(a, b).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_cat_dim0_two_vars(self):
     def f(a, b): return a.cat(b, dim=0).realize()
-    a = Tensor.rand(10, 3)
-    b = Tensor.rand(10, 3)
     for i in range(1, 5):
       for j in range(1, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        symbolic = f(a[:vi, :], b[:vj, :]).reshape(i+j, 3).numpy()
-        expected = f(a[:i, :], b[:j, :]).numpy()
+        a = Tensor.rand(i, 3)
+        b = Tensor.rand(j, 3)
+        symbolic = f(a.reshape(vi, 3), b.reshape(vj, 3)).reshape(i+j, 3).numpy()
+        expected = f(a, b).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_cat_dim1_two_vars(self):
     def f(a, b): return a.cat(b, dim=1).realize()
-    a = Tensor.rand(3, 10)
-    b = Tensor.rand(3, 10)
     for i in range(1, 5):
       for j in range(1, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        symbolic = f(a[:, :vi], b[:, :vj]).reshape(3, i+j).numpy()
-        expected = f(a[:, :i], b[:, :j]).numpy()
+        a = Tensor.rand(3, i)
+        b = Tensor.rand(3, j)
+        symbolic = f(a.reshape(3, vi), b.reshape(3, vj)).reshape(3, i+j).numpy()
+        expected = f(a, b).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_two_vars_plus1_ij(self):
     def f(a, b): return (a@b+1).realize()
-    a = Tensor.rand(10, 3)
-    b = Tensor.rand(3, 10)
     for i in range(1, 5):
       for j in range(1, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        symbolic = f(a[:vi, :], b[:, :vj]).reshape(i, j).numpy()
-        expected = f(a[:i, :], b[:, :j]).numpy()
+        a = Tensor.rand(i, 3)
+        b = Tensor.rand(3, j)
+        symbolic = f(a.reshape(vi, 3), b.reshape(3, vj)).reshape(i, j).numpy()
+        expected = f(a, b).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_two_vars_plus1_ji(self):
     # reverse the order of variables
     def f(a, b): return (a@b+1).realize()
-    a = Tensor.rand(10, 3)
-    b = Tensor.rand(3, 10)
     for i in range(1, 5):
       for j in range(1, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
-        symbolic = f(a[:vj, :], b[:, :vi]).reshape(j, i).numpy()
-        expected = f(a[:j, :], b[:, :i]).numpy()
+        a = Tensor.rand(j, 3)
+        b = Tensor.rand(3, i)
+        symbolic = f(a.reshape(vj, 3), b.reshape(3, vi)).reshape(j, i).numpy()
+        expected = f(a, b).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
-  def test_reshape_from_symbolic(self):
-    a = Tensor.rand(30)
-    for i in range(3, 5):
-      vi = Variable("i", 3, 10).bind(i)
-      symbolic = a[:vi*3].reshape((3, 3)).numpy()
-      # To match symbolic reshape (potential implicit shrink), we need a shrink
-      expected = a[:i*3].shrink(((0, 9),)).reshape((3, 3)).numpy()
-      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-
-  def test_invalid_symbolic_reshape(self):
-    a = Tensor.rand(30)
-    for i in range(1, 5):
-      vi = Variable("i", 1, 10).bind(i)
-      # Cannot reshape into symbolic from non-symbolic
-      with self.assertRaises(AssertionError): a.reshape((3, vi))
-
   def test_shrink(self):
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
@@ -184,10 +176,11 @@ class TestSymbolicOps(unittest.TestCase):
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_slice_no_start(self):
-    a = Tensor.rand(7, 11)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = a[3:5, :vi:1].reshape(2, i).numpy()
+      a = Tensor.rand(7, 11)
+      symbolic = a[3:5, :vi:1].reshape(2,i)
+      symbolic = symbolic.numpy()
       expected = a[3:5, :i:1].numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
@@ -208,75 +201,75 @@ class TestSymbolicOps(unittest.TestCase):
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_ones_sum(self):
-    t = Tensor.ones(10)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      symbolic = t[:vi].sum().item()
-      expected = t[:i].sum().item()
+      t = Tensor.ones(i)
+      symbolic = t.reshape(vi).sum().item()
+      expected = t.sum().item()
       np.testing.assert_equal(symbolic, expected)
 
   def test_mean(self):
-    a = Tensor.rand(10, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
       for axis in [None, 0, 1]:
-        expected = a[:i].mean(axis).numpy()
-        symbolic = a[:vi].mean(axis).reshape(expected.shape).numpy()
+        a = Tensor.rand(i, 3)
+        expected = a.mean(axis).numpy()
+        symbolic = a.reshape(vi, 3).mean(axis).reshape(expected.shape).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_mean_2d(self):
-    a = Tensor.rand(10, 10)
     for i in range(1, 5):
       for j in range(1, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
         for axis in [None, 0, 1]:
-          expected = a[:i, :j].mean(axis).numpy()
-          symbolic = a[:vi, :vj].mean(axis).reshape(expected.shape).numpy()
+          a = Tensor.rand(i, j)
+          expected = a.mean(axis).numpy()
+          symbolic = a.reshape(vi, vj).mean(axis).reshape(expected.shape).numpy()
           np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_var(self):
-    a = Tensor.rand(10, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
       for axis in [None, 0, 1]:
-        expected = a[:i].var(axis).numpy()
-        symbolic = a[:vi].var(axis).reshape(expected.shape).numpy()
+        a = Tensor.rand(i, 3)
+        expected = a.var(axis).numpy()
+        symbolic = a.reshape(vi, 3).var(axis).reshape(expected.shape).numpy()
         np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_var_2d(self):
-    a = Tensor.rand(10, 10)
     for i in range(1, 5):
       for j in range(1, 5):
         vi = Variable("i", 1, 10).bind(i)
         vj = Variable("j", 1, 10).bind(j)
         for axis in [None, 0, 1]:
-          expected = a[:i, :j].var(axis).numpy()
-          symbolic = a[:vi, :vj].var(axis).reshape(expected.shape).numpy()
+          a = Tensor.rand(i, j)
+          expected = a.var(axis).numpy()
+          symbolic = a.reshape(vi, vj).var(axis).reshape(expected.shape).numpy()
           np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
   def test_bitcast_down(self):
-    a = Tensor.rand(10, 3)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      expected = a[:i].bitcast(dtypes.uint8).numpy()
-      symbolic = a[:vi].bitcast(dtypes.uint8).reshape(expected.shape).numpy()
+      a = Tensor.rand(i, 3)
+      expected = a.bitcast(dtypes.uint8).numpy()
+      symbolic = a.reshape(vi, 3).bitcast(dtypes.uint8).reshape(expected.shape).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)
 
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "no uint64")
   def test_bitcast_up(self):
-    a = Tensor.rand(10, 4)
     for i in range(1, 5):
       vi = Variable("i", 1, 10).bind(i)
-      expected = a[:i].bitcast(dtypes.uint64).numpy()
-      symbolic = a[:vi].bitcast(dtypes.uint64).reshape(expected.shape).numpy()
+      a = Tensor.rand(i, 4)
+      expected = a.bitcast(dtypes.uint64).numpy()
+      symbolic = a.reshape(vi, 4).bitcast(dtypes.uint64).reshape(expected.shape).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)
 
   @unittest.expectedFailure
   def test_conv2d_ceildiv_edge_case(self):
     v = Variable('v', 11, 50_000)
     val = 39601
-    x = Tensor.randn(1, 22, 50_000)[:, :, :v.bind(val)]
+    x = Tensor.randn(1, 22, 39601).reshape(1, 22, v.bind(val))
     weight = Tensor.randn(256, 22, 12)
 
     result = x.conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))

test/test_tensor.py

@@ -415,21 +415,6 @@ class TestTinygrad(unittest.TestCase):
         data = _generate_data(depth)
         np.testing.assert_allclose(Tensor(data).numpy(), np.array(data))
 
-  def test_tensor_list_implicit_cast(self):
-    data = [True, False]
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
-    data = [-1, 0, 1, 2, 3]
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
-    data = [-3.5, -2.5, -1.5, 0, 1.5, 2.5, 3.5]
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
-    # NOTE: torch and jax raise OverflowError: Python integer -3 out of bounds for uint8
-    # np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
-    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
-
   def test_tensor_list_special_values(self):
     if is_dtype_supported(dtypes.float16):
       data = [math.nan, -math.inf, 65504, 65519, 65519.999, 65520, 65520.1]

test/test_tensor_variable.py

@@ -1,6 +1,7 @@
 import unittest
 import numpy as np
 from tinygrad import Tensor, Variable
+from tinygrad.helpers import Context
 
 class TestTensorVariable(unittest.TestCase):
   def test_add_tvar(self):
@@ -22,38 +23,43 @@ class TestTensorVariable(unittest.TestCase):
     assert (Tensor(3) * (vv * 4)).item() == 24
 
   def test_symbolic_mean(self):
-    vv = Variable("a", 1, 10).bind(2)
-    t = Tensor.ones(2, 10).contiguous()[:, :vv]
-    ret = t.mean().item()
-    assert ret == 1
+    with Context(IGNORE_OOB=1):
+      vv = Variable("a", 1, 10).bind(2)
+      t = Tensor.ones(2, 2).contiguous().reshape(2, vv)
+      ret = t.mean().item()
+      assert ret == 1
 
   def test_symbolic_mean_2d(self):
-    vv = Variable("a", 1, 10).bind(2)
-    vv2 = Variable("b", 1, 10).bind(2)
-    t = Tensor.ones(10, 10).contiguous()[:vv2, :vv]
-    ret = t.mean().item()
-    assert ret == 1
+    with Context(IGNORE_OOB=1):
+      vv = Variable("a", 1, 10).bind(2)
+      vv2 = Variable("b", 1, 10).bind(2)
+      t = Tensor.ones(2, 2).contiguous().reshape(vv2, vv)
+      ret = t.mean().item()
+      assert ret == 1
 
   def test_symbolic_mean_2d_axis_1(self):
-    vv = Variable("a", 1, 10).bind(2)
-    vv2 = Variable("b", 1, 10).bind(2)
-    t = Tensor.ones(10, 10).contiguous()[:vv2, :vv]
-    ret = t.mean(axis=1).reshape(2, 1).numpy()
-    assert np.all(ret == 1)
+    with Context(IGNORE_OOB=1):
+      vv = Variable("a", 1, 10).bind(2)
+      vv2 = Variable("b", 1, 10).bind(2)
+      t = Tensor.ones(2, 2).contiguous().reshape(vv2, vv)
+      ret = t.mean(axis=1).reshape(2, 1).numpy()
+      assert np.all(ret == 1)
 
   def test_symbolic_mean_2d_add(self):
-    add_term = Variable("c", 0, 10).bind(1)
-    vv = Variable("a", 1, 10).bind(1)
-    vv2 = Variable("b", 1, 10).bind(1)
-    t = Tensor.ones(20, 20).contiguous()[:vv2+add_term, :vv+add_term]
-    ret = t.mean().item()
-    assert ret == 1
+    with Context(IGNORE_OOB=1):
+      add_term = Variable("c", 0, 10).bind(1)
+      vv = Variable("a", 1, 10).bind(1)
+      vv2 = Variable("b", 1, 10).bind(1)
+      t = Tensor.ones(2, 2).contiguous().reshape(vv2+add_term, vv+add_term)
+      ret = t.mean().item()
+      assert ret == 1
 
   def test_symbolic_var(self):
-    vv = Variable("a", 1, 10).bind(2)
-    t = Tensor.ones(2, 10).contiguous()[:, :vv]
-    ret = t.var().item()
-    assert ret == 0
+    with Context(IGNORE_OOB=1):
+      vv = Variable("a", 1, 10).bind(2)
+      t = Tensor.ones(2, 2).contiguous().reshape(2, vv)
+      ret = t.var().item()
+      assert ret == 0
 
   def test_symbolic_pad(self):
     vv = Variable("a", 1, 10).bind(2)
@@ -86,15 +92,5 @@ class TestTensorVariable(unittest.TestCase):
     ret = Tensor.arange(begin.bind(4), end.bind(7))
     self.assertListEqual(ret.reshape(3).tolist(), [4,5,6])
 
-  def test_variable_empty(self):
-    v = Variable("i", 1, 10)
-    # TODO: Tensor creation from unbound variable should assert
-    # with self.assertRaises(AssertionError): t = Tensor.empty(3, v)
-    vb = v.bind(3)
-    t = Tensor.empty(3, vb)
-    assert t.uop.base.buffer.size == 30
-    assert t.uop.st.shape == (3, vb)
-
-
 if __name__ == '__main__':
   unittest.main()

test/test_tiny.py

@@ -30,10 +30,7 @@ class TestTiny(unittest.TestCase):
   def test_gemm(self, N=64, out_dtype=dtypes.float):
     a = Tensor.ones(N,N).contiguous()
     b = Tensor.eye(N).contiguous()
-    lst = (out:=a@b).tolist()
-    for y in range(N):
-      for x in range(N):
-        self.assertEqual(lst[y][x], 1.0, msg=f"mismatch at ({y},{x})")
+    self.assertListEqual((out:=a@b).flatten().tolist(), [1.0]*(N*N))
     if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)
 
   # *** randomness ***
@@ -76,17 +73,17 @@ class TestTiny(unittest.TestCase):
 
   def test_symbolic(self):
     i = Variable('i', 1, 10)
-    ones = Tensor.ones(10).contiguous()
-    for s in [2,5]:
-      ret = ones[:i.bind(s)] + 1
-      self.assertListEqual(ret.contiguous().reshape(s).tolist(), [2.0]*s)
+    with Context(IGNORE_OOB=1):
+      for s in [2,5]:
+        ret = Tensor.ones(s).contiguous().reshape(i.bind(s)) + 1
+        self.assertListEqual(ret.reshape(s).tolist(), [2.0]*s)
 
   def test_symbolic_reduce(self):
     i = Variable('i', 1, 10)
-    ones = Tensor.ones(10).contiguous()
-    for s in [2,5]:
-      ret = ones[:i.bind(s)].sum()
-      self.assertEqual(ret.item(), s)
+    with Context(IGNORE_OOB=1):
+      for s in [2,5]:
+        ret = Tensor.ones(s).contiguous().reshape(i.bind(s)).sum()
+        self.assertEqual(ret.item(), s)
 
   # *** a model ***
 
@@ -109,24 +106,6 @@ class TestTiny(unittest.TestCase):
     probs = Tensor.rand(1, 1, 28, 28).sequential(layers).tolist()
     self.assertEqual(len(probs[0]), 10)
 
-  # TODO: this is failing because of how swizzling rewrites the ShapeTracker of the final STORE
-  @unittest.skipIf(IMAGE>0 or (CI and Device.DEFAULT == "DSP"), "failing because of make things that can't be images not images")
-  def test_mnist_backward(self):
-    # NOTE: we don't have the whole model here for speed
-    layers = [
-      nn.Conv2d(1, 32, 5), Tensor.relu,
-      nn.Conv2d(32, 32, 5), Tensor.relu]
-
-    # replace random weights with ones
-    # TODO: there's a bug here where it's tying two of the biases together. we need UNIQUE const
-    #Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)])
-    for p in nn.state.get_parameters(layers): p.replace(Tensor.empty(p.shape))
-
-    # realize gradients
-    for x in nn.state.get_parameters(layers): x.requires_grad_()
-    Tensor.empty(4, 1, 28, 28).sequential(layers).sum().backward()
-    Tensor.realize(*[x.grad for x in nn.state.get_parameters(layers) if x.grad is not None])
-
   # *** image ***
 
   @unittest.skipIf(Device.DEFAULT != "GPU", "image only supported on GPU")

test/test_uop_graph.py

@@ -6,7 +6,7 @@ from tinygrad.helpers import DEBUG, Context
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp
 from tinygrad.uop.symbolic import sym
 from tinygrad.codegen import full_rewrite, full_rewrite_to_sink
-from tinygrad.codegen.late.expander import expander
+from tinygrad.codegen.expander import expander
 
 simple_pm = PatternMatcher([
   (UPat.cvar('x', dtypes.int), lambda x: UOp.const(dtypes.float, 1.0) + UOp.const(dtypes.float, 2.0)),
@@ -441,16 +441,18 @@ class TestUOpGraph(unittest.TestCase):
       ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 20)),))
       with self.assertRaises(RuntimeError): to_uops_list([ld0])
 
+  @unittest.skip("outdated")
   def test_in_out_of_bounds_access_gated_store(self):
     with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), src=(), arg=0)
+      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
       v = Variable("v", 0, 20)
-      st0 = UOp(Ops.STORE, dtypes.void, src=(glbl0.index(v, v<16), UOp.const(dtypes.int, 0)))
+      st0 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v), UOp.const(dtypes.int, 0), v<16))
       to_uops_list([st0])
 
-      st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v, v<20), v))
+      st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v), v, v<20))
       with self.assertRaises(RuntimeError): to_uops_list([st1])
 
+  @unittest.skip("outdated")
   def test_in_bounds_access_gated_local(self):
     with Context(IGNORE_OOB=0):
       # Define buffers
@@ -463,7 +465,7 @@ class TestUOpGraph(unittest.TestCase):
 
       gate = (gidx<400) & (lidx<8)
 
-      local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx, lidx<8), UOp.const(dtypes.uint, 1)))
+      local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx), UOp.const(dtypes.uint, 1), lidx<8))
 
       barrier = UOp(Ops.BARRIER, dtypes.void, (local_store,))
       if_barrier = UOp(Ops.IF, dtypes.void, (gate, barrier))
@@ -475,34 +477,6 @@ class TestUOpGraph(unittest.TestCase):
       global_store = UOp(Ops.STORE, dtypes.void, (gbuf.index(gidx), local_load))
       to_uops_list([global_store])
 
-  def test_load_with_float_in_index(self):
-    with Context(IGNORE_OOB=0):
-      ridx = UOp.range(20, 0)
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      i = (ridx.cast(dtypes.float)*0.68).trunc().cast(dtypes.int)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, ((0<=i)&(i<16))),))
-      to_uops_list([ld0])
-      glblfloat = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(20), (), 0)
-      ldfloat = UOp(Ops.LOAD, dtypes.float, (glblfloat.index(ridx),))
-      i = (ldfloat+3.14).cast(dtypes.int)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, ((0<=i)&(i<16))),))
-
-  def test_load_cast_to_bool(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), (), 0)
-      ridx = UOp.range(20, 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(ridx, ridx.cast(dtypes.bool).logical_not()),))
-      to_uops_list([ld0])
-
-  @unittest.skip("Bool load is not supported yet")
-  def test_load_mask(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      mask = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
-      ridx = UOp.range(20, 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask),)))
-      to_uops_list([ld0])
-
   def test_out_of_bounds_off_by_one_access(self):
     with Context(IGNORE_OOB=0):
       glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
@@ -591,9 +565,10 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_switched_range_order(self):
     glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
+    c2 = UOp.const(dtypes.int, 2)
     cf = UOp.const(dtypes.float, 0.0)
-    r1 = UOp.range(2, 0)
-    r2 = UOp.range(2, 1)
+    r1 = UOp(Ops.RANGE, dtypes.int, (c2,), 0)
+    r2 = UOp(Ops.RANGE, dtypes.int, (c2,), 1)
     alu = UOp(Ops.MUL, dtypes.int, (r2, r1))
     store = UOp(Ops.STORE, dtypes.void, (glbl.index(alu), cf))
     uops = to_uops_list([store])

test/test_uops.py

@@ -22,7 +22,7 @@ def _uops_to_prg(uops_list):
   uops = full_rewrite(ast:=UOp.sink(*uops_list), opts=Device[Device.DEFAULT].renderer)
   src = Device[Device.DEFAULT].renderer.render(uops)
   has_local = Device[Device.DEFAULT].renderer.has_local
-  return CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test", src, Device.DEFAULT, ast, uops=uops,
+  return CompiledRunner(ProgramSpec("test", src, Device.DEFAULT, ast, uops=uops,
                                 global_size=[1,1,1] if has_local else None, local_size=[1,1,1] if has_local else None))
 
 def uop(uops:list[UOp], uop:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
@@ -402,14 +402,6 @@ class TestAssembly(unittest.TestCase):
     self.assertIn(Ops.SHR, ops)
     self.assertNotIn(Ops.IDIV, ops)
 
-  def test_fast_idiv_remove_powers_of_two(self):
-    ridx = UOp.range(2**20, 0)
-    uops = to_uops_list([ridx//(7*64)], opts=Device[Device.DEFAULT].renderer)
-    ops = [x.op for x in uops]
-    # this requires shifting out the powers of two before doing fast_idiv
-    # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
-    self.assertNotIn(Ops.CAST, ops)
-
   def test_mulacc_unrolled(self):
     # test that     acc = acc + a0*b0 + a1*b1 + a2*b2 + a3*b3
     # is not        acc = acc + (a0*b0 + a1*b1 + a2*b2 + a3*b3)
@@ -447,7 +439,7 @@ class TestUOpMethod(unittest.TestCase):
   def test_uop_variables(self):
     a = UOp.variable("a", 1, 10)
     uop_var = Tensor(a.bind(1))
-    st_var = Tensor.empty((2, 10))[:, :a.bind(1)]
+    st_var = Tensor.empty((2, 1)).reshape((2, a.bind(1)))
     _, var_vals = (uop_var+st_var).schedule_with_vars()
     self.assertEqual(len(var_vals), 1)
     self.assertEqual(list(var_vals)[0], a)

test/test_uops_stats.py

@@ -20,7 +20,6 @@ def get_stats(x:Tensor):
   ei = lower_schedule_item(si)
   return ei.prg.estimates.ops, ei.prg.estimates.mem
 
-@unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu does extra load/store for packed types")
 class TestMemoryCount(unittest.TestCase):
   def test_add(self):
     a = Tensor.empty(1024, 1024, dtype=dtypes.uint8)

test/unit/test_attention.py

@@ -1,6 +1,5 @@
 import unittest
-from tinygrad import Tensor, dtypes, TinyJit, UOp
-from tinygrad.apps.llm import apply_rope
+from tinygrad import Tensor, dtypes
 
 # TODO: test_scheduler, but just in uint
 class TestAttention(unittest.TestCase):
@@ -17,29 +16,5 @@ class TestAttention(unittest.TestCase):
     for si in softmax_inputs:
       assert all(b.dtype == dtypes.half for b in si.bufs), f"non half {si.bufs=}"
 
-  def test_apply_rope(self):
-    x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32)
-    result = apply_rope(x, 0)
-    self.assertEqual(result.shape, x.shape)
-    self.assertEqual(result.dtype, x.dtype)
-    self.assertGreater((result - apply_rope(x, 5)).abs().max().item(), 1e-6)
-    with self.assertRaises(AssertionError): apply_rope(Tensor.randn(1, 1, 4, 7, dtype=dtypes.float32), 0)
-
-  def test_apply_rope_jit_prune(self):
-    def rope_fn(x_in, pos): return apply_rope(x_in, pos)
-    rope_noprune = TinyJit(rope_fn)
-    rope_prune = TinyJit(rope_fn, prune=True)
-
-    v_pos = UOp.variable("start_pos", 0, 100)
-    for _ in range(3):
-      rope_noprune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1))
-      rope_prune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1))
-    noprune_size = len(rope_noprune.captured.jit_cache)
-    prune_size = len(rope_prune.captured.jit_cache)
-
-    self.assertGreater(noprune_size, prune_size)
-    self.assertGreaterEqual(noprune_size, 3)
-    self.assertEqual(prune_size, 1)
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_block_reorder.py

@@ -1,7 +1,7 @@
 import unittest, random
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import print_uops, UOp, Ops
-from tinygrad.codegen.late.linearize import block_reorder
+from tinygrad.codegen.linearize import block_reorder
 from tinygrad.renderer.cstyle import OpenCLRenderer
 
 def is_toposorted(lst:list[UOp]):

test/unit/test_dtype.py

@@ -53,14 +53,5 @@ class TestCastConvenienceMethod(unittest.TestCase):
       self.assertEqual(t.float().dtype, dtypes.float)
       self.assertEqual(t.double().dtype, dtypes.double)
 
-class TestDtypeTolist(unittest.TestCase):
-  def test_bfloat16(self):
-    self.assertEqual(Tensor([-60000, 1.5, 3.1, 60000], device="PYTHON", dtype=dtypes.bfloat16).tolist(), [-59904.0, 1.5, 3.09375, 59904.0])
-  def test_fp8(self):
-    # 448
-    self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e4m3).tolist(), [-448.0, 1.5, 3.0, 448.0])
-    # 57344
-    self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e5m2).tolist(), [-28672.0, 1.5, 3.0, 28672.0])
-
 if __name__ == "__main__":
   unittest.main()

test/unit/test_dtype_spec.py

@@ -1,6 +1,6 @@
-import unittest, math, operator, subprocess, struct
+import unittest, math, operator, subprocess
 from tinygrad.tensor import Tensor, dtypes, Device
-from tinygrad.dtype import DType, DTYPES_DICT, truncate, truncate_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
+from tinygrad.dtype import DType, DTYPES_DICT, truncate, truncate_fp16, truncate_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, CI, DEBUG
 from hypothesis import given, settings, strategies as strat
@@ -26,9 +26,6 @@ def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float
   except AssertionError as e:
     raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e
 
-def u32_to_f32(u): return struct.unpack('f', struct.pack('I', u))[0]
-def f32_to_u32(f): return struct.unpack('I', struct.pack('f', f))[0]
-
 class TestHelpers(unittest.TestCase):
   signed_ints = (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64)
   uints = (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
@@ -105,79 +102,18 @@ class TestHelpers(unittest.TestCase):
     self.assertEqual(truncate_fp16(65504), 65504)
     self.assertEqual(truncate_fp16(65519.999), 65504)
     self.assertEqual(truncate_fp16(65520), math.inf)
-    self.assertEqual(truncate_fp16(1e-8), 0.0)
-    self.assertEqual(truncate_fp16(-65504), -65504)
-    self.assertEqual(truncate_fp16(-65519.999), -65504)
-    self.assertEqual(truncate_fp16(-65520), -math.inf)
-    self.assertTrue(math.isnan(truncate_fp16(math.nan)))
 
-  def test_float_to_bf16(self):
-    # TODO: fuzz this better
+  def test_truncate_bf16(self):
+    self.assertEqual(truncate_bf16(1), 1)
+    self.assertAlmostEqual(truncate_bf16(1.1), 1.09375, places=7)
+    for a in [1234, 23456, -777.777]:
+      self.assertEqual(truncate_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
+    # TODO: torch bfloat 1.1 gives 1.1015625 instead of 1.09375
     max_bf16 = torch.finfo(torch.bfloat16).max
-    for a in [1, 1.1, 1234, 23456, -777.777, max_bf16, max_bf16 * 1.00001, -max_bf16, -max_bf16 * 1.00001, math.inf, -math.inf]:
-      self.assertEqual(float_to_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
-    self.assertTrue(math.isnan(float_to_bf16(math.nan)))
-
-  def test_float_to_bf16_nan(self):
-    # In f32, NaN = exp 0xFF and mantissa ≠ 0. Quiet-vs-signaling is bit 22 of the mantissa: 1 = qNaN, 0 = sNaN.
-    # qNaN(+/-), sNaN(+/-) overflow(+/-)
-    patterns = [0x7FC00001, 0xFFC00001, 0x7F800001, 0xFF800001, 0x7FFFFFFF, 0xFFFFFFFF]
-    for u in patterns:
-      x = u32_to_f32(u)
-      y = float_to_bf16(x)
-      t = torch.tensor([x], dtype=torch.bfloat16).item()
-      self.assertTrue(math.isnan(y))
-      self.assertTrue(math.isnan(t))
-
-  def test_float_to_bf16_round(self):
-    # round_to_nearest_even
-    uppers = [0x3f800000, 0x41230000, 0xC1460000] # 1.0, 10.1875, -12.375
-    for upper in uppers:
-      base = upper & 0xFFFF0000
-      base_f32 = u32_to_f32(base)
-      base_f32_round_up = u32_to_f32(base + 0x00010000)
-
-      # low < 0x8000(0.5ULP) -> round down
-      x = u32_to_f32(base | 0x00007000)
-      self.assertEqual(float_to_bf16(x), base_f32)
-      self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32)
-
-      # low > 0x8000(0.5ULP) -> round up
-      x = u32_to_f32(base | 0x0000C000)
-      self.assertEqual(float_to_bf16(x), base_f32_round_up)
-      self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32_round_up)
-
-      # low == 0x8000(0.5ULP) and LSB even -> round down
-      if ((upper >> 16) & 1) == 0:
-        x = u32_to_f32(base | 0x00008000)
-        self.assertEqual(float_to_bf16(x), base_f32)
-        self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32)
-      # low == 0x8000(0.5ULP) and LSB odd -> round up
-      else:
-        x = u32_to_f32(base | 0x00008000)
-        self.assertEqual(float_to_bf16(x), base_f32_round_up)
-        self.assertEqual(torch.tensor([x], dtype=torch.bfloat16).item(), base_f32_round_up)
-
-  def test_float_to_bf16_boundary(self):
-    # bf16 max finite: exp=0xFE, faction=0x7F => 0x7F7F0000(f32)
-    # bf16 inf(+/-):   exp=0xFF
-    base = 0x7F7F0000
-    inf_u32 = 0x7F800000
-
-    # low < 0.5ULP
-    x = u32_to_f32(base | 0x00007FFF)
-    self.assertEqual(f32_to_u32(float_to_bf16(x)), base)
-    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), base)
-
-    # low > 0.5ULP -> overflows to +inf
-    x = u32_to_f32(base | 0x0000C000)
-    self.assertEqual(f32_to_u32(float_to_bf16(x)), inf_u32)
-    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), inf_u32)
-
-    # low == 0.5ULP and LSB odd -> overflows to +inf
-    x = u32_to_f32(base | 0x00008000)
-    self.assertEqual(f32_to_u32(float_to_bf16(x)), inf_u32)
-    self.assertEqual(f32_to_u32(torch.tensor([x], dtype=torch.bfloat16).item()), inf_u32)
+    self.assertEqual(truncate_bf16(max_bf16), max_bf16)
+    self.assertEqual(truncate_bf16(min_bf16:=-max_bf16), min_bf16)
+    self.assertEqual(truncate_bf16(max_bf16 * 1.00001), math.inf)
+    self.assertEqual(truncate_bf16(min_bf16 * 1.00001), -math.inf)
 
   @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
   def test_truncate_fp8e4m3(self, x):

test/unit/test_gguf.py

@@ -53,37 +53,11 @@ class TestGGUF(unittest.TestCase):
   def test_load_tinyllama_q4_0(self): self._test_gguf_load("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q4_0.gguf?download=true")
   def test_load_gpt2_q4_1(self): self._test_gguf_load("https://huggingface.co/PrunaAI/gpt2-GGUF-smashed/resolve/main/gpt2.Q4_1.gguf?download=true")
   def test_load_sample_q6_k(self): self._test_gguf_load("https://huggingface.co/Isotr0py/test-gguf-sample/resolve/main/Quant_Q6_K_1024.gguf?download=true")
-  def test_load_sample_mxfp4(self): self._test_gguf_load("https://huggingface.co/ngxson/boring-testing-tiny/resolve/main/stories260K-mxfp4.gguf?download=true")
 
   def test_dequantization_q4_0(self): self._test_dequantization(ggml.GGML_TYPE_Q4_0)
   def test_dequantization_q4_1(self): self._test_dequantization(ggml.GGML_TYPE_Q4_1)
   def test_dequantization_q8_0(self): self._test_dequantization(ggml.GGML_TYPE_Q8_0)
   def test_dequantization_q6_k(self): self._test_dequantization(ggml.GGML_TYPE_Q6_K)
-  def test_dequantization_mxfp4(self):
-    MXFP4 = 39
-
-    def encode(nibbles, E):
-      packed = [(low & 0xF) | ((high & 0xF) << 4) for low, high in zip(nibbles[:16], nibbles[16:])]
-      return np.array([E] + packed, dtype=np.uint8)
-
-    def decode(code, E):
-      sign = -1.0 if code * 0b1000 else 1.0
-      exp = (code >> 1) & 0b11
-      mant = code & 0b1
-      val = (1.0 + 0.5 * mant) * np.exp2(exp - 1) if exp else 0.5 * mant
-      scale = np.exp2(E - 128) if E >= 2 else np.exp2(-127 if E == 1 else -128)
-      return sign * val * scale
-
-    blocks, expected = [], []
-    rng = np.random.default_rng(42)
-    for _ in range(4):
-      E = rng.integers(0, 256)
-      codes = rng.integers(0, 16, size=32, dtype=np.uint8)
-      blocks.append(encode(codes, E))
-      expected.extend(decode(c, E) for c in codes)
-    tensor = Tensor(np.concatenate(blocks))
-    out = ggml_data_to_tensor(tensor, len(expected), MXFP4)
-    self.assertListEqual(out.numpy().tolist(), np.array(expected, dtype=np.float32).tolist())
 
   def test_expected_failure_unknown_type(self):
     with self.assertRaises(ValueError):

test/unit/test_graph_rewrite.py

@@ -65,21 +65,21 @@ class TestFoldingAndReduction(unittest.TestCase):
   def test_full_graph_rewrite_reduction_with_unused_range(self):
     const1 = UOp.const(dtypes.int32, 15)
     const2 = UOp.const(dtypes.int32, 25)
-    rng = UOp.range(10, idx=0)
+    rng = UOp.range(dtypes.int32, 10, idx=0)
     optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng))
     expected_sum = 10 * (15 + 25)
     self.assertEqual(optimized_sink.arg, expected_sum)
 
   @unittest.skip("currently failing")
   def test_full_graph_rewrite_range_reduction(self):
-    simple_range = UOp.range(5, idx=0)
+    simple_range = UOp.range(dtypes.int32, 5, idx=0)
     optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range))
     expected_sum = sum(range(5))
     self.assertEqual(optimized_sink.arg, expected_sum)
 
   @unittest.skip("currently failing")
   def test_full_graph_rewrite_simple_reduction_folding(self):
-    simple_range = UOp.range(4, idx=0)
+    simple_range = UOp.range(dtypes.int32, 4, idx=0)
     add_uop = simple_range + UOp.const(dtypes.int32, 1)
     optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range))
     expected_sum = sum(i + 1 for i in range(4))
@@ -87,8 +87,8 @@ class TestFoldingAndReduction(unittest.TestCase):
 
   @unittest.skip("currently failing")
   def test_full_graph_rewrite_nested_loop_collapse(self):
-    outer_range = UOp.range(8, 0)
-    inner_range = UOp.range(4, 1)
+    outer_range = UOp.range(dtypes.int32, 8, 0)
+    inner_range = UOp.range(dtypes.int32, 4, 1)
     expr = (outer_range * 10) + inner_range
     optimized_reduce_uop = apply_rewrite(expr.reduce(Ops.ADD, outer_range, inner_range))
     self.assertEqual(optimized_reduce_uop.op, Ops.CONST)

test/unit/test_hashing.py

@@ -2,6 +2,7 @@ from typing_extensions import Callable
 import hashlib, random, unittest
 from tinygrad import Tensor, Device, getenv, dtypes
 from tinygrad.device import is_dtype_supported
+from tinygrad.helpers import CI
 
 @unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64")
 @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI")
@@ -11,7 +12,7 @@ class TestHashing(unittest.TestCase):
     chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks]
     return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
 
-  @unittest.skip("very slow")
+  @unittest.skipIf(CI, "very slow")
   def test_abc(self):
     expected = self._python_hash_1mb(b"abc" + b"\x00" * (2**20 - 3))
     out = Tensor(b"abc").hash()
@@ -64,7 +65,7 @@ class TestKeccak(unittest.TestCase):
     data = b"\x00" * 4
     self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
 
-    data = b"\x00" * 1000
+    data = b"\x00" * (1000 if CI else 4096)
     self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
 
 if __name__ == "__main__":

test/unit/test_rewrite_tracked_childen.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor
 from tinygrad.uop.ops import PatternMatcher, Ops, UPat, graph_rewrite, RewriteContext, UOp
-from tinygrad.schedule.kernelize import kernelize_sym, merge_views
+from tinygrad.schedule.kernelize import sym, merge_views
 
 class TestRewriteTrackedChildren(unittest.TestCase):
   @unittest.skip("track_children no longer supported")
@@ -57,7 +57,7 @@ class TestRewriteTrackedChildren(unittest.TestCase):
     extra = PatternMatcher([(UPat(Ops.REDUCE_AXIS, name="r"), print_children)])
     a = Tensor.empty(3, 3)
     r = (a+0).sum()
-    graph_rewrite(r.uop, merge_views+kernelize_sym+extra, track_children=True)
+    graph_rewrite(r.uop, merge_views+sym+extra, track_children=True)
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_shapetracker.py

@@ -6,7 +6,7 @@ from tinygrad.helpers import prod
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad import Variable
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
-from tinygrad.codegen.late.devectorizer import sym
+from tinygrad.codegen.devectorizer import sym
 from itertools import product
 
 def shapetracker_getitem(st:ShapeTracker, val:int):
@@ -839,22 +839,25 @@ class TestRender(unittest.TestCase):
     self.assertEqual(idx.render(), "((ridx0*3)+ridx1)")
     self.assertEqual(valid.render(), "(ridx0<2)")
 
-class TestVariableShrink(unittest.TestCase):
-  def test_shrink(self):
-    st = ShapeTracker.from_shape((10,))
-    st = st.shrink(((0, Variable("i", 1, 10)),))
+class TestVariableReshape(unittest.TestCase):
+  def test_reshape(self):
+    st = ShapeTracker.from_shape((3,))
+    st = st.reshape((Variable("i", 1, 10),))
     assert len(st.views) == 1
 
-  def test_shrink_bound(self):
-    st = ShapeTracker.from_shape((10,))
-    st = st.shrink(((0, Variable("i", 1, 10).bind(3)),))
+  def test_reshape_stride_0(self):
+    st = ShapeTracker.from_shape((3,), (0,))
+    st = st.reshape((Variable("i", 1, 10).bind(3),))
+    assert len(st.views) == 1, f"multiview {st}"
+
+  def test_reshape_bound(self):
+    st = ShapeTracker.from_shape((3,))
+    st = st.reshape((Variable("i", 1, 10).bind(3),))
     assert len(st.views) == 1
 
-class TestVariableMerge(unittest.TestCase):
-  def test_add_reshape(self):
-    vi = Variable("i", 1, 10)
-    st1 = ShapeTracker.from_shape((vi,))
-    st2 = ShapeTracker.from_shape((1, vi,))
+  def test_add(self):
+    st1 = ShapeTracker.from_shape((3,))
+    st2 = ShapeTracker.from_shape((Variable("i", 1, 10),))
     st = st1+st2
     assert len(st.views) == 1
 
@@ -864,17 +867,15 @@ class TestVariableMerge(unittest.TestCase):
     st = st1+st2
     assert len(st.views) == 1, f"multiview {st}"
 
-  def test_add_reshape_bound(self):
-    vi = Variable("i", 1, 10).bind(3)
-    st1 = ShapeTracker.from_shape((vi,))
-    st2 = ShapeTracker.from_shape((1, vi,))
+  def test_add_bound(self):
+    st1 = ShapeTracker.from_shape((3,))
+    st2 = ShapeTracker.from_shape((Variable("i", 1, 10).bind(3),))
     st = st1+st2
     assert len(st.views) == 1
 
   def test_simplify(self):
-    vi = Variable("i", 1, 10).bind(3)
-    st1 = ShapeTracker.from_shape((vi,))
-    st2 = ShapeTracker.from_shape((1, vi,))
+    st1 = ShapeTracker.from_shape((3,))
+    st2 = ShapeTracker.from_shape((Variable("i", 1, 10).bind(3),))
     st = ShapeTracker((st1.views[0], st2.views[0]))
     st = st.simplify()
     assert len(st.views) == 1

test/unit/test_shapetracker_math.py

@@ -87,6 +87,20 @@ class TestShapeTrackerAdd(unittest.TestCase):
     assert not (st_equal(st1, st2))
 
 class TestShapeTrackerAddVariable(unittest.TestCase):
+  def test_self_add(self):
+    j = Variable("j", 0, 20).bind(10)
+    a = ShapeTracker.from_shape((10,10))
+    x = a.reshape((10, j))
+    out = x + x
+    assert out == x
+
+  def test_self_add_reshape(self):
+    j = Variable("j", 0, 20).bind(10)
+    a = ShapeTracker.from_shape((10,10))
+    x = a.reshape((10, j))
+    out = x.reshape((5, 2, j)) + x
+    assert out == x
+
   def test_merge_symbolic_views(self):
     var_i = Variable('i', 1, 10)
     var_j = Variable('i', 1, 10)

test/unit/test_simplify_valid_idx.py

@@ -19,7 +19,7 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO
 
 def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.int, (), (expr, nmax))
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
-def Range(n, nmax): return UOp.range(nmax, n)
+def Range(n, nmax): return UOp(Ops.RANGE, dtypes.int, arg=n, src=(UOp.const(dtypes.int, nmax),))
 
 class TestHelpers(unittest.TestCase):
   def test_is_increasing(self):

test/unit/test_symbolic_shapetracker.py

@@ -48,11 +48,11 @@ class TestSymbolic(unittest.TestCase):
     i = Variable("i", 1, 5).bind(3)
     j = Variable("j", 1, 5).bind(3)
     k = Variable("k", 1, 5).bind(3)
-    t = Tensor.rand(5, 4)[:i].cat(Tensor.rand(5, 4)[:j], dim=0).cat(Tensor.rand(5, 4)[:k], dim=0)
+    t = Tensor.rand(3, 4).reshape(i, 4).cat(Tensor.rand(3, 4).reshape(j, 4), dim=0).cat(Tensor.rand(3, 4).reshape(k, 4), dim=0)
     st = t.uop.st
     self.assert_tuple_equal(st.shape, (i+j+k, 4))
     assert st.real_strides() == (4, 1)
-    t = Tensor.rand(5, 3)[:i].cat(Tensor.rand(5, 3)[:i], dim=0).cat(Tensor.rand(3, 3), dim=0)
+    t = Tensor.rand(3, 3).reshape(i, 3).cat(Tensor.rand(3, 3).reshape(i, 3), dim=0).cat(Tensor.rand(3, 3), dim=0)
     st = t.uop.st
     self.assert_tuple_equal(st.shape, (2*i+3, 3))
     assert st.real_strides() == (3, 1)
@@ -61,7 +61,7 @@ class TestSymbolic(unittest.TestCase):
     i = Variable("i", 1, 5).bind(4)
     j = Variable("j", 1, 5).bind(4)
     k = Variable("k", 1, 5).bind(4)
-    t = Tensor.rand(3, 5)[:, :i].cat(Tensor.rand(3, 5)[:, :j], dim=1).cat(Tensor.rand(3, 5)[:, :k], dim=1)
+    t = Tensor.rand(3, 4).reshape(3, i).cat(Tensor.rand(3, 4).reshape(3, j), dim=1).cat(Tensor.rand(3, 4).reshape(3, k), dim=1)
     st = t.uop.st
     self.assert_tuple_equal(st.shape, (3, i+j+k))
     self.assert_tuple_equal(st.real_strides(), (i+j+k, 1))
@@ -109,44 +109,60 @@ class TestShapeTrackerUnbind(unittest.TestCase):
     assert unbound_view == View.create(shape=(v, 4))
     assert var_val == {v: 3}
 
+  def test_reshape_unbind(self):
+    v = Variable("v", 1, 100)
+    bv = Variable("v", 1, 100).bind(3)
+    t = Tensor.rand(3, 4).reshape(bv, 4)
+    unbound_st, var_val = t.uop.st.unbind()
+    assert unbound_st == ShapeTracker((View.create(shape=(v, 4)),))
+    assert var_val == {v: 3}
+
   def test_shrink_unbind(self):
     v = Variable("v", 1, 100)
     bv = Variable("v", 1, 100).bind(2)
-    t = Tensor.rand(3, 4).shrink(((0,bv),(0,4)))
-    unbound_st, var_val = t.uop.st.unbind()
-    assert unbound_st == ShapeTracker((View.create(shape=(v, 4)),))
-    assert var_val == {v: 2}
     t = Tensor.rand(3, 4).shrink(((bv, bv+1), (0, 4)))
     unbound_st, var_val = t.uop.st.unbind()
     assert unbound_st == ShapeTracker((View.create(shape=(1, 4), offset=4*v),))
     assert var_val == {v: 2}
 
-class TestSymbolicReshape(unittest.TestCase):
-  def test_reshape(self):
-    a = Tensor.rand(5, 4)
-    b = Tensor.rand(5, 6)
+class TestSymbolicReshapeFromContiguous(unittest.TestCase):
+  def test_reshape_into_symbols_simple(self):
     for i in range(1, 6):
       vi = Variable("i", 1, 5).bind(i)
-      ret = a[:vi]
-      ret = ret.reshape((vi, 4))
-      assert ret.shape == (vi, 4)
-      ret = b[:vi]
-      ret = ret.reshape((vi, 2, 3))
-      assert ret.shape == (vi, 2, 3)
+      t = Tensor.rand(i, 4).reshape(vi, 4)
+      assert t.shape == (vi, 4)
+      t = Tensor.rand(i, 6).reshape(vi, 2, 3)
+      assert t.shape == (vi, 2, 3)
+
+  def test_reshape_symbols_reshape_ints(self):
+    for i in range(1, 6):
+      vi = Variable("i", 1, 5).bind(i)
+      t = Tensor.rand(i, 4).reshape(vi, 4)
+      assert t.shape == (vi, 4)
+      t = t.reshape(i, 4)
+      assert t.shape == (i, 4)
+
+  @unittest.skip("works now")
+  def test_reshape_into_symbols_bad_shape(self):
+    vi = Variable("i", 1, 10).bind(4)
+    # TODO: this never actually worked, it relied on lazy
+    #with self.assertRaises(ValueError):
+    #  Tensor.rand(4, 6).reshape(vi, 6).reshape(1, 77) # reshape to a different size new shape through symbolic shape
+    with self.assertRaises(AssertionError):
+      Tensor.rand(3, 4).reshape(3, (vi+1)) # reshape into non-Variable Node
 
   def test_two_symbol_reshape(self):
-    t = Tensor.rand(5, 5)
     for i in range(1, 6):
       for j in range(1, 6):
         vi = Variable("i", 1, 5).bind(i)
         vj = Variable("j", 1, 5).bind(j)
-        ret = t[:vi, :vj]
-        ret = ret.reshape(vj, vi)
-        assert ret.shape == (vj, vi)
-        ret = ret.reshape(vi, vj)
-        assert ret.shape == (vi, vj)
-        ret = ret.reshape(1, vi*vj)
-        assert ret.shape == (1, vi*vj)
+        t = Tensor.rand(i, j).reshape(vi, vj)
+        assert t.shape == (vi, vj)
+        # NOTE: this is currently not allowed
+        # t = t.reshape(1, vi*vj)
+        # assert t.shape == (1, vi*vj)
+        t = t.reshape(vj, vi)
+        assert t.shape == (vj, vi)
 
   def test_symbolic_mask(self):
     # taken from gpt2 single kvcache
@@ -159,6 +175,41 @@ class TestSymbolicReshape(unittest.TestCase):
     new_shape = (2, (Variable('start_pos', 1, 128)+1), 16, 64)
     assert view.reshape(new_shape) is None
 
+class TestSymbolicReshapeFromNonContiguous(unittest.TestCase):
+  def test_reshape_from_const(self):
+    vi = Variable("i", 1, 5).bind(4)
+    t = Tensor.ones(3, 4).reshape(3, vi)
+    assert t.shape == (3, vi)
+    assert not t.uop.st.contiguous
+    assert len(t.uop.st.views) == 1
+
+  def test_reshape_not_allowed(self):
+    vi = Variable("i", 1, 5).bind(4)
+    with self.assertRaises(ValueError):
+      # different shape length  # TODO: cases where contractions matched might be fine
+      Tensor.ones(3, 4, 1).reshape(3, vi)
+    with self.assertRaises(ValueError):
+      # size matched, but dimensions do not match
+      Tensor.ones(4, 3).reshape(3, vi)
+
+  def test_reshape_from_padded(self):
+    vi = Variable("i", 1, 5).bind(4)
+    t = Tensor.ones(3, 4).contiguous().expand(2, 3, 4).pad(((1, 1), None, None)).shrink((None, None, (1, 3)))
+    st = t.uop.st
+    assert len(st.views) == 1
+    view = st.views[0]
+    assert view.shape == (4, 3, 2)
+    t = t.reshape(vi, 3, 2)
+    st2 = t.uop.st
+    assert len(st2.views) == 1
+    view2 = st2.views[0]
+    # check only shape changed. strides, offset, mask, contiguous remained the same
+    assert view2.shape == (vi, 3, 2)
+    assert view.strides == view2.strides == (0, 4, 1)
+    assert view.offset == view2.offset == 1
+    assert view.mask == view2.mask == ((1, 3), (0, 3), (0, 2))
+    assert not view.contiguous and not view2.contiguous
+
 class TestSymbolicExpand(unittest.TestCase):
   def test_expand_into_symbols(self):
     vi = Variable("i", 1, 5).bind(3)
@@ -169,12 +220,11 @@ class TestSymbolicExpand(unittest.TestCase):
     assert a.shape == (3, vi, vj)
 
   def test_plus_expands_constant(self):
-    a = Tensor.rand(3, 5)
     for i in range(1, 6):
       vi = Variable("i", 1, 5).bind(i)
-      ret = a[:, :vi]
-      ret = ret + 1
-      self.assertTupleEqual(ret.shape, (3, vi))
+      a = Tensor.rand(3, i).reshape(3, vi)
+      a = a + 1
+      self.assertTupleEqual(a.shape, (3, vi))
 
   def test_pad_then_expand_into_symbols(self):
     vi = Variable("i", 1, 10).bind(3)
@@ -184,11 +234,6 @@ class TestSymbolicExpand(unittest.TestCase):
     self.assertEqual(a.reshape(vi*25).shape, (vi*25,))
 
 class TestSymbolicShrink(unittest.TestCase):
-  def test_shrink_symbols_simple(self):
-    vi = Variable("i", 1, 5)
-    t = Tensor.rand(5, 5).shrink(((0, 5),(0,vi)))
-    assert t.shape == (5, vi)
-
   def test_shrink_symbols(self):
     vi = Variable("i", 1, 5)
     t = Tensor.rand(3, 5).shrink(((0, 2), (vi, vi+1)))
@@ -197,10 +242,10 @@ class TestSymbolicShrink(unittest.TestCase):
 class TestSymbolicPad(unittest.TestCase):
   def test_pad(self):
     v = Variable("v", 1, 100).bind(5)
-    t = Tensor.ones(100)[:v].pad(((4, 0),))
-    t = t.reshape(9)
-    assert t.tolist() == [0,0,0,0,1,1,1,1,1]
-
+    t = Tensor.ones(5).reshape(v).pad(((4, 0),)).reshape(9)
+    assert t.shape == (9,)
+    st = t.uop.st
+    print(st)
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_tensor_uop_representation.py

@@ -97,7 +97,7 @@ class TestTensorUopRepresentation(unittest.TestCase):
     is_pattern(a, UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER),)))
     vi = UOp.variable("i", 1, 3).bind(1)
     a = Tensor.empty(3, vi)
-    is_pattern(a, UPat(Ops.RESHAPE, src=(UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),))),))
+    is_pattern(a, UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER),)))
     self.assertEqual(a.uop.base.buffer.size, 9)
 
 if __name__ == '__main__':

test/unit/test_transcendental_helpers.py

@@ -2,7 +2,7 @@ import unittest, math
 import numpy as np
 from tinygrad import dtypes
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.uop.decompositions import TRANSCENDENTAL_DTYPES, payne_hanek_reduction, cody_waite_reduction
+from tinygrad.uop.decompositions import TRANSCENDENTAL_SUPPORTED_DTYPES, payne_hanek_reduction, cody_waite_reduction
 from tinygrad.uop.decompositions import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if
 from test.helpers import eval_uop
 
@@ -89,7 +89,7 @@ class TestTranscendentalVectorizedFunctions(unittest.TestCase):
       assert u1.op == u2.op, f'expected {u1.op=} but got {u2.op=} for UOps\n{u1=}\n{u2}'
     [self._check_uops_match(x1, x2) for x1, x2 in zip((u1 if isinstance(u1, tuple) else u1.src), (u2 if isinstance(u2, tuple) else u2.src))]
 
-  def _test_vectorized(self, fxn, scalar_dtypes=TRANSCENDENTAL_DTYPES, vals=[-2,1.3,194], vcounts=[1,4,19]):
+  def _test_vectorized(self, fxn, scalar_dtypes=TRANSCENDENTAL_SUPPORTED_DTYPES, vals=[-2,1.3,194], vcounts=[1,4,19]):
     for scalar_dtype in scalar_dtypes:
       for val in vals:
         for vcount in vcounts:

test/unit/test_uop_spec.py

@@ -81,16 +81,5 @@ class TestUOpSpec(unittest.TestCase):
     with self.assertRaisesRegex(RuntimeError, "UOp verification failed"):
       type_verify([a], tensor_uop_spec)
 
-class TestUOpSink(unittest.TestCase):
-  def test_0(self):
-    s = UOp.sink()
-    self.assertEqual(len(s.src), 0)
-
-  def test_1(self):
-    a = UOp.const(dtypes.int, 0)
-    s1 = UOp.sink(a)
-    s2 = a.sink()
-    self.assertIs(s1, s2)
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_uop_symbolic.py

@@ -4,11 +4,11 @@ import z3
 
 from tinygrad.dtype import dtypes, ConstType
 from tinygrad.codegen import full_rewrite
-from tinygrad.codegen.late.devectorizer import sym
+from tinygrad.codegen.devectorizer import sym
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad import Variable
-from tinygrad.uop.spec import uops_to_z3
+from tinygrad.uop.spec import z3_renderer
 
 def render(self) -> tuple[str, ConstType, ConstType]:
   # NOTE: we need STORE so the ALU op has children
@@ -30,15 +30,16 @@ class TestSymbolicPickle(unittest.TestCase):
 
 class TestSymbolic(unittest.TestCase):
   def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    if test_z3:
-      solver = z3.Solver()
-      expr, expr_simplified = uops_to_z3(solver, v, v.simplify())
-      self.assertEqual(solver.check(expr != expr_simplified), z3.unsat, "simplified expression not equal to original")
     rendered, nmin, nmax = render(v)
     if isinstance(s, tuple): self.assertIn(rendered, s)
     else: self.assertEqual(rendered, s)
     self.assertEqual(nmin, n)
     self.assertEqual(nmax, m)
+    if test_z3:
+      solver = z3.Solver()
+      z3_sink = graph_rewrite(v.sink(v.simplify()), z3_renderer, ctx=(solver, {}))
+      expr, epxr_simplified = z3_sink.src[0].arg, z3_sink.src[1].arg
+      self.assertEqual(solver.check(expr != epxr_simplified), z3.unsat, "simplified expression not equal to original")
 
   def test_cmp_simple(self):
     self.helper_test_variable(Variable("a", 3, 8) < 4, 0, 1, "(a<4)")
@@ -127,8 +128,6 @@ class TestSymbolic(unittest.TestCase):
     b = Variable("b", 0, 8)
     self.helper_test_variable(a+a, 0, 16, "(a*2)")
     self.helper_test_variable((a+b)+b, 0, 24, "(a+(b*2))")
-    self.helper_test_variable((a*3+b)+a, 0, 40, "(b+(a*4))")
-    self.helper_test_variable((a+b)+a*3, 0, 40, "(b+(a*4))")
 
   def test_sub_self(self):
     a = Variable("a", 0, 8)
@@ -163,6 +162,10 @@ class TestSymbolic(unittest.TestCase):
   def test_div_remove(self):
     self.helper_test_variable(Variable("a", 0, 7) // 20, 0, 0, "0")
 
+  def test_div_min_max(self):
+    self.helper_test_variable(Variable("a", 1, 7) // 2, 0, 3, "(a//2)")
+    self.helper_test_variable(Variable("a", 0, 6) // 2, 0, 3, "(a//2)")
+
   def test_div_neg_min_max(self):
     self.helper_test_variable(Variable("a", 1, 7) // -2, -3, 0, "((a//2)*-1)")
     self.helper_test_variable(Variable("a", 0, 6) // -2, -3, 0, "((a//2)*-1)")
@@ -208,18 +211,6 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual((Variable("x", -10, 0)%Variable("y", -10, -1))._min_max, (-9, 0))
     self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (-9, 0))
 
-  def test_div_min_max(self):
-    self.helper_test_variable(Variable("a", 2, 7) // 2, 1, 3, "(a//2)")
-    self.helper_test_variable(Variable("a", 0, 6) // 2, 0, 3, "(a//2)")
-
-    self.helper_test_variable(Variable("x", 0, 10)//Variable("y", 1, 10), 0, 10, "(x//y)")
-    self.helper_test_variable(Variable("x", -10, 0)//Variable("y", 1, 10), -10, 0, "(((x*-1)//y)*-1)")
-    self.helper_test_variable(Variable("x", 0, 10)//Variable("y", -10, -1), -10, 0, "((x//(y*-1))*-1)")
-    self.helper_test_variable(Variable("x", -10, 0)//Variable("y", -10, -1), 0, 10, "((x*-1)//(y*-1))")
-
-    self.helper_test_variable(Variable("x", -10, 10)//Variable("y", 1, 10), -10, 10, "(x//y)")
-    self.helper_test_variable(Variable("x", -10, 10)//Variable("y", -10, -1), -10, 10, "((x//(y*-1))*-1)")
-
   def test_mod_factor(self):
     self.helper_test_variable(usum([Variable("a", 0, 7)*100, Variable("b", 0, 3)*50]) % 100, 0, 50, "((b%2)*50)")
 
@@ -275,16 +266,6 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(((5*Variable("a", 0, 31)) % 12) % 5, 0, 4, "(((a*5)%12)%5)")
     self.helper_test_variable((Variable("a", 0, 31) % 4) % 12, 0, 3, "(a%4)")
 
-  def test_mod_mod_wrong_sign(self):
-    v1=Variable("v1", 0, 128)
-    v3=Variable("v3", 0, 7)
-    self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), -4, 4, "(((((v1%2)*2)+((v3+-1)%5))+-2)%5)")
-
-  def test_mod_mod_wrong_sign2(self):
-    v2=Variable("v2", 0, 8)
-    v3=Variable("v3", 0, 4)
-    self.helper_test_variable((((((v3+3)%7)+(v2+-2))%7)%7), -6, 6, "(((v2+((v3+3)%7))+-2)%7)")
-
   def test_mul_mul(self):
     self.helper_test_variable((Variable("a", 0, 5)*10)*9, 0, 5*10*9, "(a*90)")
 
@@ -394,17 +375,6 @@ class TestSymbolic(unittest.TestCase):
   def test_mul_div(self):
     self.helper_test_variable((Variable("a", 0, 10)*4)//4, 0, 10, "a")
 
-  def test_div_drop_small_terms(self):
-    # from openpilot, shouldnt simplify
-    gidx0 = UOp.variable("gidx0", 0, 10)
-    gidx1 = UOp.variable("gidx1", 0, 10)
-    lidx0 = UOp.variable("lidx0", 0, 1)
-    lidx1 = UOp.variable("lidx1", 0, 1)
-    ridx1005 = UOp.variable("ridx1005", 0, 2)
-    ridx1006 = UOp.variable("ridx1006", 0, 2)
-    self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -2, 20,
-      "(((((lidx1+(((gidx1*18)+(ridx1005*18))+(lidx0*162)))+(gidx0*2))+(ridx1006*2))+-40)//18)")
-
   def test_add_div(self):
     # careful about the lower bounds and upper bounds
     self.helper_test_variable((Variable("a", 0, 5)-2)//4, 0, 0, "0")
@@ -449,13 +419,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((-Variable("a", 10, 10))%7, -3, -3, "-3")
 
   def test_div_numerator_negative(self):
-    with Context(CORRECT_DIVMOD_FOLDING=1):
-      self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -8, 0, "(((idx*10)//11)*-1)")
-
-  def test_nest_div_negative_factor(self):
-    ridx0=UOp.variable("ridx0", 0, 9)
-    ridx1=UOp.variable("ridx1", 0, 6)
-    self.helper_test_variable(((((ridx0*-7)+ridx1)+63)//35), 0, 1, "(((ridx0//5)*-1)+1)")
+    self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -8, 0, "(((idx*10)//11)*-1)")
 
   def test_div_into_mod(self):
     self.helper_test_variable((Variable("idx", 0, 16)*4)%8//4, 0, 1, "(idx%2)")
@@ -632,23 +596,20 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(cond, 0, 1, "(a<2)")
     self.helper_test_variable(cond.where(u1, u0), 0, 1, "(a<2)")
     self.helper_test_variable(cond.where(u1, u0).where(u1, u0), 0, 1, "(a<2)")
-    self.helper_test_variable(cond.where(u0, u1), 0, 1, "((a<2)!=True)")
-    self.helper_test_variable(cond.where(u0, u1).where(u0, u1), 0, 1, "(a<2)")
 
   def test_where_combine(self):
     cond = Variable("x", 0, 3) < 2
     a = Variable("a", 0, 3)
     b = Variable("b", 0, 3)
-    c = Variable("c", 0, 3)
     aa = cond.where(a, a.ufix(0))
     bb = cond.where(b, b.ufix(1))
     self.helper_test_variable(aa, 0, 3, "(a if (x<2) else 0)")
     self.helper_test_variable(bb, 0, 3, "(b if (x<2) else 1)")
     self.helper_test_variable(aa+bb, 0, 6, "((a+b) if (x<2) else 1)")
     self.helper_test_variable(aa.maximum(bb), 0, 3, "(max(a, b) if (x<2) else 1)")
-    self.helper_test_variable((c+aa)+bb, 0, 9, "(c+((a+b) if (x<2) else 1))")
 
     # not combining because it increased total ALU
+    c = Variable("c", 0, 3)
     cc = cond.where(c, c+1)
     self.helper_test_variable(bb+cc, 0, 7, "((b if (x<2) else 1)+(c if (x<2) else (c+1)))")
 
@@ -718,10 +679,6 @@ class TestSymbolic(unittest.TestCase):
     # TODO: should z3 work?
     self.helper_test_variable(2*(2*a).reciprocal(), -math.inf, math.inf, "(1/a)", test_z3=False)
 
-  def test_trunc_noop(self):
-    a = Variable("a", 1, 10, dtypes.int)
-    self.helper_test_variable(a.trunc(), 1, 10, "a", test_z3=False)
-
 class TestSymbolicNumeric(unittest.TestCase):
   def helper_test_numeric(self, f):
     MIN, MAX = 0, 10

test/unit/test_viz.py

@@ -1,11 +1,11 @@
-import unittest, decimal, json, struct
+import unittest, decimal, json
 from dataclasses import dataclass
 
 from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher
 from tinygrad.uop.ops import graph_rewrite, track_rewrites, TRACK_MATCH_STATS
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context
+from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent
 from tinygrad.device import Buffer
 
 @track_rewrites(name=True)
@@ -76,7 +76,7 @@ class TestViz(BaseTestViz):
     self.assertEqual(lineno, inner.__code__.co_firstlineno)
 
   def test_exceptions(self):
-    # VIZ tracks rewrites up to and including the error
+    # VIZ tracks rewrites up to the error
     def count_3(x:UOp):
       assert x.arg <= 3
       return x.replace(arg=x.arg+1)
@@ -85,7 +85,7 @@ class TestViz(BaseTestViz):
     with self.assertRaises(AssertionError): exec_rewrite(a, [err_pm])
     lst = get_viz_list()
     err_step = lst[0]["steps"][0]
-    self.assertEqual(err_step["match_count"], 4) # 3 successful rewrites + 1 err
+    self.assertEqual(err_step["match_count"], 3)
 
   def test_default_name(self):
     a = UOp.variable("a", 1, 10)
@@ -240,84 +240,34 @@ class TestVizIntegration(BaseTestViz):
     self.assertEqual(lst[0]["name"], "Schedule 1 Kernel n1")
     self.assertEqual(lst[1]["name"], prg.name)
 
-  def test_metadata_tracing(self):
-    with Context(TRACEMETA=2):
-      a = Tensor.empty(1)
-      b = Tensor.empty(1)
-      metadata = (alu:=a+b).uop.metadata
-      alu.kernelize()
-      graph = next(get_details(tracked_ctxs[0][0]))["graph"]
-    self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1)
-
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
 
-class TinyUnpacker:
-  def __init__(self, buf): self.buf, self.offset = buf, 0
-  def __call__(self, fmt:str) -> tuple:
-    ret = struct.unpack_from(fmt, self.buf, self.offset)
-    self.offset += struct.calcsize(fmt)
-    return ret
-
-# 0 means None, otherwise it's an enum value
-def option(i:int) -> int|None: return None if i == 0 else i-1
-
-def load_profile(lst:list[ProfileEvent]) -> dict:
-  ret = get_profile(lst)
-  u = TinyUnpacker(ret)
-  dur, global_peak, index_len, layout_len = u("<IQII")
-  strings, dtypes = json.loads(ret[u.offset:u.offset+index_len]).values()
-  u.offset += index_len
-  layout:dict[str, dict] = {}
-  for _ in range(layout_len):
-    klen = u("<B")[0]
-    k = ret[u.offset:u.offset+klen].decode()
-    u.offset += klen
-    layout[k] = v = {"shapes":[]}
-    event_type, event_count = u("<BI")
-    if event_type == 0:
-      for _ in range(event_count):
-        name, ref, st, dur, cat, _ = u("<IIIfBI")
-        v["shapes"].append({"name":strings[name], "ref":option(ref), "st":st, "dur":dur, "cat":option(cat)})
-    else:
-      v["peak"] = u("<Q")[0]
-      for _ in range(event_count):
-        alloc, ts, key = u("<BII")
-        if alloc: v["shapes"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-        else: v["shapes"].append({"event":"free", "ts":ts, "key":key})
-  return {"dur":dur, "peak":global_peak, "layout":layout}
-
 class TestVizProfiler(unittest.TestCase):
   def test_perfetto_node(self):
     prof = [ProfileRangeEvent(device='NV', name='E_2', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=False),
             ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100))]
 
-    j = load_profile(prof)
+    j = json.loads(get_profile(prof))
 
-    dev_events = j['layout']['NV']['shapes']
+    dev_events = j['layout']['NV']['timeline']['shapes']
     self.assertEqual(len(dev_events), 1)
     event = dev_events[0]
     self.assertEqual(event['name'], 'E_2')
     self.assertEqual(event['st'], 0)
     self.assertEqual(event['dur'], 10)
-    assert event['ref'] is None
 
   def test_perfetto_copy_node(self):
     prof = [ProfileRangeEvent(device='NV', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=True),
-            ProfileRangeEvent(device='NV:2', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=True),
-            ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
-            ProfileDeviceEvent(device='NV:2', comp_tdiff=decimal.Decimal(-800), copy_tdiff=decimal.Decimal(-80))]
+            ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100))]
 
-    j = load_profile(prof)
+    j = json.loads(get_profile(prof))
 
-    event = j['layout']['NV']['shapes'][0]
+    event = j['layout']['NV']['timeline']['shapes'][0]
     self.assertEqual(event['name'], 'COPYxx')
-    self.assertEqual(event['st'], 0)   # first event
+    self.assertEqual(event['st'], 900) # diff clock
     self.assertEqual(event['dur'], 10)
 
-    event2 = j['layout']['NV:2']['shapes'][0]
-    self.assertEqual(event2['st'], 20) # second event, diff clock
-
   def test_perfetto_graph(self):
     prof = [ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
             ProfileDeviceEvent(device='NV:1', comp_tdiff=decimal.Decimal(-500), copy_tdiff=decimal.Decimal(-50)),
@@ -326,44 +276,28 @@ class TestVizProfiler(unittest.TestCase):
                               deps=[[], [0]],
                               sigs=[decimal.Decimal(1000), decimal.Decimal(1002), decimal.Decimal(1004), decimal.Decimal(1008)])]
 
-    j = load_profile(prof)
+    j = json.loads(get_profile(prof))
 
-    tracks = list(j['layout'])
-    self.assertEqual(tracks[0], 'NV Graph')
-    self.assertEqual(tracks[1], 'NV')
-    self.assertEqual(tracks[2], 'NV:1')
+    devices = list(j['layout'])
+    self.assertEqual(devices[0], 'NV Graph')
+    self.assertEqual(devices[1], 'NV')
+    self.assertEqual(devices[2], 'NV:1')
 
-    nv_events = j['layout']['NV']['shapes']
+    nv_events = j['layout']['NV']['timeline']['shapes']
     self.assertEqual(nv_events[0]['name'], 'E_25_4n2')
     self.assertEqual(nv_events[0]['st'], 0)
     self.assertEqual(nv_events[0]['dur'], 2)
     #self.assertEqual(j['devEvents'][6]['pid'], j['devEvents'][0]['pid'])
 
-    nv1_events = j['layout']['NV:1']['shapes']
+    nv1_events = j['layout']['NV:1']['timeline']['shapes']
     self.assertEqual(nv1_events[0]['name'], 'NV -> NV:1')
     self.assertEqual(nv1_events[0]['st'], 954)
     #self.assertEqual(j['devEvents'][7]['pid'], j['devEvents'][3]['pid'])
 
-    graph_events = j['layout']['NV Graph']['shapes']
+    graph_events = j['layout']['NV Graph']['timeline']['shapes']
     self.assertEqual(graph_events[0]['st'], nv_events[0]['st'])
     self.assertEqual(graph_events[0]['st']+graph_events[0]['dur'], nv1_events[0]['st']+nv1_events[0]['dur'])
 
-  def test_bytes_per_kernel(self):
-    step = 10
-    n_events = 1_000
-    prof = [ProfileRangeEvent("CPU", name="k_test", st=decimal.Decimal(ts:=i*step), en=decimal.Decimal(ts)+step) for i in range(n_events)]
-    sz = len(get_profile(prof))
-    self.assertLessEqual(sz/n_events, 26)
-
-  # can pack up to 1hr 11 min of trace events
-  def test_trace_duration(self):
-    dur_mins = 72
-    n_events = 1_000
-    step = decimal.Decimal(dur_mins*60*1e6//n_events)
-    prof = [ProfileRangeEvent("CPU", name="k_test", st=decimal.Decimal(ts:=i*step), en=decimal.Decimal(ts)+step) for i in range(n_events)]
-    with self.assertRaises(struct.error):
-      get_profile(prof)
-
 def _alloc(b:int):
   a = Tensor.empty(b, device="NULL", dtype=dtypes.char)
   a.uop.buffer.allocate()
@@ -373,29 +307,38 @@ class TestVizMemoryLayout(BaseTestViz):
   def test_double_alloc(self):
     a = _alloc(1)
     _b = _alloc(1)
-    profile_ret = load_profile(Buffer.profile_events)
-    ret = profile_ret["layout"][f"{a.device} Memory"]
+    profile_ret = json.loads(get_profile(Buffer.profile_events))
+    ret = profile_ret["layout"][a.device]["mem"]
     self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["shapes"]), 2)
+    self.assertEqual(ret["shapes"][0]["x"], [0, 2])
+    self.assertEqual(ret["shapes"][1]["x"], [1, 2])
 
   def test_del_once(self):
     a = _alloc(1)
     del a
     b = _alloc(1)
-    profile_ret = load_profile(Buffer.profile_events)
-    ret = profile_ret["layout"][f"{b.device} Memory"]
+    profile_ret = json.loads(get_profile(Buffer.profile_events))
+    ret = profile_ret["layout"][b.device]["mem"]
     self.assertEqual(ret["peak"], 1)
-    self.assertEqual(len(ret["shapes"]), 3)
+    self.assertEqual(ret["shapes"][0]["x"], [0, 2])
+    self.assertEqual(ret["shapes"][1]["x"], [2, 3])
+    self.assertEqual(ret["shapes"][0]["y"], [0, 0])
+    self.assertEqual(ret["shapes"][1]["y"], [0, 0])
 
   def test_alloc_free(self):
     a = _alloc(1)
     _b = _alloc(1)
     del a
     c = _alloc(1)
-    profile_ret = load_profile(Buffer.profile_events)
-    ret = profile_ret["layout"][f"{c.device} Memory"]
+    profile_ret = json.loads(get_profile(Buffer.profile_events))
+    ret = profile_ret["layout"][c.device]["mem"]
     self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["shapes"]), 4)
+    self.assertEqual(ret["shapes"][0]["x"], [0, 3])
+    self.assertEqual(ret["shapes"][1]["x"], [1, 3, 3, 4])
+    self.assertEqual(ret["shapes"][0]["y"], [0, 0])
+    self.assertEqual(ret["shapes"][1]["y"], [1, 1, 0, 0])
+    self.assertEqual(ret["shapes"][2]["x"], [3, 4])
+    self.assertEqual(ret["shapes"][2]["y"], [1, 1])
 
 if __name__ == "__main__":
   unittest.main()

test_driven_development.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+python3 test/external/process_replay/reset.py
+CAPTURE_PROCESS_REPLAY=1 pytest -n auto test/test_tiny.py test/test_uop_graph.py test/test_ops.py test/test_linearizer.py
+while true; do
+  if python3 test/test_tiny.py; then
+    PYTHONPATH="." python3 test/external/process_replay/process_replay.py
+  fi
+done

tinygrad/apps/llm.py

@@ -53,15 +53,17 @@ class SimpleTokenizer:
     try: return [ self._normal_tokens[p] for p in parts ]
     except KeyError: raise RuntimeError("token not found")
 
-def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor:
+def apply_rope(x:Tensor, start_pos:int|UOp, base:int=10000):
   B, H, T, Hd = x.shape
-  assert (Hd & 1) == 0, "RoPE requires an even head dimension"
-  half = Hd // 2
-  angles = (Tensor.arange(T, dtype="float32") + start_pos)[:, None] * (base ** (-(Tensor.arange(half, dtype="float32") / half)))[None, :]
-  cos, sin = angles.cos().reshape(1, 1, T, half).cast(x.dtype), angles.sin().reshape(1, 1, T, half).cast(x.dtype)
-  x_pairs = x.reshape(B, H, T, half, 2)
-  return Tensor.stack(x_pairs[..., 0] * cos - x_pairs[..., 1] * sin,
-                      x_pairs[..., 0] * sin + x_pairs[..., 1] * cos, dim=-1).reshape(B, H, T, Hd)
+  # NOTE: this is usually in a RoPE cache, but tinygrad JIT should prune it outside the kernel
+  # TODO: make it do that
+  freq = base ** (-Tensor.arange(0, 1, 2/Hd, dtype='float32'))
+  angles = Tensor.arange(start_pos, start_pos+T, dtype='float32')[None, None, :, None] * freq
+  cos, sin = angles.cos(), angles.sin()
+  x = x.reshape(B, H, T, Hd // 2, 2)    # split into pairs
+  y1 = x[..., 0] * cos - x[..., 1] * sin
+  y2 = x[..., 0] * sin + x[..., 1] * cos
+  return Tensor.stack(y1, y2, dim=-1).reshape(B, H, T, Hd)
 
 class TransformerBlock:
   def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, max_context:int=0):

tinygrad/codegen/__init__.py

@@ -1,7 +1,7 @@
 from typing import Any, Callable
 import functools
 from dataclasses import dataclass
-from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, RANGEIFY, POSTOPT
+from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp
 from tinygrad.uop.spec import type_verify
 from tinygrad.renderer import Renderer
@@ -12,14 +12,12 @@ from tinygrad.codegen.quantize import pm_quant
 from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
 from tinygrad.uop.decompositions import get_late_rewrite_patterns
-from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_expander
-from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
+from tinygrad.codegen.expander import migrate_indexing, expander
+from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render
-from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
-from tinygrad.codegen.opt import pm_get_optimization, pm_do_optimize
+from tinygrad.codegen.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
+from tinygrad.codegen.opt import pm_optimize
 from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
-from tinygrad.codegen.opt.postrange import pm_postrange_opt
-from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen
 
 @dataclass
 class RewriteStep:
@@ -46,10 +44,10 @@ rewrites_for_linearizer = [
 
 def get_rewrites_for_renderer(opts:Renderer, linearizer:bool=True) -> list[RewriteStep]:
   # cache with the values of the context vars
-  return _get_rewrites_for_renderer(opts, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value, RANGEIFY.value, POSTOPT.value)
+  return _get_rewrites_for_renderer(opts, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value)
 
 @functools.cache
-def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL, _RANGEIFY, _POSTOPT) -> list[RewriteStep]:
+def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
   # ** lowerer (rewrite_shapetracker_with_index) **
   ret: list[RewriteStep] = []
 
@@ -57,28 +55,22 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   ret.extend(rewrites_for_views)
 
   # this is kernel.py
-  if not _RANGEIFY: ret.append(RewriteStep(pm_get_optimization, ctx=lambda _: opts, name="get optimization"))
-  if not _POSTOPT and not _RANGEIFY: ret.append(RewriteStep(pm_do_optimize, ctx=lambda _: opts, name="optimize ast"))
+  ret.append(RewriteStep(pm_optimize, ctx=lambda _: opts, name="optimize ast"))
 
   if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
   ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))
 
-  if _POSTOPT or _RANGEIFY: ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
-
   # ** expander (expand_rewrite) **
   ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic"))
 
   # expand
-  ret.append(RewriteStep(sym+pm_pre_expander+expander, name="expander"))
-
-  # add locals
-  ret.append(RewriteStep(pm_add_buffers_local+rangeify_codegen, name="add local buffers"))
+  ret.append(RewriteStep(sym+expander, name="expander"))
 
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce
   ret.append(RewriteStep(pm_reduce+gep_pushing, lambda _: ReduceContext(), name="remove_reduce"))
 
-  # add gpu dims (late). this works after devectorize, but it's faster here
+  # add gpu dims (late)
   ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
 
   # devectorize (TODO: does this need opts?)
@@ -95,7 +87,7 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
 
   # decompositions
   pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, _TRANSCENDENTAL>=2)
-  ret.append(RewriteStep(pm_decomp, lambda _: opts.device, name="decompositions"))
+  ret.append(RewriteStep(pm_decomp, name="decompositions"))
 
   # final rules for the renderer (without sym)
   pm_final_rewrite = pm_decomp+pm_render+extra_matcher