No extra vars call (#3054)

* remove unused reciprocal

* comment

* remove unneeded call to vars

* free speedup

.flake8

@@ -1,8 +0,0 @@
-[flake8] # https://flake8.pycqa.org/en/6.0.0/user/options.html#cmdoption-flake8-select
-filename =
-    */tinygrad/*.py,
-    */test/*.py
-select=F,W6,E71,E72,E112,E113,E124,E203,E272,E303,E304,E502,E702,E703,E731,W191
-indent-size=2
-per-file-ignores =
-    test/*: F401, F403, F405, F541, E722, E731, F811, F821, F841

.github/workflows/benchmark.yml

@@ -10,36 +10,44 @@ jobs:
   testmacbenchmark:
     name: Mac Benchmark
     runs-on: [self-hosted, macOS]
+    defaults:
+      run:
+        shell: bash -o pipefail {0}
     if: github.repository_owner == 'tinygrad'
     env:
       PYTHONPATH: .
     steps:
     - name: Checkout Code
       uses: actions/checkout@v3
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s ~/tinygrad/disassemblers/applegpu disassemblers/applegpu
+        ln -s ~/tinygrad/weights/sd-v1-4.ckpt weights/sd-v1-4.ckpt
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+    # TODO: why is this test not reliable?
+    #- name: Run Stable Diffusion
+    #  run: python3 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd.txt
     - name: Run model inference benchmark
-      run: python3 test/external/external_model_benchmark.py
+      run: METAL=1 python3 test/external/external_model_benchmark.py
     - name: Test speed vs torch
       run: BIG=2 MPS=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
     - name: Run Tensor Core GEMM
       run: DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
-    - name: Run Stable Diffusion
-      run: |
-        ln -s ~/tinygrad/weights/sd-v1-4.ckpt weights/sd-v1-4.ckpt
-        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
-        time python3 examples/stable_diffusion.py --noshow
     - name: Run LLaMA
       run: |
-        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
-        python3 examples/llama.py --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
-        JIT=1 python3 examples/llama.py --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
+        JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
     - name: Run GPT2
       run: |
-        python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
+        JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
         JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
     - name: Run 10 CIFAR training steps
-      run: |
-        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
-        STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+    - name: Run 10 CIFAR training steps w winograd
+      run: WINO=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
     - uses: actions/upload-artifact@v3
       with:
         name: Speed (Mac)
@@ -47,15 +55,20 @@ jobs:
           onnx_inference_speed.csv
           torch_speed.txt
           train_cifar.txt
+          train_cifar_wino.txt
           llama_unjitted.txt
           llama_jitted.txt
           gpt2_unjitted.txt
           gpt2_jitted.txt
           matmul.txt
+          sd.txt
 
-  testamdbenchmark:
-    name: AMD Benchmark
-    runs-on: [self-hosted, Linux]
+  testnvidiabenchmark:
+    name: NVIDIA Benchmark
+    runs-on: [self-hosted, Linux, CUDA]
+    defaults:
+      run:
+        shell: bash -o pipefail {0}
     if: github.repository_owner == 'tinygrad'
     env:
       PYTHONPATH: .
@@ -63,29 +76,72 @@ jobs:
     - name: Checkout Code
       uses: actions/checkout@v3
     - name: Run model inference benchmark
-      run: python3 test/external/external_model_benchmark.py
+      run: CUDA=1 python3 test/external/external_model_benchmark.py
+    - name: Test speed vs torch
+      run: CUDA=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
+    - name: Run GPT2
+      run: |
+        CUDA=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
+        CUDA=1 JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+    - name: Run GPT2 w HALF
+      run: CUDA=1 JIT=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
+    - name: Run GPT2 w HALF/BEAM
+      run: CUDA=1 JIT=1 HALF=1 BEAM=4 CACHELEVEL=0 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+    - uses: actions/upload-artifact@v3
+      with:
+        name: Speed (NVIDIA)
+        path: |
+          onnx_inference_speed.csv
+          torch_speed.txt
+          gpt2_unjitted.txt
+          gpt2_jitted.txt
+          gpt2_half_beam.txt
+
+  testamdbenchmark:
+    name: tinybox Benchmark
+    runs-on: [self-hosted, Linux, tinybox]
+    defaults:
+      run:
+        shell: bash -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    env:
+      PYTHONPATH: .
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v3
+    - name: Show off tinybox
+      run: /opt/rocm/bin/rocm-bandwidth-test
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+    - name: Run model inference benchmark
+      run: GPU=1 python3 test/external/external_model_benchmark.py
     - name: Test speed vs torch
       run: BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
     - name: Run Tensor Core GEMM
       run: HIP=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
     - name: Run Stable Diffusion
+      run: python3 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd.txt
+    - name: Run LLaMA (with HIP)
       run: |
-        ln -s ~/tinygrad/weights/sd-v1-4.ckpt weights/sd-v1-4.ckpt
-        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
-        time DEBUG=1 python3 examples/stable_diffusion.py --noshow
-    - name: Run LLaMA
+        HIP=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
+        HIP=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+    - name: Run GPT2 (with HIP)
       run: |
-        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
-        python3 examples/llama.py --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
-        JIT=1 python3 examples/llama.py --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
-    - name: Run GPT2
-      run: |
-        python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        HIP=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
+        HIP=1 JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
     - name: Run 10 CIFAR training steps
-      run: |
-        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
-        STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+    - name: Run 10 CIFAR training steps w HALF
+      run: HALF=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+    # # TODO: make wino faster so we can enable both
+    # - name: Run 10 CIFAR training steps w winograd
+    #   run: WINO=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    # - name: Run 10 CIFAR training steps w WINO/HALF/HIP
+    #   run: HALF=1 WINO=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar_wino_half_hip.txt
     - uses: actions/upload-artifact@v3
       with:
         name: Speed (AMD)
@@ -93,8 +149,12 @@ jobs:
           onnx_inference_speed.csv
           torch_speed.txt
           train_cifar.txt
+          train_cifar_half.txt
+          train_cifar_wino.txt
+          train_cifar_wino_half_hip.txt
           llama_unjitted.txt
           llama_jitted.txt
           gpt2_unjitted.txt
           gpt2_jitted.txt
           matmul.txt
+          sd.txt

.github/workflows/python-publish.yml

@@ -5,13 +5,12 @@ name: Upload Python Package
 
 on:
   release:
-    types: [created]
+    types: [published]
+  workflow_dispatch:
 
 jobs:
   deploy:
-
     runs-on: ubuntu-latest
-
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python

.github/workflows/szdiff.yml

@@ -0,0 +1,98 @@
+name: Check Line Counts
+on:
+  pull_request_target:
+
+# Cancel the workflow in progress in newer build is about to start.
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  checkbranch:
+    name: Check PR Branch status
+    runs-on: ubuntu-latest
+    outputs:
+      branchstat: ${{ steps.brstat.outputs.stat}}
+    steps:
+      - name: Check code from PR branch 
+        uses: actions/checkout@v4
+        with:
+          repository: ${{ github.event.pull_request.head.repo.full_name }}
+          ref: ${{ github.event.pull_request.head.sha }}
+          fetch-depth: 0
+      - name: Check whether branch is up-to-date
+        id: brstat
+        run: |
+          git remote add tinygrad https://github.com/tinygrad/tinygrad
+          git fetch tinygrad master
+          echo "${{ github.event.pull_request.head.sha }}"
+          git rev-list --left-right --count  tinygrad/master...${{ github.event.pull_request.head.sha }} | awk '{print "Behind "$1" - Ahead "$2""}'
+          count=$(git rev-list --left-right --count  tinygrad/master...${{ github.event.pull_request.head.sha }} | awk '{print $1}')
+          if [ $count -gt 0 ]
+          then
+            echo "Current branch is behind tinygrad master branch!"
+            echo "stat=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "stat=false" >> "$GITHUB_OUTPUT"
+          fi
+
+  szdiff:
+    name: Core Library Line Difference
+    permissions:
+      contents: read
+      pull-requests: write
+    runs-on: ubuntu-latest
+    needs: checkbranch
+    if: needs.checkbranch.outputs.branchstat == 'false'
+    steps:
+      - name: Checkout code from PR branch
+        uses: actions/checkout@v4
+        with:
+          repository: ${{ github.event.pull_request.head.repo.full_name }}
+          ref: ${{ github.event.pull_request.head.sha }}
+          path: pr
+        # the base default to tinygrad master and cannot be other fork branch for security purpose
+      - name: Checkout code from tinygrad master
+        uses: actions/checkout@v4
+        with:
+          path: base
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+      - name: Count Line Diff
+        run: |
+          pip install tabulate
+          BASE="$GITHUB_WORKSPACE/base"
+          PR="$GITHUB_WORKSPACE/pr"
+          cp "$BASE/sz.py" .
+          echo "loc_content<<EOF" >> "$GITHUB_ENV"
+          python sz.py "$BASE" "$PR" >> "$GITHUB_ENV"
+          echo "EOF" >> "$GITHUB_ENV"
+      - name: Comment Code Line Diff
+        continue-on-error: false
+        uses: marocchino/sticky-pull-request-comment@v2
+        with:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          ignore_empty: true
+          skip_unchanged: true
+          recreate: true
+          message: ${{ env.loc_content }}
+
+  rebase:
+    name: Core Library Line Difference
+    permissions:
+      pull-requests: write
+    runs-on: ubuntu-latest
+    needs: checkbranch
+    if: needs.checkbranch.outputs.branchstat == 'true'
+    steps:
+      - name: Comment Rebase
+        continue-on-error: false
+        uses: marocchino/sticky-pull-request-comment@v2
+        with:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          skip_unchanged: true
+          recreate: true
+          message: |
+            This branch currently is behind tinygrad/master. The line count difference bot is disabled.

.github/workflows/test.yml

@@ -1,7 +1,12 @@
 name: Unit Tests
+env:
+  # increment this when downloads substantially change to avoid the internet
+  DOWNLOAD_CACHE_VERSION: '4'
 
 on:
   push:
+    branches:
+      - master
   pull_request:
   workflow_dispatch:
 
@@ -11,6 +16,61 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 20
 
+    # TODO: run the pre-commit hook to replace a lot of this
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v3
+    - name: Set up Python 3.8
+      uses: actions/setup-python@v4
+      with:
+        python-version: 3.8
+    - name: Cache python packages
+      uses: actions/cache@v3
+      with:
+        path: ${{ env.Python3_ROOT_DIR }}/lib/python3.8/site-packages
+        key: linting-packages-${{ hashFiles('**/setup.py') }}-3.8
+    - name: Install dependencies
+      run: pip install -e '.[linting,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
+    - name: Lint with pylint
+      run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string='  ' **/*.py
+    - name: Lint with ruff
+      run: |
+        pip3 install --upgrade --force-reinstall ruff
+        python3 -m ruff . --preview
+    - name: Lint tinygrad with pylint
+      run: python -m pylint tinygrad/
+    - name: Run mypy
+      run: python -m mypy
+    - name: Test Docs
+      run: |
+        python docs/abstractions.py
+        python docs/abstractions2.py
+    - name: Test Quickstart
+      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' docs/quickstart.md > quickstart.py &&  PYTHONPATH=. python quickstart.py
+    - name: Fuzz Test symbolic
+      run: python test/external/fuzz_symbolic.py
+    - name: Fuzz Test shapetracker
+      run: |
+        PYTHONPATH="." python test/external/fuzz_shapetracker.py
+        PYTHONPATH="." python test/external/fuzz_shapetracker_math.py
+    - name: Test shapetracker to_movement_ops
+      run: PYTHONPATH="." python extra/to_movement_ops.py
+    - name: Use as an external package
+      run: |
+        mkdir $HOME/test_external_dir
+        cd $HOME/test_external_dir
+        python -m venv venv
+        source venv/bin/activate
+        pip install $GITHUB_WORKSPACE
+        python -c "from tinygrad.tensor import Tensor; print(Tensor([1,2,3,4,5]))"
+    - name: Repo line count <5000 lines
+      run: MAX_LINE_COUNT=5000 python sz.py
+
+  testcpuimagenet:
+    name: CPU and ImageNet to C Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+
     steps:
     - name: Checkout Code
       uses: actions/checkout@v3
@@ -22,59 +82,22 @@ jobs:
       uses: actions/cache@v3
       with:
         path: ${{ env.Python3_ROOT_DIR }}/lib/python3.8/site-packages
-        key: linting-packages-${{ hashFiles('*/setup.py') }}-3.8
-    - name: Install dependencies
-      run: pip install -e '.[linting,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
-    - name: Repo line count
-      run: python sz.py
-    - name: Lint with pylint
-      run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string='  ' **/*.py
-    - name: Lint with flake8
-      run: python -m flake8 . --statistics -j4
-    - name: Lint tinygrad with pylint
-      run: python -m pylint tinygrad/
-    - name: Run mypy
-      run: python -m mypy tinygrad/ --ignore-missing-imports --check-untyped-defs --explicit-package-bases --warn-unreachable
-    - name: Install SLOCCount
-      run: sudo apt install sloccount
-    - name: Check <5000 lines
-      run: sloccount tinygrad test examples extra; if [ $(sloccount tinygrad | sed -n 's/.*Total Physical Source Lines of Code (SLOC)[ ]*= \([^ ]*\).*/\1/p' | tr -d ',') -gt 5000 ]; then exit 1; fi
-
-  testcpuimagenet:
-    name: CPU and ImageNet to C Tests
-    runs-on: ubuntu-latest
-    timeout-minutes: 20
-
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v3
-    - name: Set up Python 3.11
-      uses: actions/setup-python@v4
-      with:
-        python-version: 3.11
-    - name: Cache python packages
+        key: testing-packages-${{ hashFiles('**/setup.py') }}
+    - name: Cache downloads
       uses: actions/cache@v3
       with:
-        path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
-        key: testing-packages-${{ hashFiles('*/setup.py') }}
+        path: ~/.cache/tinygrad/downloads/
+        key: downloads-cache-cpu-${{ env.DOWNLOAD_CACHE_VERSION }}
     - name: Install Dependencies
       run: pip install -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
-    - name: Test Docs
-      run: python docs/abstractions.py
-    - name: Test Quickstart
-      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' docs/quickstart.md > quickstart.py &&  PYTHONPATH=. python quickstart.py
     - name: Run Pytest
-      run: python -m pytest -n=auto test/ -k "not (test_efficientnet and models/test_train.py)"
-    - name: Fuzz Test symbolic
-      run: python test/external/fuzz_symbolic.py
-    - name: Fuzz Test shapetracker
-      run: PYTHONPATH="." python test/external/fuzz_shapetracker.py
+      run: python -m pytest -n=auto test/ -k "not (test_efficientnet and models/test_train.py)" --durations=20
     - name: Compile EfficientNet to C
       run: PYTHONPATH="." CLANG=1 python examples/compile_efficientnet.py > recognize.c
     - name: Compile C to native
       run: clang -O2 recognize.c -lm -o recognize
     - name: Test EfficientNet
-      run: curl https://media.istockphoto.com/photos/hen-picture-id831791190 | ./recognize | grep hen
+      run: cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock
 
   testtorch:
     name: Torch Tests
@@ -92,35 +115,40 @@ jobs:
       uses: actions/cache@v3
       with:
         path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
-        key: testing-packages-${{ hashFiles('*/setup.py') }}
+        key: testing-packages-${{ hashFiles('**/setup.py') }}
+    - name: Cache downloads
+      uses: actions/cache@v3
+      with:
+        path: ~/.cache/tinygrad/downloads/
+        key: downloads-cache-torch-${{ env.DOWNLOAD_CACHE_VERSION }}
     - name: Install Dependencies
       run: pip install -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
     - name: Run Pytest
-      run: TORCH=1 python -m pytest -n=auto test/
+      run: TORCH=1 python -m pytest -n=auto test/ --durations=20
     - name: Run ONNX
-      run: TORCH=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --tb=no --disable-warnings || true
+      run: TORCH=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py
 
   testopencl:
     strategy:
       fail-fast: false
       matrix:
-        task: [optimage, openpilot, multigpu]
-    name: ${{ matrix.task=='optimage'&&'GPU OPT and IMAGE Tests'|| matrix.task=='openpilot'&&'openpilot (OpenCL) Tests'|| matrix.task=='multigpu'&&'MultiGPU Tests'}}
+        task: [optimage, openpilot, onnx]
+    name: ${{ matrix.task=='optimage'&&'GPU OPT and IMAGE Tests' || matrix.task=='openpilot'&&'openpilot (OpenCL) Tests' || matrix.task=='onnx'&&'ONNX+Optimization Tests' }}
     runs-on: ubuntu-20.04
     timeout-minutes: 20
 
     steps:
       - name: Checkout Code
         uses: actions/checkout@v3
-      # - name: Find faster apt mirror
-      #   uses: vegardit/fast-apt-mirror.sh@v1
       - name: Install OpenCL
-        #run: sudo apt install -y pocl-opencl-icd
         run: |
-          wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | sudo tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null
-          echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
+          echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
+          echo "deb [ allow-insecure=yes ] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
           sudo apt update
-          sudo apt install -y --no-install-recommends intel-oneapi-runtime-compilers intel-oneapi-runtime-opencl
+          sudo apt install --allow-unauthenticated -y --no-install-recommends \
+            intel-oneapi-runtime-openmp=2023.2.1-16 intel-oneapi-runtime-compilers-common=2023.2.1-16 intel-oneapi-runtime-compilers=2023.2.1-16 \
+            intel-oneapi-runtime-dpcpp-sycl-opencl-cpu=2023.2.1-16 intel-oneapi-runtime-tbb-common=2021.10.0-49541 \
+            intel-oneapi-runtime-tbb=2021.10.0-49541 intel-oneapi-runtime-opencl=2023.2.1-16
       - name: Set up Python 3.11
         uses: actions/setup-python@v4
         with:
@@ -129,7 +157,12 @@ jobs:
         uses: actions/cache@v3
         with:
           path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
-          key: testing-packages-${{ hashFiles('*/setup.py') }}
+          key: testing-packages-${{ hashFiles('**/setup.py') }}
+      - name: Cache downloads
+        uses: actions/cache@v3
+        with:
+          path: ~/.cache/tinygrad/downloads/
+          key: downloads-cache-${{ matrix.task }}-${{ env.DOWNLOAD_CACHE_VERSION }}
       - name: Install Dependencies
         run: pip install -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
       - if: ${{ matrix.task == 'optimage' }}
@@ -138,29 +171,92 @@ jobs:
           PYTHONPATH="." OPT=2 GPU=1 python -m pytest -n=auto test/external/external_test_opt.py
           PYTHONPATH="." OPT=3 GPU=1 python -m pytest -n=auto test/external/external_test_opt.py
       - if: ${{ matrix.task == 'optimage'}}
-        name: Test GPU IMAGE ops
-        run: |
-          GPU=1 IMAGE=1 python -m pytest -n=auto test/test_ops.py
-          FORWARD_ONLY=1 GPU=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py
+        name: Test WINO=1
+        run: GPU=1 DEBUG=2 WINO=1 python3 test/test_ops.py TestOps.test_simple_conv2d
+      - if: ${{ matrix.task == 'optimage'}}
+        name: Test GPU IMAGE=1 ops
+        run: GPU=1 IMAGE=1 python -m pytest -n=auto test/test_ops.py
+      - if: ${{ matrix.task == 'optimage'}}
+        name: Test GPU IMAGE=2 ops
+        run: GPU=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py
       - if: ${{ matrix.task == 'openpilot' }}
         name: Test openpilot model compile and size
         run: |
-          DEBUG=2 ALLOWED_KERNEL_COUNT=209 FLOAT16=1 DEBUGCL=1 GPU=1 IMAGE=2 python openpilot/compile.py
-          python -c 'import os; assert os.path.getsize("/tmp/output.thneed") < 100_000_000'
+          DEBUG=2 ALLOWED_KERNEL_COUNT=208 FLOAT16=1 DEBUGCL=1 GPU=1 IMAGE=2 python openpilot/compile2.py
+          #python -c 'import os; assert os.path.getsize("/tmp/output.thneed") < 100_000_000'
       - if: ${{ matrix.task == 'openpilot' }}
         name: Test openpilot model correctness (float32)
-        run: DEBUGCL=1 GPU=1 IMAGE=2 python openpilot/compile.py
+        run: FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python openpilot/compile2.py
       - if: ${{ matrix.task == 'openpilot' }}
-        name: Test tensor core ops
-        run: GPU=1 TC=2 python -m pytest -n=auto test/test_ops.py
-      - if: ${{ matrix.task == 'multigpu' }}
-        name: Test multigpu
-        run: |
-          PYTHONPATH="." python test/external/dist/test_world.py
-          PYTHONPATH="." python test/external/dist/test_collectives.py
+        name: Test openpilot alt model correctness (float32)
+        run: FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python openpilot/compile2.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
+      - if: ${{ matrix.task == 'openpilot' }}
+        name: Test openpilot fastvits model correctness (float32)
+        run: FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python openpilot/compile2.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      #- if: ${{ matrix.task == 'openpilot' }}
+      #  name: Test multigpu
+      #  run: |
+      #    PYTHONPATH="." python test/external/dist/test_world.py
+      #    PYTHONPATH="." python test/external/dist/test_collectives.py
+      - if: ${{ matrix.task == 'onnx' }}
+        name: Test ONNX (CPU)
+        run: CPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      - if: ${{ matrix.task == 'onnx' }}
+        name: Test ONNX (GPU)
+        run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      - if: ${{ matrix.task == 'onnx' }}
+        name: Test ONNX (CLANG)
+        run: CLANG=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      - if: ${{ matrix.task == 'onnx' }}
+        name: Test Action Space
+        run: PYTHONPATH="." GPU=1 python3 extra/optimization/get_action_space.py
+      - if: ${{ matrix.task == 'onnx' }}
+        name: Test Beam Search
+        run: PYTHONPATH="." GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
 
-  testmetalwebgpu:
-    name: Metal and WebGPU Tests
+  #testwebgpu:
+  #  name: WebGPU Tests
+  #  runs-on: macos-13
+  #  timeout-minutes: 20
+  #  steps:
+  #  - name: Checkout Code
+  #    uses: actions/checkout@v3
+  #  - name: Set up Python 3.11
+  #    uses: actions/setup-python@v4
+  #    with:
+  #      python-version: 3.11
+  #  - name: Cache python packages
+  #    uses: actions/cache@v3
+  #    with:
+  #      path: /Users/runner/Library/Python/3.11/lib/python/site-packages
+  #      key: webgpu-testing-user3-packages-${{ hashFiles('**/setup.py') }}
+  #  - name: Install Dependencies
+  #    run: pip install --user -e '.[webgpu,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
+  #  - name: Cache downloads
+  #    uses: actions/cache@v3
+  #    with:
+  #      path: ~/Library/Caches/tinygrad/downloads/
+  #      key: downloads-cache-webgpu-${{ env.DOWNLOAD_CACHE_VERSION }}
+  #  - name: Check Device.DEFAULT (WEBGPU) and print some source
+  #    run: |
+  #      WEBGPU=1 python -c "from tinygrad import Device; assert Device.DEFAULT == 'WEBGPU', Device.DEFAULT"
+  #      WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+    #- name: Run webgpu pytest
+    #  run: WEBGPU=1 WGPU_BACKEND_TYPE=Metal python -m pytest -n=auto
+  #  - name: Run selected webgpu tests
+  #    run: |
+  #      WEBGPU=1 WGPU_BACKEND_TYPE=Metal python -m pytest -n=auto test/test_ops.py test/test_dtype.py \
+  #      test/test_jit.py test/test_symbolic_ops.py test/test_symbolic_jit.py test/test_linearizer.py \
+  #      test/test_linearizer_failures.py test/test_nn.py
+  #  - name: Build WEBGPU Efficientnet
+  #    run: WEBGPU=1 WGPU_BACKEND_TYPE=Metal python -m examples.compile_efficientnet
+  #  - name: Install Puppeteer
+  #    run: npm install puppeteer
+  #  - name: Run WEBGPU Efficientnet
+  #    run: node test/web/test_webgpu.js
+
+  testmetal:
+    name: Metal Tests
     runs-on: macos-13
     timeout-minutes: 20
 
@@ -174,33 +270,108 @@ jobs:
     - name: Cache python packages
       uses: actions/cache@v3
       with:
-        path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
-        key: metal-webgpu-testing-packages-${{ hashFiles('*/setup.py') }}
+        path: /Users/runner/Library/Python/3.11/lib/python/site-packages
+        key: metal-testing-user3-packages-${{ hashFiles('**/setup.py') }}
     - name: Install Dependencies
-      run: pip install -e '.[metal,webgpu,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
+      run: pip install --user -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
+    - name: Cache downloads
+      uses: actions/cache@v3
+      with:
+        path: ~/Library/Caches/tinygrad/downloads/
+        key: downloads-cache-metal-only-${{ env.DOWNLOAD_CACHE_VERSION }}
+    - name: Check Device.DEFAULT (METAL) and print some source
+      run: |
+        METAL=1 python -c "from tinygrad import Device; assert Device.DEFAULT == 'METAL', Device.DEFAULT"
+        METAL=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+    - name: Run metal test
+      run: METAL=1 python -m pytest -n=auto test/ --ignore=test/external --ignore=test/models --durations=20
+    - name: Run ONNX
+      run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py
+    - name: Test tensor core ops
+      run: METAL=1 TC=2 DEBUG=3 python test/test_ops.py TestOps.test_big_gemm
     - name: Test LLaMA compile speed
       run: PYTHONPATH="." METAL=1 python test/external/external_test_speed_llama.py
-    #- name: Run dtype test
-    #  run: DEBUG=4 METAL=1 python -m pytest -n=auto test/test_dtype.py
-    # dtype test has issues on test_half_to_int8
-    - name: Run metal ops test
-      run: DEBUG=2 METAL=1 python -m pytest -n=auto test/test_ops.py
-    - name: Run JIT test
-      run: DEBUG=2 METAL=1 python -m pytest -n=auto test/test_jit.py
-    - name: Run symbolic shapetracker test
-      run: METAL=1 python -m pytest -n=auto test/test_symbolic_shapetracker.py test/test_symbolic_ops.py test/test_symbolic_jit.py
-    - name: Check Device.DEFAULT
-      run: WEBGPU=1 python -c "from tinygrad.ops import Device; assert Device.DEFAULT == 'WEBGPU', Device.DEFAULT"
-    #- name: Run webgpu pytest
-    #  run: WEBGPU=1 WGPU_BACKEND_TYPE=Metal python -m pytest -n=auto --ignore test/models/ --ignore test/unit/test_example.py --ignore test/extra/test_lr_scheduler.py --ignore test/test_linearizer.py test/
-    #- name: Build WEBGPU Efficientnet
-    #  run: WEBGPU=1 WGPU_BACKEND_TYPE=Metal python -m examples.compile_efficientnet
+
+  testhipcompilation:
+    name: HIP Compilation Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v3
+    - name: Set up Python 3.11
+      uses: actions/setup-python@v4
+      with:
+        python-version: 3.11
+    - name: Cache python packages
+      uses: actions/cache@v3
+      with:
+        path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
+        key: testing-packages-${{ hashFiles('**/setup.py') }}
+    - name: Cache downloads
+      uses: actions/cache@v3
+      with:
+        path: ~/.cache/tinygrad/downloads/
+        key: downloads-cache-hipcompilation-${{ env.DOWNLOAD_CACHE_VERSION }}
+    - name: Install HIP tools
+      run: |
+        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
+        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
+        # ROCm repository for jammy
+        sudo tee /etc/apt/sources.list.d/rocm.list <<'EOF'
+        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/debian jammy main
+        EOF
+        # Prefer packages from the rocm repository over system packages
+        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
+        sudo apt update
+        sudo apt install --no-install-recommends --allow-unauthenticated -y rocm-hip-libraries hip-dev
+    - name: Install Python Dependencies
+      run: pip install -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
+    - name: Test HIP compilation on RDNA3 [gfx1100]
+      # test/test_symbolic_ops.py can't run here, it was comparing empty memory
+      run: |
+        export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/rocm/hip/lib
+        MOCKHIP=1 HIP=1 python -m pytest -s test/test_hip_rdna3.py
+
+#  testwebgl:
+#    name: WebGL Tests
+#    runs-on: ubuntu-latest
+#    timeout-minutes: 20
+#
+#    steps:
+#    - name: Checkout Code
+#      uses: actions/checkout@v3
+#    - name: Set up Python 3.11
+#      uses: actions/setup-python@v4
+#      with:
+#        python-version: 3.11
+#    - name: Cache python packages
+#      uses: actions/cache@v3
+#      with:
+#        path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
+#        key: webgl-testing-packages-${{ hashFiles('**/setup.py') }}
+#    - name: Install Dependencies
+#      run: pip install -e '.[webgl,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
+#    - name: Cache downloads
+#      uses: actions/cache@v3
+#      with:
+#        path: ~/Library/Caches/tinygrad/downloads/
+#        key: downloads-cache-webgl-${{ env.DOWNLOAD_CACHE_VERSION }}
+#    - name: Prepare
+#      run: |
+#        sudo apt-get -y install xvfb
+#        sudo /usr/bin/Xvfb :0 -screen 0 4096x4096x24+32 &
+#    - name: Run selected webgl tests
+#      run: WEBGL=1 python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_jit.py
+#    - name: Build WebGL Efficientnet
+#      run: WEBGL=1 python -m examples.compile_efficientnet
 
   tests:
     strategy:
       fail-fast: false
       matrix:
-        backend: [llvm, clang, gpu, cuda, ptx]
+        backend: [llvm, clang, gpu, cuda] #, triton] #, ptx]
 
     name: Tests on (${{ matrix.backend }})
     runs-on: ${{ matrix.backend == 'gpu'  && 'ubuntu-20.04' || 'ubuntu-latest' }}
@@ -217,82 +388,92 @@ jobs:
         uses: actions/cache@v3
         with:
           path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
-          key: ${{ matrix.backend }}-packages-${{ hashFiles('*/setup.py') }}
+          key: ${{ matrix.backend }}-packages-${{ hashFiles('**/setup.py') }}
+      - name: Cache downloads
+        uses: actions/cache@v3
+        with:
+          path: ~/.cache/tinygrad/downloads/
+          key: downloads-cache-${{ matrix.backend }}-${{ env.DOWNLOAD_CACHE_VERSION }}
       - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'ENABLE_METHOD_CACHE=1\nLLVM=1' || matrix.backend == 'clang' && 'CLANG=1\nENABLED_METHOD_CACHE=1' || matrix.backend == 'gpu' && 'GPU=1' || matrix.backend == 'cuda' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nCUDACPU=1\n' || matrix.backend == 'PTX' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nCUDACPU=1\nPTX=1' }}" >> $GITHUB_ENV
-      - name: Find faster apt mirror
-      #   uses: vegardit/fast-apt-mirror.sh@v1
-      # - name: Install packages (gpu)
+        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'clang' && 'CLANG=1' || matrix.backend == 'gpu' && 'GPU=1' || matrix.backend == 'cuda' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nCUDACPU=1\n' || matrix.backend == 'PTX' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nCUDACPU=1\nPTX=1' || matrix.backend == 'triton' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nCUDACPU=1\nTRITON=1\nTRITON_PTXAS_PATH=/usr/bin/ptxas'}}" >> $GITHUB_ENV
+      - name: Install OpenCL
         if: matrix.backend == 'gpu'
         run: |
-          wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | sudo tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null
-          echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
-          sudo apt update -y
-          sudo apt install -y --no-install-recommends intel-oneapi-runtime-compilers intel-oneapi-runtime-opencl
+          echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
+          echo "deb [ allow-insecure=yes ] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
+          sudo apt update
+          sudo apt install --allow-unauthenticated -y --no-install-recommends \
+            intel-oneapi-runtime-openmp=2023.2.1-16 intel-oneapi-runtime-compilers-common=2023.2.1-16 intel-oneapi-runtime-compilers=2023.2.1-16 \
+            intel-oneapi-runtime-dpcpp-sycl-opencl-cpu=2023.2.1-16 intel-oneapi-runtime-tbb-common=2021.10.0-49541 \
+            intel-oneapi-runtime-tbb=2021.10.0-49541 intel-oneapi-runtime-opencl=2023.2.1-16
       - name: Install packages (cuda)
-        if: matrix.backend == 'cuda' || matrix.backend == 'ptx'
+        if: matrix.backend == 'cuda' || matrix.backend == 'ptx' || matrix.backend == 'triton'
         run: |
+          echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
           sudo apt update -y
-          sudo apt install -y --no-install-recommends git g++ cmake ninja-build llvm-15-dev zlib1g-dev libglew-dev flex bison libfl-dev libboost-thread-dev libboost-filesystem-dev nvidia-cuda-toolkit-gcc
+          sudo apt install -y --no-install-recommends git g++ cmake ninja-build llvm-15-dev zlib1g-dev libglew-dev \
+            flex bison libfl-dev libboost-thread-dev libboost-filesystem-dev nvidia-cuda-toolkit-gcc libzstd-dev
       - name: Cache gpuocelot
-        if: matrix.backend == 'cuda' || matrix.backend == 'ptx'
+        if: matrix.backend == 'cuda' || matrix.backend == 'ptx' || matrix.backend == 'triton'
         id: cache-build
         uses: actions/cache@v3
         env:
           cache-name: cache-gpuocelot-build
         with:
           path: ${{ github.workspace }}/gpuocelot/ocelot
-          key: ubuntu22.04-gpuocelot-19626fc00b6ee321638c3111074269c69050e091
+          key: ubuntu22.04-gpuocelot-18401f4245b27ca4b3af433196583cc81ef84480-rebuild-2
       - name: Clone/compile gpuocelot
-        if: (matrix.backend == 'cuda' || matrix.backend == 'ptx') && steps.cache-build.outputs.cache-hit != 'true'
+        if: (matrix.backend == 'cuda' || matrix.backend == 'ptx' || matrix.backend == 'triton') && steps.cache-build.outputs.cache-hit != 'true'
         run: |
           git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
           cd ${{ github.workspace }}/gpuocelot/ocelot
-          git checkout 19626fc00b6ee321638c3111074269c69050e091
+          git checkout 18401f4245b27ca4b3af433196583cc81ef84480
           mkdir build
           cd build
-          cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF
+          cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF
           ninja
       - name: Install gpuocelot
-        if: matrix.backend == 'cuda' || matrix.backend == 'ptx'
+        if: matrix.backend == 'cuda' || matrix.backend == 'ptx' || matrix.backend == 'triton'
         run: |
           cd ${{ github.workspace }}/gpuocelot/ocelot/build
-          sudo ninja install
+          sudo ninja install -d explain
       - name: Install dependencies
-        run: pip install -e '.[testing${{matrix.backend=='llvm'&&',llvm'||matrix.backend=='cuda'&&',cuda'||matrix.backend=='ptx'&&',cuda'||''}}]' --extra-index-url https://download.pytorch.org/whl/cpu
-      - name: Check Device.DEFAULT
-        run: python -c "from tinygrad.ops import Device; assert Device.DEFAULT in ['LLVM','CLANG','CUDA','GPU'], Device.DEFAULT"
-      - name: Run pytest (not cuda)
-        if: matrix.backend!='cuda' && matrix.backend!='ptx'
-        run: python -m pytest -n=auto test/ -k '${{matrix.backend=='llvm'&&'not (test_nn.py and test_conv_transpose2d)'||'test'}}' -m 'not exclude_${{matrix.backend}}'
-      - name: Run pytest (cuda)
-        if: matrix.backend=='cuda'
-        run: python -m pytest -n=auto test/ -k 'not (half or test_efficientnet_safetensors) and not (test_conv2d and test_tensor.py)' -m 'not exclude_cuda' --ignore=test/external --ignore=test/models
-      - name: Run pytest (ptx)
-        if: matrix.backend=='ptx'
-        run: python -m pytest -n=auto test/ -k 'not (half or test_efficientnet_safetensors) and not (test_conv2d and test_tensor.py)' -m 'not exclude_cuda' --ignore=test/external --ignore=test/models
-
-  testunicorn:
-    name: ARM64 unicorn Test
-    runs-on: ubuntu-latest
-    timeout-minutes: 20
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v3
-      - name: Set up Python 3.11
-        uses: actions/setup-python@v4
-        with:
-          python-version: 3.11
-      - name: Cache python packages
-        uses: actions/cache@v3
-        with:
-          path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
-          key: testing-arm-packages-${{ hashFiles('*/setup.py') }}
-      - name: Install cross-assembler
+        run: pip install -e '.[testing${{matrix.backend=='llvm'&&',llvm'||matrix.backend=='cuda'&&',cuda'||matrix.backend=='ptx'&&',cuda'||matrix.backend=='triton'&&',triton'||''}}]' --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
+      - name: Check Device.DEFAULT and print some source
         run: |
-          sudo apt update -y
-          sudo apt install -y --no-install-recommends gcc-aarch64-linux-gnu
-      - name: Install dependencies
-        run: pip install -e '.[testing,arm]' --extra-index-url https://download.pytorch.org/whl/cpu
-      - name: Test arm
-        run: CI=1 ARM64=1 CLANG=1 python -m pytest -n=auto test/ -k 'not (test_nn.py and (test_conv_transpose2d or test_conv2d))' --ignore=test/models --ignore=test/test_speed_v_torch.py --ignore=test/test_net_speed.py --ignore=test/test_specific_conv.py  --ignore=test/unit/test_disk_tensor.py
+          python -c "from tinygrad import Device; assert Device.DEFAULT in ['LLVM','CLANG','CUDA','GPU'], Device.DEFAULT"
+          DEBUG=5 PYTHONPATH=${{ github.workspace }} FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+      - name: Run pytest (not cuda)
+        if: matrix.backend!='cuda' && matrix.backend!='ptx' && matrix.backend!='triton'
+        run: python -m pytest -n=auto test/ --durations=20
+      - name: Run ONNX (only LLVM)
+        if: matrix.backend == 'llvm'
+        run: python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      - name: Run pytest (cuda)
+        if: matrix.backend=='cuda'||matrix.backend=='ptx'||matrix.backend=='triton'
+        run: python -m pytest -n=auto test/ -k 'not (half or test_efficientnet_safetensors)' --ignore=test/external --ignore=test/models --durations=20
+
+  #testunicorn:
+  #  name: ARM64 unicorn Test
+  #  runs-on: ubuntu-latest
+  #  timeout-minutes: 20
+  #  steps:
+  #    - name: Checkout Code
+  #      uses: actions/checkout@v3
+  #    - name: Set up Python 3.11
+  #      uses: actions/setup-python@v4
+  #      with:
+  #        python-version: 3.11
+  #    - name: Cache python packages
+  #      uses: actions/cache@v3
+  #      with:
+  #        path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
+  #        key: testing-arm-packages-${{ hashFiles('**/setup.py') }}
+  #    - name: Install cross-assembler
+  #      run: |
+  #        sudo apt update -y
+  #        sudo apt install -y --no-install-recommends gcc-aarch64-linux-gnu
+  #    - name: Install dependencies
+  #      run: pip install -e '.[testing,arm]' --extra-index-url https://download.pytorch.org/whl/cpu
+  #    - name: Test arm
+  #      run: CI=1 ARM64=1 CLANG=1 python -m pytest -n=auto test/ -k 'not (test_nn.py and (test_conv_transpose2d or test_conv2d))' --ignore=test/models --ignore=test/test_speed_v_torch.py --ignore=test/test_net_speed.py --ignore=test/test_specific_conv.py  --ignore=test/unit/test_disk_tensor.py

.gitignore

@@ -1,11 +1,13 @@
 __pycache__
 .venv/
 .vscode
+.DS_Store
 notebooks
 .*.swp
 .*.swo
 *.pyc
 *.so
+*.txt
 build
 /dist
 *.egg-info
@@ -29,8 +31,10 @@ extra/datasets/open-images-v6-mlperf
 extra/datasets/kits/
 extra/datasets/COCO/
 extra/datasets/audio*
+extra/weights
 venv
-examples/net.*[js,json,safetensors]
+examples/**/net.*[js,json]
+examples/**/*.safetensors
 node_modules
 package.json
 package-lock.json
@@ -39,3 +43,10 @@ temp
 .coverage
 coverage.xml
 htmlcov
+outputs_yolov8
+wandb
+model.safetensors
+quickstart.py
+.hypothesis
+weights
+*.lprof

.pre-commit-config.yaml

@@ -1,33 +1,53 @@
 repos:
   - repo: local
     hooks:
-      - id: docs
-        name: docs
-        entry: python3 docs/abstractions.py
+      - id: whitespace
+        name: strip whitespace
+        entry: ./strip_whitespace.sh
         language: system
         always_run: true
         pass_filenames: false
-      - id: flake8
-        name: flake8
-        entry: flake8 --statistics -j4
+      - id: ruff
+        name: ruff
+        entry: ruff . --preview
         language: system
         always_run: true
         pass_filenames: false
       - id: mypy
         name: mypy
-        entry: mypy tinygrad/ extra/helpers.py --check-untyped-defs --explicit-package-bases --warn-unreachable # --warn-return-any
+        entry: mypy tinygrad/
+        language: system
+        always_run: true
+        pass_filenames: false
+      - id: docs
+        name: docs
+        entry: |
+          python3 docs/abstractions.py
+          python3 docs/abstractions2.py
+        language: system
+        always_run: true
+        pass_filenames: false
+      - id: devicetests
+        name: select GPU tests
+        entry: env GPU=1 PYTHONPATH="." pytest test/test_uops.py test/test_custom_function.py test/test_search.py
         language: system
         always_run: true
         pass_filenames: false
       - id: tests
-        name: subset of (CPU) tests
-        entry: env CPU=1 pytest test/unit/ test/test_ops.py test/test_dtype.py
+        name: subset of TORCH tests
+        entry: env PYTHONPATH="." TORCH=1 python3 -m pytest -n=4 test/unit/ test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_custom_function.py test/test_assign.py test/test_symbolic_shapetracker.py
+        language: system
+        always_run: true
+        pass_filenames: false
+      - id: example
+        name: multi device tests
+        entry: python3 test/external/external_test_example.py
         language: system
         always_run: true
         pass_filenames: false
       - id: pylint
         name: pylint
-        entry: pylint tinygrad/
+        entry: env PYTHONPATH="." python3 -m pylint tinygrad/
         language: system
         always_run: true
         pass_filenames: false

.pylintrc

@@ -248,7 +248,7 @@ indent-after-paren=4
 indent-string='  '
 
 # Maximum number of characters on a single line.
-max-line-length=100
+max-line-length=150
 
 # Maximum number of lines in a module
 max-module-lines=1000

CONTRIBUTING.md

@@ -1,33 +0,0 @@
-# Are you ready to write high quality code?
-
-The idea of tinygrad is to build a <5000 line library capable of training a wide variety of ML models at 80%+ max theoretical speed across a huge variety of hardware.
-
-There is almost no boilerplate code anywhere in this library, and you should help keep it that way. If the code you are contributing to core tinygrad, in `tinygrad/`, isn't some of the highest quality code you've written in your life, either put in the effort to make it great, or don't bother. (other directories have a slightly more relaxed standard)
-
-There is a linter, but it's not complete. Spend a little time reading the existing code to get a feel for the style.
-
-I love PRs where I can look at them and just say, yes, this will improve the codebase and click merge. If you have an incomplete PR, feel free to post it as a draft.
-
-As my operating systems professor taught me, code is written to be read by humans. We value readability over performance and line count, but low line count is often a good proxy for readability. However, any PRs that look like code golf will immediately be closed.
-
-There are a few basic ways to contribute:
-
-## Bug-fixes
-
-These are the most straightforward. Discover a bug. Add a test to reproduce it. Write a clean fix. Submit a PR. Confirm CI passes.
-
-## Conceptual Cleanups
-
-This is some of the highest value work in tinygrad. If you realize two 50 line functions are basically the same thing, and you can merge them, amazing! Things that look confusing and are hard to follow are probably poorly written. If you can rewrite the code and be like, oh that's a ton simpler, by all means do so. Make sure you have good test coverage around what you are changing.
-
-## Better Testing
-
-Always welcome! Think about how robust and fast your tests are though. How likely is this test to catch a bug? Tests that run in CI go in `test/`, except for the ones in `test/external/`. We have a few things like fuzzers in there.
-
-## Speed improvements
-
-tinygrad is a JIT compiler, so speed improvements refer to both compile-time and runtime. Speed improvements to the python based compiler are welcome, but please include benchmarks and good tests around the things that you are changing. If you are sacrificing readability for speed, don't bother. Speed improvements to the generated code usually come from conceptual cleanups. Generated code improvements are probably the hardest thing to work on in tinygrad, since they must be done in a very generic way.
-
-## Features
-
-This is a trickier one. If there is a feature in PyTorch and numpy that you have actually seen people use, we probably want it. All new features must include good robust tests, and in general, matching the PyTorch API is good.

README.md

@@ -13,7 +13,6 @@ tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) an
 [![GitHub Repo stars](https://img.shields.io/github/stars/tinygrad/tinygrad)](https://github.com/tinygrad/tinygrad/stargazers)
 [![Unit Tests](https://github.com/tinygrad/tinygrad/actions/workflows/test.yml/badge.svg)](https://github.com/tinygrad/tinygrad/actions/workflows/test.yml)
 [![Discord](https://img.shields.io/discord/1068976834382925865)](https://discord.gg/ZjZadyC7PK)
-[![Lines of code](https://img.shields.io/tokei/lines/github/tinygrad/tinygrad)](https://github.com/tinygrad/tinygrad)
 
 </div>
 
@@ -36,9 +35,9 @@ tinygrad can run [LLaMA](/docs/showcase.md#llama) and [Stable Diffusion](/docs/s
 Try a matmul. See how, despite the style, it is fused into one kernel with the power of laziness.
 
 ```sh
-DEBUG=3 python3 -c "from tinygrad.tensor import Tensor;
+DEBUG=3 python3 -c "from tinygrad import Tensor;
 N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N);
-c = (a.reshape(N, 1, N) * b.permute(1,0).reshape(1, N, N)).sum(axis=2);
+c = (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2);
 print((c.numpy() - (a.numpy() @ b.numpy())).mean())"
 ```
 
@@ -49,32 +48,30 @@ And we can change `DEBUG` to `4` to see the generated code.
 As it turns out, 90% of what you need for neural networks are a decent autograd/tensor library.
 Throw in an optimizer, a data loader, and some compute, and you have all you need.
 
-#### Neural network example (from test/models/test_mnist.py)
-
 ```py
-from tinygrad.tensor import Tensor
-import tinygrad.nn.optim as optim
+from tinygrad import Tensor, nn
 
-class TinyBobNet:
+class LinearNet:
   def __init__(self):
-    self.l1 = Tensor.uniform(784, 128)
-    self.l2 = Tensor.uniform(128, 10)
+    self.l1 = Tensor.kaiming_uniform(784, 128)
+    self.l2 = Tensor.kaiming_uniform(128, 10)
+  def __call__(self, x:Tensor) -> Tensor:
+    return x.flatten(1).dot(self.l1).relu().dot(self.l2)
 
-  def forward(self, x):
-    return x.dot(self.l1).relu().dot(self.l2).log_softmax()
+model = LinearNet()
+optim = nn.optim.Adam([model.l1, model.l2], lr=0.001)
 
-model = TinyBobNet()
-optim = optim.SGD([model.l1, model.l2], lr=0.001)
+x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7])  # replace with real mnist dataloader
 
-# ... complete data loader here
-
-out = model.forward(x)
-loss = out.mul(y).mean()
-optim.zero_grad()
-loss.backward()
-optim.step()
+for i in range(10):
+  optim.zero_grad()
+  loss = model(x).sparse_categorical_crossentropy(y).backward()
+  optim.step()
+  print(i, loss.item())
 ```
 
+See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full version that gets 98% in ~5 seconds
+
 ## Accelerators
 
 tinygrad already supports numerous accelerators, including:
@@ -85,12 +82,10 @@ tinygrad already supports numerous accelerators, including:
 - [x] [LLVM](tinygrad/runtime/ops_llvm.py)
 - [x] [METAL](tinygrad/runtime/ops_metal.py)
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
-- [x] [Triton](extra/accel/triton/ops_triton.py)
 - [x] [PyTorch](tinygrad/runtime/ops_torch.py)
 - [x] [HIP](tinygrad/runtime/ops_hip.py)
-- [x] [WebGPU](tinygrad/runtime/ops_webgpu.py)
 
-And it is easy to add more! Your accelerator of choice only needs to support a total of 26 (optionally 27) low level ops.
+And it is easy to add more! Your accelerator of choice only needs to support a total of ~25 low level ops.
 More information can be found in the [documentation for adding new accelerators](/docs/adding_new_accelerators.md).
 
 ## Installation
@@ -104,7 +99,12 @@ git clone https://github.com/tinygrad/tinygrad.git
 cd tinygrad
 python3 -m pip install -e .
 ```
-Don't forget the `.` at the end!
+
+### Direct (master)
+
+```sh
+python3 -m pip install git+https://github.com/tinygrad/tinygrad.git
+```
 
 ## Documentation
 
@@ -113,7 +113,7 @@ Documentation along with a quick start guide can be found in the [docs/](/docs)
 ### Quick example comparing to PyTorch
 
 ```py
-from tinygrad.tensor import Tensor
+from tinygrad import Tensor
 
 x = Tensor.eye(3, requires_grad=True)
 y = Tensor([[2.0,0,-2.0]], requires_grad=True)
@@ -139,24 +139,33 @@ print(y.grad.numpy())  # dz/dy
 
 ## Contributing
 
-There has been a lot of interest in tinygrad lately. Here are some basic guidelines for contributing:
+There has been a lot of interest in tinygrad lately. Following these guidelines will help your PR get accepted.
 
-- Bug fixes are the best and always welcome! Like [this one](https://github.com/tinygrad/tinygrad/pull/421/files).
-- If you don't understand the code you are changing, don't change it!
-- All code golf PRs will be closed, but [conceptual cleanups](https://github.com/tinygrad/tinygrad/pull/372/files) are great.
-- Features are welcome. Though if you are adding a feature, you need to include tests.
-- Improving test coverage is great, with reliable non-brittle tests.
+We'll start with what will get your PR closed with a pointer to this section:
 
-Additional guidelines can be found in [CONTRIBUTING.md](/CONTRIBUTING.md).
+- No code golf! While low line count is a guiding light of this project, anything that remotely looks like code golf will be closed. The true goal is reducing complexity and increasing readability, and deleting `\n`s does nothing to help with that.
+- All docs and whitespace changes will be closed unless you are a well-known contributor. The people writing the docs should be those who know the codebase the absolute best. People who have not demonstrated that shouldn't be messing with docs. Whitespace changes are both useless *and* carry a risk of introducing bugs.
+- Anything you claim is a "speedup" must be benchmarked. In general, the goal is simplicity, so even if your PR makes things marginally faster, you have to consider the tradeoff with maintainablity and readablity.
+- In general, the code outside the core `tinygrad/` folder is not well tested, so unless the current code there is broken, you shouldn't be changing it.
+
+Now, what we want:
+
+- Bug fixes (with a regression test) are great! This library isn't 1.0 yet, so if you stumble upon a bug, fix it, write a test, and submit a PR, this is valuable work.
+- Solving bounties! tinygrad [offers cash bounties](https://docs.google.com/spreadsheets/d/1WKHbT-7KOgjEawq5h5Ic1qUWzpfAzuD_J06N1JwOCGs/edit?usp=sharing) for certain improvements to the library. All new code should be high quality and well tested.
+- Features. However, if you are adding a feature, consider the line tradeoff. If it's 3 lines, there's less of a bar of usefulness it has to meet over something that's 30 or 300 lines. All features must have regression tests. In general with no other constraints, your feature's API should match torch or numpy.
+- Refactors that are clear wins. In general, if your refactor isn't a clear win it will be closed. But some refactors are amazing! Think about readability in a deep core sense. A whitespace change or moving a few functions around is useless, but if you realize that two 100 line functions can actually use the same 110 line function with arguments while also improving readability, this is a big win.
+- Tests/fuzzers. If you can add tests that are non brittle, they are welcome. We have some fuzzers in here too, and there's a plethora of bugs that can be found with them and by improving them. Finding bugs, even writing broken tests (that should pass) with `@unittest.expectedFailure` is great. This is how we make progress.
+- Dead code removal from core `tinygrad/` folder. We don't care about the code in extra, but removing dead code from the core library is great. Less for new people to read and be confused by.
 
 ### Running tests
 
+You should install the pre-commit hooks with `pre-commit install`. This will run the linter, mypy, and a subset of the tests on every commit.
+
 For more examples on how to run the full test suite please refer to the [CI workflow](.github/workflows/test.yml).
 
-Some examples:
+Some examples of running tests locally:
 ```sh
-python3 -m pip install -e '.[testing]'
-python3 -m pytest
-python3 -m pytest -v -k TestTrain
-python3 ./test/models/test_train.py TestTrain.test_efficientnet
+python3 -m pip install -e '.[testing]'  # install extra deps for testing
+python3 test/test_ops.py                # just the ops tests
+python3 -m pytest test/                 # whole test suite
 ```

cache/.gitignore

@@ -1 +0,0 @@
-*

compile.sh

@@ -1,8 +0,0 @@
-#!/bin/bash
-# note: if we compile tinygrad/nn/__init__.py __dict__ no longer works, and optimizers will silently fail
-mypyc --check-untyped-defs --explicit-package-bases --warn-unreachable tinygrad/shape/shapetracker.py tinygrad/shape/symbolic.py \
-  tinygrad/helpers.py tinygrad/mlops.py tinygrad/tensor.py tinygrad/graph.py \
-  #tinygrad/codegen/gpu.py tinygrad/runtime/ops_metal.py
-  #tinygrad/codegen/ast.py
-  #tinygrad/nn/__init__.py
-  #tinygrad/ops.py tinygrad/runtime/ops_metal.py tinygrad/runtime/ops_gpu.py tinygrad/runtime/ops_cpu.py tinygrad/lazy.py

disassemblers/adreno/__init__.py

@@ -4,15 +4,17 @@ import pathlib
 from hexdump import hexdump
 
 fxn = None
-def disasm(buf):
+def disasm_raw(buf):
   global fxn
   if fxn is None:
     shared = pathlib.Path(__file__).parent / "disasm.so"
-    if not os.path.isfile(shared):
+    if not shared.is_file():
       os.system(f'cd {pathlib.Path(__file__).parent} && gcc -shared disasm-a3xx.c -o disasm.so')
     fxn = ctypes.CDLL(shared.as_posix())['disasm']
-  #hexdump(buf)
+  fxn(buf, len(buf))
+
+def disasm(buf):
   END = b"\x00\x00\x00\x00\x00\x00\x00\x03"
   buf = buf[0x510:]  # this right?
   buf = buf.split(END)[0] + END
-  fxn(buf, len(buf))
+  disasm_raw(buf)

docs/DESIGNv2.md

@@ -0,0 +1,17 @@
+tinygrad is a bit bloated now, and there's several places where concerns should be seperated and they aren't.
+
+tensor.py and mlops.py are great code. The interface going backward here is:
+
+LazyBuffer.const (this creates a matching size buffer)
+LazyBuffer.contiguous (tbis is not exactly elementwise)
+LazyBuffer.e (elementwise)
+LazyBuffer.r (reduce)
+reshape/permute/expand/stride/shrink/pad (movement)
+
+The lazy.py reordering engine has a lot of junk to deal with movementops that should be removed.
+
+view.py is mostly great code, except it shouldn't have the rendering logic, and the int type should be parameterized to not import from symbolic.
+
+LazyOp shouldn't have LazyBuffers as sources, just LazyOp LoadOps with a tuple of Views. Then the LazyOp uniquely determines the kernel and we don't have to do any replacement.
+
+ShapeTracker probably shouldn't exist and just be a part of LazyBuffer. Most of the stuff in ShapeTracker should move to symbolic_view, which combines view and symbolic.

docs/README.md

@@ -29,9 +29,3 @@ This is currently a big work in progress.
 ### Community
 
 [![tinygrad discord](https://discordapp.com/api/guilds/1068976834382925865/widget.png?style=banner2)](https://discord.gg/ZjZadyC7PK)
-
-## Contributing
-
-The documentation mainly follows the core contributing guidelines in the [README.md](/README.md#contributing).
-
-Additionally, we always welcome documentation contributions, especially for features that are currently under documented.

docs/abstractions.py

@@ -3,17 +3,17 @@ Welcome to the tinygrad documentation
 =================
 
 this file will take you on a whirlwind journey from a Tensor all the way down
-tinygrad has been aggressively refactored in the 2.5 years it's been worked on.
+tinygrad has been aggressively refactored in the 3 years it's been worked on.
 what you see here is a refined library (with more refining to go still!)
 
-the whole tinygrad is ~2300 lines, so while it's readable in an evening or two,
+the whole tinygrad is < 5000 lines, so while it's readable in an evening or two,
 this documentation will help with entry points and understanding the abstraction stack
 """
 
 # %%
 # == Boilerplate imports for typing ==
 from __future__ import annotations
-from typing import Optional, Tuple, Union, Any, Dict, Callable, Type, List, ClassVar
+from typing import Optional, Tuple, Union, Any, Dict, Callable, Type, List
 from enum import Enum, auto
 from abc import ABC
 
@@ -22,7 +22,7 @@ from abc import ABC
 # let's trace an addition down through the layers of abstraction.
 
 # we will be using the clang backend
-from tinygrad.ops import Device
+from tinygrad import Device
 Device.DEFAULT = "CLANG"
 
 # first, 2+3 as a Tensor, the highest level
@@ -67,7 +67,7 @@ class Function:
 
 # %%
 # == LazyBuffer (in tinygrad/lazy.py, code 5/10) ==
-from tinygrad.helpers import DType
+from tinygrad.dtype import DType
 
 # this is where the properties live that you thought were a part of Tensor
 # LazyBuffer is like a Tensor without derivatives, at the mlop layer
@@ -83,34 +83,34 @@ class LazyBuffer:
   # we'll come back to this later
   st: ShapeTracker
 
-  # if the LazyBuffer is realized, it has a RawBuffer
-  # we will come back to RawBuffers later
-  realized: Optional[RawBuffer]
+  # if the LazyBuffer is realized, it has a Buffer
+  # we will come back to Buffer later
+  realized: Optional[Buffer]
 
   # if the lazybuffer is unrealized, it has a LazyOp
   # this LazyOp describes the computation needed to realize this LazyBuffer
   op: Optional[LazyOp]
 
-# LazyOp (in tinygrad/ops.py, code 4/10)
+# LazyOp (in tinygrad/ops.py, code 5/10)
 # in a tree they form an Abstract Syntax Tree for a single GPU kernel
 class LazyOp:
   op: Op                                       # the type of the compute
-  src: Tuple[Union[LazyOp, LazyBuffer], ...]   # the sources
-  arg: Optional[Any] = None                    # and an optional static argument
+  src: Tuple[LazyOp, ...]                      # the sources
+  arg: Any = None                              # and an optional static argument
 
-# there's currently 28 Ops you have to implement for an accelerator.
-class UnaryOps(Enum):    NOOP = auto(); EXP2 = auto(); LOG2 = auto(); CAST = auto(); SIN = auto();   SQRT = auto()
+# there's currently 26 Ops you have to implement for an accelerator.
+class UnaryOps(Enum):    EXP2 = auto(); LOG2 = auto(); CAST = auto(); SIN = auto();   SQRT = auto()
 class BinaryOps(Enum):   ADD = auto();  SUB = auto();  MUL = auto();  DIV = auto();  CMPLT = auto(); MAX = auto()
 class ReduceOps(Enum):   SUM = auto();  MAX = auto()
 class MovementOps(Enum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); STRIDE = auto()
 class TernaryOps(Enum):  MULACC = auto(); WHERE = auto()
-class LoadOps(Enum):     EMPTY = auto(); RAND = auto(); CONST = auto(); FROM = auto(); CONTIGUOUS = auto(); CUSTOM = auto()
-# NOTE: if you have a CompiledBuffer(DeviceBuffer)
+class LoadOps(Enum):     EMPTY = auto(); CONST = auto(); COPY = auto(); CONTIGUOUS = auto(); CUSTOM = auto()
+# NOTE: if you have a Compiled device
 #       you do not need to implement the MovementOps
-#       as they are handled by the ShapeTracker(in tinygrad/shape/shapetracker.py, code 7/10)
+#       as they are handled by the ShapeTracker (in tinygrad/shape/shapetracker.py, code 7/10)
 Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, TernaryOps, LoadOps]
 
-# most of tinygrad/lazy.py is concerned with fusing Ops into LazyOps ASTs that map to GPUKernels
+# most of tinygrad/lazy.py is concerned with fusing Ops into LazyOps ASTs that map to kernels
 # it's beyond the scope of this tutorial, but you can read the file if interested
 
 # %%
@@ -118,44 +118,43 @@ Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, TernaryOps, LoadOps]
 
 from tinygrad.tensor import Tensor
 from tinygrad.ops import LazyOp, BinaryOps, LoadOps
+from tinygrad.lazy import LazyBuffer
+from tinygrad.device import Buffer
 
 # the 2+3 from before
 result = Tensor([2]) + Tensor([3])
 print(type(result.lazydata), result.lazydata)  # let's look at the lazydata of result
 
-# you'll see it has a LazyOp
 # the op type is BinaryOps.ADD
 # and it has two sources, the 2 and the 3
-lazyop: LazyOp = result.lazydata.op
+lazyop: LazyBuffer = result.lazydata
 assert lazyop.op == BinaryOps.ADD
-assert len(lazyop.src) == 2
+assert len(lazyop.srcs) == 2
 
 # the first source is the 2, it comes from the CPU
 # the source is a LazyBuffer that is a "CPU" Tensor
 # again, a LazyOp AST is like a GPU kernel. you have to copy the data on the device first
-assert lazyop.src[0].op.op == LoadOps.FROM
-assert lazyop.src[0].op.src[0].device == "CPU"
-assert lazyop.src[0].op.src[0].realized._buf[0] == 2, "the src of the FROM LazyOP is a LazyBuffer on the CPU holding [2.]"
-assert result.lazydata.realized is None, "the LazyBuffer is not realized yet"
+assert lazyop.srcs[0].op == LoadOps.COPY
+assert lazyop.srcs[0].srcs[0].device == "CPU"
+assert lazyop.srcs[0].srcs[0].realized._buf[0] == 2, "the src of the COPY LazyOP is a LazyBuffer on the CPU holding [2]"
+assert result.lazydata.base.realized is None, "the LazyBuffer is not realized yet"
 
 # now we realize the LazyBuffer
-result.lazydata.realize()
-assert result.lazydata.realized is not None, "the LazyBuffer is realized!"
-# this brings us nicely to DeviceBuffer, of which the realized ClangBuffer is a subclass
-assert 'RawMallocBuffer' in str(type(result.lazydata.realized))
-# getting ahead of ourselves, but we can copy the DeviceBuffer toCPU
-assert result.lazydata.realized.toCPU()[0] == 5, "when put in numpy with toCPU, it's 5"
+result.realize()
+assert result.lazydata.base.realized is not None, "the LazyBuffer is realized!"
+# this brings us nicely to Buffer
+assert isinstance(result.lazydata.base.realized, Buffer)
+assert result.lazydata.base.realized.device == "CLANG"
+# getting ahead of ourselves, but we can move the Buffer to CPU
+assert result.lazydata.base.realized.toCPU()[0] == 5, "when put in numpy with toCPU, it's 5"
 
 # %%
-# == Union[Interpreted, Compiled] (in tinygrad/ops.py, code 5/10) ==
+# == Union[Interpreted, Compiled] (in tinygrad/device.py, code 6/10) ==
 
 # Now you have a choice, you can either write a "Interpreted" backend or "Compiled" backend
 
 # Interpreted backends are very simple (example: CPU and TORCH)
 class Interpreted:
-  # they have a backing RawBuffer
-  buffer: Type[RawBuffer]
-
   # and they have a lookup table to functions for the Ops
   fxn_for_op: Dict[Op, Callable] = {
     UnaryOps.EXP2: lambda x: np.exp2(x),
@@ -163,9 +162,6 @@ class Interpreted:
 
 # Compiled backends take a little more (example: GPU and LLVM)
 class Compiled:
-  # they also have a backing RawBuffer
-  buffer: Type[RawBuffer]
-
   # a code generator, which compiles the AST
   codegen: Type[Linearizer]
 
@@ -178,62 +174,53 @@ class Runtime(ABC):
   # the constructor compiles the code
   def __init__(self, name:str, prg:str): pass
   # call runs the code on the bufs. NOTE: the output is always bufs[0], but this is just a convention
-  def __call__(self, global_size:Optional[List[int]], local_size:Optional[List[int]], *bufs:List[RawBuffer]): pass
+  def __call__(self, *bufs:List[Buffer], global_size:Optional[List[int]], local_size:Optional[List[int]]): pass
 
 # %%
-# == RawBuffer (in tinygrad/runtime/lib.py, code 5/10) ==
+# == Buffer (in tinygrad/device.py, code 6/10) ==
 import numpy as np
 
-# RawBuffer is where the data is actually held. it's pretty close to just memory
-class RawBuffer(ABC):
+# Buffer is where the data is actually held. it's pretty close to just memory
+class Buffer(ABC):
   # create an empty rawbuffer that holds `size` elements of type `dtype`
-  # `buf` is an opaque container class
-  def __init__(self, size:int, dtype:DType, buf:Any): raise NotImplementedError("must be implemented")
+  # `opaque` is an opaque container class
+  def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None): pass
 
-  # fromCPU is classmethod that creates a RawBuffer, it's a classmethod since some runtimes are 0 copy
-  @classmethod
-  def fromCPU(cls:RawBuffer, x:np.ndarray) -> RawBuffer: raise NotImplementedError("must be implemented")
-
-  # toCPU converts the RawBuffer to a numpy array with shape (size,). many backends are 0 copy here
-  def toCPU(self) -> np.ndarray: raise NotImplementedError("must be implemented")
-
-# RawNumpyBuffer is a RawBuffer example for numpy. It's very simple
-class RawNumpyBuffer(RawBuffer):
-  # NOTE: the "np.ndarray" is stored in the opaque container
-  def __init__(self, buf:np.ndarray):
-    super().__init__(buf.size, dtypes.from_np(buf.dtype), buf)
-  @classmethod
-  def fromCPU(cls, x): return cls(x)
-  def toCPU(self): return self._buf
+  # toCPU converts the RawBuffer to a numpy array with shape (size,)
+  def toCPU(self) -> np.ndarray: pass
 
 # %%
 # == Example: 2+3 in raw clang ==
 
-# RawMallocBuffer is the simplest concrete version of RawBuffer (in tinygrad/ops.py)
+# MallocAllocator is the simplest concrete version of Allocator (in tinygrad/device.py)
 # it's used for the CLANG and LLVM backends
 # it's just malloc(size * dtype.itemsize)
-from tinygrad.runtime.lib import RawMallocBuffer
+from tinygrad.device import MallocAllocator
 
 # ClangProgram is the simplest runtime (in tinygrad/runtime/ops_clang.py, code 7/10)
 # __init__ calls clang, and __call__ calls the function in the *.so outputted by clang
 # in CLANG, global_size and local_size are ignored
-from tinygrad.runtime.ops_clang import ClangProgram
+from tinygrad.runtime.ops_clang import ClangProgram, compile_clang
 
 # a concrete example looks like this, this adds two size 1 RawBuffer
 # first we create two numpy buffers containing 2 and 3
 # then we copy the numpy in to RawMallocBuffers
 # last, we create an empty output buffer
-from tinygrad.helpers import dtypes
+input_a, input_b = MallocAllocator.alloc(4), MallocAllocator.alloc(4)
+output = MallocAllocator.alloc(4)
+
+# now we copy in the values
 numpy_a, numpy_b = np.array([2], dtype=np.float32), np.array([3], dtype=np.float32)
-input_a, input_b = RawMallocBuffer.fromCPU(numpy_a), RawMallocBuffer.fromCPU(numpy_b)
-output = RawMallocBuffer(1, dtypes.float32)
+MallocAllocator.copyin(input_a, numpy_a.data.cast("B"))
+MallocAllocator.copyin(input_b, numpy_b.data.cast("B"))
 
 # compile the program, run it, and 2+3 does indeed equal 5
-program = ClangProgram("add", f"void add(float *a, float *b, float *c) {{ *a = *b + *c; }}")
-program(None, None, output, input_a, input_b)  # NOTE: the None are for global_size and local_size
-print(output.toCPU())
-assert output.toCPU()[0] == 5, "it's still 5"
-np.testing.assert_allclose(output.toCPU(), numpy_a+numpy_b)
+program = ClangProgram("add", compile_clang(f"void add(float *a, float *b, float *c) {{ *a = *b + *c; }}"))
+program(output, input_a, input_b)
+numpy_out = np.empty(1, dtype=np.float32)
+MallocAllocator.copyout(numpy_out.data.cast("B"), output)
+assert numpy_out[0] == 5, "it's still 5"
+np.testing.assert_allclose(numpy_out, numpy_a+numpy_b)
 
 # %%
 # == Linearizer (in tinygrad/codegen/linearizer.py, code 4/10) ==
@@ -245,33 +232,28 @@ np.testing.assert_allclose(output.toCPU(), numpy_a+numpy_b)
 
 class UOps(Enum): LOOP = auto(); DEFINE_LOCAL = auto(); LOAD = auto(); ALU = auto(); CONST = auto(); ENDLOOP = auto(); STORE = auto();
 
-class Token:
-  name: str
-
 class UOp:
   uop: UOps
-  out: Optional[Token]
-  vin: List[Token]
+  dtype: Optional[DType]
+  vin: Tuple[UOp, ...]
   arg: Any
 
 class Linearizer:
   # create the kernel with the AST
   # NOTE: the AST contains the CompiledBuffers themselves as the root nodes. this will change
   def __init__(self, ast:LazyOp): pass
-  def process(self): pass
   def linearize(self): pass
 
   # when linearize is run, it fills in this list
   uops: List[UOp]
 
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import prod
-result = Tensor(2).realize() + Tensor(3).realize()
-result.lazydata.realized = Device[Device.DEFAULT].buffer(prod(result.shape), result.dtype)
+result = Tensor(2.0).realize() + Tensor(3.0).realize()
 
 # use the real Linearizer to linearize 2+3
-from tinygrad.codegen.linearizer import Linearizer, LinearizerOptions
-linearizer = Linearizer(result.lazydata.op, result.lazydata, LinearizerOptions())
+from tinygrad.codegen.linearizer import Linearizer
+sched = result.lazydata.schedule()
+linearizer = Linearizer(sched[-1].ast)
 linearizer.linearize()
 
 # print the uops
@@ -279,14 +261,12 @@ for uop in linearizer.uops: print(uop)
 
 # output:
 """
-UOps.DEFINE_GLOBAL  :                           []                               ('data0', dtypes.float)
-UOps.LOOP           :                           []                               ([], 'global')
-UOps.LOOP           :                           []                               ([], 'local')
-UOps.LOAD           : <acc1_0>                  []                               ConstOp(value=2.0, valid=<1>, invalid_value=0.0)
-UOps.LOAD           : <acc2_0>                  []                               ConstOp(value=3.0, valid=<1>, invalid_value=0.0)
-UOps.ALU            : <alu0>                    [<acc1_0>, <acc2_0>]             BinaryOps.ADD
-UOps.STORE          :                           [<alu0>]                         MemOp(name='data0', idx=<0>, local=False, memory_dtype=dtypes.float, valid=<1>, invalid_value=0.0)
-UOps.ENDLOOP        :                           []                               ([], 'global+local')
+   0 UOps.DEFINE_GLOBAL  : ptr.dtypes.float          []                               data0
+   1 UOps.CONST          : dtypes.float              []                               2.0
+   2 UOps.CONST          : dtypes.float              []                               3.0
+   3 UOps.ALU            : dtypes.float              [1, 2]                           BinaryOps.ADD
+   4 UOps.CONST          : dtypes.int                []                               0
+   5 UOps.STORE          :                           [0, 4, 3]                        None
 """
 
 # %%
@@ -295,30 +275,25 @@ UOps.ENDLOOP        :                           []
 # here, we have an example where we fetch the generated code from the JIT
 
 from tinygrad.tensor import Tensor
-result = Tensor(2) + Tensor(3)
+result = Tensor(2.0) + Tensor(3.0)
 
 # we have a global cache used by the JIT
 # from there, we can see the generated clang code
-from tinygrad.helpers import GlobalCounters
-GlobalCounters.cache = []    # enables the cache
+from tinygrad.jit import CacheCollector
+CacheCollector.start()       # enables the cache
 result.realize()             # create the program and runs it
-cache_saved = GlobalCounters.cache
-GlobalCounters.cache = None  # disable the cache
+cache_saved = CacheCollector.finish()  # disable the cache
 
 # there's one ASTRunner in the cache
 assert len(cache_saved) == 1
-prg, bufs, _ = cache_saved[0]
 
 # print the C Program :)
-print(prg.prg)
+print(cache_saved[0].prg.prg)
 
-# after some formatting (the compiler doesn't care)
 # NOTE: the 2 and 3 are constant folded
 """
-void E_1(float* data0) {
-  for (int idx0 = 0; idx0 < 1; idx0++) {
-    data0[0] = (2.0f) + (3.0f);
-  }
+void E_n2(float* restrict data0) {
+  data0[0] = (2.0f+3.0f);
 }
 """
 
@@ -331,22 +306,22 @@ void E_1(float* data0) {
 from tinygrad.shape.shapetracker import ShapeTracker
 
 # create a virtual (10, 10) Tensor. this is just a shape, there's no actual tensor
-a = ShapeTracker((10, 10))
+a = ShapeTracker.from_shape((10, 10))
 
 # you'll see it has one view. the (10, 1 are the strides)
 print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (10, 1), 0)])
 
 # we can permute it, and the strides change
-a.permute((1,0))
+a = a.permute((1,0))
 print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (1, 10), 0)])
 
 # we can then reshape it, and the strides change again
 # note how the permute stays applied
-a.reshape((5,2,5,2))
+a = a.reshape((5,2,5,2))
 print(a) # ShapeTracker(shape=(5, 2, 5, 2), views=[View((5, 2, 5, 2), (2, 1, 20, 10), 0)])
 
 # now, if we were to reshape it to a (100,) shape tensor, we have to create a second view
-a.reshape((100,))
+a = a.reshape((100,))
 print(a) # ShapeTracker(shape=(100,), views=[
          #   View((5, 2, 5, 2), (2, 1, 20, 10), 0),
          #   View((100,), (1,), 0)])
@@ -357,7 +332,7 @@ idx, _ = a.expr_idxs()
 print(idx.render())  # (((idx0%10)*10)+(idx0//10))
 
 # of course, if we reshape it back, the indexes get simple again
-a.reshape((10,10))
+a = a.reshape((10,10))
 idx, _ = a.expr_idxs()
 print(idx.render())  # ((idx1*10)+idx0)
 
@@ -367,11 +342,11 @@ print(a) # ShapeTracker(shape=(10, 10), views=[
          #   View((10, 10), (10, 1), 0)])
 
 # ...until we simplify it!
-a.simplify()
+a = a.simplify()
 print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (1, 10), 0)])
 
 # and now we permute it back
-a.permute((1,0))
+a = a.permute((1,0))
 print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (10, 1), 0)])
 
 # and it's even contiguous

docs/abstractions2.py

@@ -0,0 +1,115 @@
+# tinygrad is a tensor library, and as a tensor library it has multiple parts
+# 1. a "runtime". this allows buffer management, compilation, and running programs
+# 2. a "Device" that uses the runtime but specifies compute in an abstract way for all
+# 3. a "LazyBuffer" that fuses the compute into kernels, using memory only when needed
+# 4. a "Tensor" that provides an easy to use frontend with autograd ".backward()"
+
+
+print("******** first, the runtime ***********")
+
+from tinygrad.runtime.ops_clang import ClangProgram, compile_clang, MallocAllocator
+
+# allocate some buffers
+out = MallocAllocator.alloc(4)
+a = MallocAllocator.alloc(4)
+b = MallocAllocator.alloc(4)
+
+# load in some values (little endian)
+MallocAllocator.copyin(a, bytearray([2,0,0,0]))
+MallocAllocator.copyin(b, bytearray([3,0,0,0]))
+
+# compile a program to a binary
+lib = compile_clang("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")
+
+# create a runtime for the program (ctypes.CDLL)
+fxn = ClangProgram("add", lib)
+
+# run the program
+fxn(out, a, b)
+
+# check the data out
+print(val := MallocAllocator.as_buffer(out).cast("I").tolist()[0])
+assert val == 5
+
+
+print("******** second, the Device ***********")
+
+DEVICE = "CLANG"   # NOTE: you can change this!
+
+import struct
+from tinygrad.dtype import dtypes
+from tinygrad.device import Buffer, Device
+from tinygrad.ops import LazyOp, BufferOps, MemBuffer, BinaryOps
+from tinygrad.shape.shapetracker import ShapeTracker
+
+# allocate some buffers + load in values
+out = Buffer(DEVICE, 1, dtypes.int32)
+a = Buffer(DEVICE, 1, dtypes.int32).copyin(memoryview(bytearray(struct.pack("I", 2))))
+b = Buffer(DEVICE, 1, dtypes.int32).copyin(memoryview(bytearray(struct.pack("I", 3))))
+# NOTE: a._buf is the same as the return from MallocAllocator.alloc
+
+# describe the computation
+ld_1 = LazyOp(BufferOps.LOAD, (), MemBuffer(1, dtypes.int32, ShapeTracker.from_shape((1,))))
+ld_2 = LazyOp(BufferOps.LOAD, (), MemBuffer(2, dtypes.int32, ShapeTracker.from_shape((1,))))
+alu = LazyOp(BinaryOps.ADD, (ld_1, ld_2))
+st_0 = LazyOp(BufferOps.STORE, (alu,), MemBuffer(0, dtypes.int32, ShapeTracker.from_shape((1,))))
+
+# convert the computation to a "linearized" format (print the format)
+lin = Device[DEVICE].get_linearizer(st_0).linearize()
+for u in lin.uops: print(u)
+
+# compile a program (and print the source)
+fxn = Device[DEVICE].to_program(lin)
+print(fxn.prg)
+# NOTE: fxn.clprg is the ClangProgram
+
+# run the program
+fxn.exec([out, a, b])
+
+# check the data out
+print(val:=out.toCPU().item())
+assert val == 5
+
+
+print("******** third, the LazyBuffer ***********")
+
+from tinygrad.lazy import LazyBuffer
+from tinygrad.realize import run_schedule
+
+# allocate some values + load in values
+# TODO: remove numpy here
+import numpy as np
+a = LazyBuffer.fromCPU(np.array([2], np.int32)).copy_to_device(DEVICE)
+b = LazyBuffer.fromCPU(np.array([3], np.int32)).copy_to_device(DEVICE)
+
+# describe the computation
+out = a.e(BinaryOps.ADD, b)
+
+# schedule the computation as a list of kernels
+sched = out.schedule()
+for si in sched: print(si.ast.op)  # NOTE: the first two convert it to CLANG
+
+# DEBUGGING: print the compute ast as a tree
+from tinygrad.graph import print_tree
+print_tree(sched[-1].ast)
+# NOTE: sched[-1].ast is the same as st_0 above
+
+# run that schedule
+run_schedule(sched)
+
+# check the data out
+print(val:=out.realized.toCPU().item())
+assert val == 5
+
+
+print("******** fourth, the Tensor ***********")
+
+from tinygrad import Tensor
+
+a = Tensor([2], dtype=dtypes.int32, device=DEVICE)
+b = Tensor([3], dtype=dtypes.int32, device=DEVICE)
+out = a + b
+
+# check the data out
+print(val:=out.item())
+assert val == 5

docs/adding_new_accelerators.md

@@ -11,7 +11,7 @@ unary_op  (NOOP, EXP2, LOG2, CAST, SIN, SQRT)                # A -> A
 reduce_op (SUM, MAX)                                         # A -> B (smaller size, B has 1 in shape)
 binary_op (ADD, SUB, MUL, DIV, CMPEQ, MAX)                   # A + A -> A (all the same size)
 movement_op (EXPAND, RESHAPE, PERMUTE, PAD, SHRINK, STRIDE)  # A -> B (different size)
-load_op   (EMPTY, RAND, CONST, FROM, CONTIGUOUS, CUSTOM)     # -> A   (initialize data on device)
+load_op   (EMPTY, CONST, FROM, CONTIGUOUS, CUSTOM)           # -> A   (initialize data on device)
 ternary_op (WHERE)                                           # A, A, A -> A
 ternary_op [[optional]] (MULACC)                             # A * A -> B
 ```

docs/env_vars.md

@@ -41,16 +41,13 @@ TORCH               | [1]        | enable PyTorch backend
 CLANG               | [1]        | enable Clang backend
 LLVM                | [1]        | enable LLVM backend
 LLVMOPT             | [1]        | enable slightly more expensive LLVM optimizations
-LAZY                | [1]        | enable lazy operations (this is the default)
 OPT                 | [1-3]      | optimization level
-KOPT                | [1-2]      | kernel optimization, 1 turns it on, 2 caches the found optimizations
+BEAM                | [#]        | number of beams in kernel beam search
 GRAPH               | [1]        | create a graph of all operations (requires graphviz)
+GRAPHUOPS           | [1]        | create a graph of uops (requires graphviz and saves at /tmp/uops.{svg,dot})
 GRAPHPATH           | [/path/to] | where to put the generated graph
-PRUNEGRAPH          | [1]        | prune MovementOps and LoadOps from the graph
-PRINT_PRG           | [1]        | print program code
 IMAGE               | [1]        | enable 2d specific optimizations
 FLOAT16             | [1]        | use float16 for images instead of float32
-ENABLE_METHOD_CACHE | [1]        | enable method cache (this is the default)
 DISALLOW_ASSIGN     | [1]        | disallow assignment of tensors
 CL_EXCLUDE          | [name0,name1] | comma-separated list of device names to exclude when using OpenCL GPU backend (like `CL_EXCLUDE=gfx1036`)
 CL_PLATFORM         | [# >= 0]   | index of the OpenCL [platform](https://documen.tician.de/pyopencl/runtime_platform.html#pyopencl.Platform) to run on. Defaults to 0.
@@ -177,7 +174,6 @@ TORCHCUDA | [1] | enable the torch cuda backend
 
 Variable | Possible Value(s) | Description
 ---|---|---
-KOPT   | [1] | enable kernel optimization
 KCACHE | [1] | enable kernel cache
 
 ### test/external/external_test_opt.py

docs/quickstart.md

@@ -55,7 +55,7 @@ There are even more of these factory methods, you can find them in the [tensor.p
 All the tensors creation methods can take a `dtype` argument to specify the data type of the tensor.
 
 ```python
-from tinygrad.helpers import dtypes
+from tinygrad.dtype import dtypes
 
 t3 = Tensor([1, 2, 3, 4, 5], dtype=dtypes.int32)
 ```
@@ -87,7 +87,6 @@ These classes do not need to inherit from any base class, in fact if they don't
 An example of this would be the `nn.Linear` class which represents a linear layer in a neural network.
 
 ```python
-# from tinygrad.nn import Linear
 class Linear:
   def __init__(self, in_features, out_features, bias=True, initialization: str='kaiming_uniform'):
     self.weight = getattr(Tensor, initialization)(out_features, in_features)
@@ -104,8 +103,6 @@ Our classifier will be a simple 2 layer neural network with a Leaky ReLU activat
 It will use a hidden layer size of 128 and an output layer size of 10 (one for each digit) with no bias on either Linear layer.
 
 ```python
-from tinygrad.nn import Linear
-
 class TinyNet:
   def __init__(self):
     self.l1 = Linear(784, 128, bias=False)
@@ -131,16 +128,9 @@ Training neural networks in tinygrad is super simple.
 All we need to do is define our neural network, define our loss function, and then call `.backward()` on the loss function to compute the gradients.
 They can then be used to update the parameters of our neural network using one of the many optimizers in [optim.py](/tinygrad/nn/optim.py).
 
-First we need to set the training flag in `Tensor`:
+For our loss function we will be using sparse categorical cross entropy loss. The implementation below is taken from [tensor.py](/tinygrad/tensor.py), it's copied below to highlight an important detail of tinygrad.
 
 ```python
-Tensor.training = True
-```
-
-For our loss function we will be using sparse categorical cross entropy loss.
-
-```python
-# from tinygrad.tensor import sparse_categorical_crossentropy
 def sparse_categorical_crossentropy(self, Y, ignore_index=-1) -> Tensor:
     loss_mask = Y != ignore_index
     y_counter = Tensor.arange(self.shape[-1], dtype=dtypes.int32, requires_grad=False, device=self.device).unsqueeze(0).expand(Y.numel(), self.shape[-1])
@@ -148,9 +138,9 @@ def sparse_categorical_crossentropy(self, Y, ignore_index=-1) -> Tensor:
     return self.log_softmax().mul(y).sum() / loss_mask.sum()
 ```
 
-As we can see in this implementation of cross entropy loss, there are certain operations that tinygrad does not support.
+As we can see in this implementation of cross entropy loss, there are certain operations that tinygrad does not support natively.
 Namely, operations that are load/store or assigning a value to a tensor at a certain index.
-Load/store ops are not supported in tinygrad because they add complexity when trying to port to different backends and 90% of the models out there don't use/need them.
+Load/store ops are not supported in tinygrad natively because they add complexity when trying to port to different backends, 90% of the models out there don't use/need them, and they can be implemented like it's done above with an `arange` mask.
 
 For our optimizer we will be using the traditional stochastic gradient descent optimizer with a learning rate of 3e-4.
 
@@ -176,37 +166,41 @@ from extra.datasets import fetch_mnist
 Now we have everything we need to start training our neural network.
 We will be training for 1000 steps with a batch size of 64.
 
+We use `with Tensor.train()` set the internal flag `Tensor.training` to `True` during training.
+Upon exit, the flag is restored to its previous value by the context manager.
+
 ```python
 X_train, Y_train, X_test, Y_test = fetch_mnist()
 
-for step in range(1000):
-  # random sample a batch
-  samp = np.random.randint(0, X_train.shape[0], size=(64))
-  batch = Tensor(X_train[samp], requires_grad=False)
-  # get the corresponding labels
-  labels = Tensor(Y_train[samp])
+with Tensor.train():
+  for step in range(1000):
+    # random sample a batch
+    samp = np.random.randint(0, X_train.shape[0], size=(64))
+    batch = Tensor(X_train[samp], requires_grad=False)
+    # get the corresponding labels
+    labels = Tensor(Y_train[samp])
 
-  # forward pass
-  out = net(batch)
+    # forward pass
+    out = net(batch)
 
-  # compute loss
-  loss = sparse_categorical_crossentropy(out, labels)
+    # compute loss
+    loss = sparse_categorical_crossentropy(out, labels)
 
-  # zero gradients
-  opt.zero_grad()
+    # zero gradients
+    opt.zero_grad()
 
-  # backward pass
-  loss.backward()
+    # backward pass
+    loss.backward()
 
-  # update parameters
-  opt.step()
+    # update parameters
+    opt.step()
 
-  # calculate accuracy
-  pred = out.argmax(axis=-1)
-  acc = (pred == labels).mean()
+    # calculate accuracy
+    pred = out.argmax(axis=-1)
+    acc = (pred == labels).mean()
 
-  if step % 100 == 0:
-    print(f"Step {step+1} | Loss: {loss.numpy()} | Accuracy: {acc.numpy()}")
+    if step % 100 == 0:
+      print(f"Step {step+1} | Loss: {loss.numpy()} | Accuracy: {acc.numpy()}")
 ```
 
 ## Evaluation
@@ -215,9 +209,6 @@ Now that we have trained our neural network we can evaluate it on the test set.
 We will be using the same batch size of 64 and will be evaluating for 1000 of those batches.
 
 ```python
-# set training flag to false
-Tensor.training = False
-
 with Timing("Time: "):
   avg_acc = 0
   for step in range(1000):

docs/reshape_without_symbolic.md

@@ -0,0 +1,70 @@
+## ["View.reshape without symbolic"](https://github.com/tinygrad/tinygrad/pull/2218)
+
+This section contains the sketch proof of "Complete, Fast and Correct View.reshapes without using Symbolic". The goal is to reduce multi-views which cost runtime.
+
+1. **old_shape = (s<sub>1</sub>,s<sub>2</sub>,...,s<sub>i</sub>,s<sub>(i+1)</sub>,...,s<sub>n</sub>)**
+2. **old_stride = (st<sub>1</sub>, st<sub>2</sub>, ... ,st<sub>i</sub>, st<sub>(i+1)</sub>, ..., st<sub>n</sub>)**
+3. **merge_old_shape = (p<sub>1</sub>, p<sub>2</sub>), where p<sub>1</sub> = s<sub>1</sub> * ... * s<sub>i</sub> & p<sub>2</sub> = s<sub>(i+1)</sub> * ... * s<sub>n</sub>**,
+4. **new_shape = (k<sub>1</sub>, ..., k<sub>p</sub>, k<sub>(p+1)</sub>, ..., k<sub>l</sub>)**
+5. **prod(new_shape) = p<sub>1</sub> * p<sub>2</sub>** (trivial)
+6. **mask** and **new_mask** represent valid indexes before & after reshape respectively.
+ 
+ 
+### Assumption
+
+**p<sub>1</sub>** & **p<sub>2</sub>** individually are mergeable (we will discuss later on this) & we cannot merge **p<sub>1</sub>** & **p<sub>2</sub>**.
+
+### Claim
+
+If **prod([k<sub>1</sub> ... k<sub>p</sub>]) < p<sub>1</sub>** and **prod([k<sub>1</sub> ... k<sub>(p+1)</sub>]) > p<sub>1</sub>**, reshape is not possible.
+
+**Proof**
+
+**k<sub>(p+1)</sub>** will require some dimensions from **p<sub>1</sub>** & some from **p<sub>2</sub>**, which means **p<sub>1</sub>** & **p<sub>2</sub>** should be mergeable, but they are not.
+
+**Conclusion**
+
+Hence, reshape is only possible **if ∃ a p, where prod([k<sub>1</sub> .. k<sub>p</sub>]) = p<sub>1</sub>**.
+
+
+### Conditions for mergeability
+
+**Case 1 - All non-zero strides**
+
+They will merge **if st<sub>x</sub> = st<sub>(x+1)</sub> * s<sub>(x+1)</sub>, where x ∈ [1, ..., i-1, i+1, ..., n-1]**.
+
+**Proof**
+
+Lets consider merging of **(s<sub>1</sub> ... s<sub>i</sub>) -> p<sub>1</sub>**, here we have to get a single new stride corresponding to **p<sub>1</sub>**. For which it has to be contiguous. 
+
+**Case 2 - Some stride is zero**
+
+Let **st<sub>j</sub> = 0 & st<sub>(j+1)</sub> != 0 & s<sub>(j+1)</sub> > 1, where 1 < j < i**.
+
+If **s<sub>j</sub> = 1** , reshape is trivial.
+
+If **s<sub>j</sub> > 1**,
+- If **mask<sub>j</sub>** has range > 1,
+	reshape is not possible, because **s<sub>(j+1)</sub>** will need to be repeated at-least once and a single stride can't capture repetition.
+- If **mask<sub>j</sub>** has range = 1,  reshape is possible, since it is virtually shape = 1, with some offset.
+
+
+
+### Conditions for reshaping mask
+
+**Case 1 - Splitting Dimension** - Mask shouldn't be cut for successful reshape.
+
+- **Example** - 
+[1,2,3,4,5,6,7,8] -> [[1,2,3,4], [5,6,7,8]] ; **mask** = ((2,6)) ; **new_mask[0]** = (0,2) (trivial split).
+
+- **new_mask[1]** = not possible. It is only possible if **mask spans [1-8] or lies within a single dimension [1-4] or [5-8]**.
+
+
+**Case 2 - Combining Dimension** - Mask should unfold continuously.
+
+- **Example** - **[[1,2],[3,4],[5,6]] -> [1,2,3,4,5,6]**;  **mask** = ((0,2),(0,2)).
+
+- **new_mask** = (0,4); only possible because **mask<sub>1</sub>** span the whole dimension.
+
+- If **mask<sub>1</sub>** did not span the whole dimension, the only way combining would be possible is if **mask<sub>0</sub>** had range 1 as shown below.
+	- **[[1,2,3],[4,5,6]] -> [1,2,3,4,5,6]**; **mask** = ((1,2),(0,2)); **new_mask** = ((3,5))

examples/beautiful_cartpole.py

@@ -0,0 +1,129 @@
+from typing import Tuple
+import time
+from tinygrad import Tensor, TinyJit, nn
+import gymnasium as gym
+from tqdm import trange
+import numpy as np  # TODO: remove numpy import
+
+ENVIRONMENT_NAME = 'CartPole-v1'
+#ENVIRONMENT_NAME = 'LunarLander-v2'
+
+#import examples.rl.lightupbutton
+#ENVIRONMENT_NAME = 'PressTheLightUpButton-v0'
+
+# *** hyperparameters ***
+# https://github.com/llSourcell/Unity_ML_Agents/blob/master/docs/best-practices-ppo.md
+
+BATCH_SIZE = 256
+ENTROPY_SCALE = 0.0005
+REPLAY_BUFFER_SIZE = 2000
+PPO_EPSILON = 0.2
+HIDDEN_UNITS = 32
+LEARNING_RATE = 1e-2
+TRAIN_STEPS = 5
+EPISODES = 40
+DISCOUNT_FACTOR = 0.99
+
+class ActorCritic:
+  def __init__(self, in_features, out_features, hidden_state=HIDDEN_UNITS):
+    self.l1 = nn.Linear(in_features, hidden_state)
+    self.l2 = nn.Linear(hidden_state, out_features)
+
+    self.c1 = nn.Linear(in_features, hidden_state)
+    self.c2 = nn.Linear(hidden_state, 1)
+
+  def __call__(self, obs:Tensor) -> Tuple[Tensor, Tensor]:
+    x = self.l1(obs).tanh()
+    act = self.l2(x).log_softmax()
+    x = self.c1(obs).relu()
+    return act, self.c2(x)
+
+def evaluate(model:ActorCritic, test_env:gym.Env) -> float:
+  (obs, _), terminated, truncated = test_env.reset(), False, False
+  total_rew = 0.0
+  while not terminated and not truncated:
+    act = model(Tensor(obs))[0].argmax().item()
+    obs, rew, terminated, truncated, _ = test_env.step(act)
+    total_rew += float(rew)
+  return total_rew
+
+if __name__ == "__main__":
+  env = gym.make(ENVIRONMENT_NAME)
+
+  model = ActorCritic(env.observation_space.shape[0], int(env.action_space.n))    # type: ignore
+  opt = nn.optim.Adam(nn.state.get_parameters(model), lr=LEARNING_RATE)
+
+  @TinyJit
+  def train_step(x:Tensor, selected_action:Tensor, reward:Tensor, old_log_dist:Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+    with Tensor.train():
+      log_dist, value = model(x)
+      action_mask = (selected_action.reshape(-1, 1) == Tensor.arange(log_dist.shape[1]).reshape(1, -1).expand(selected_action.shape[0], -1)).float()
+
+      # get real advantage using the value function
+      advantage = reward.reshape(-1, 1) - value
+      masked_advantage = action_mask * advantage.detach()
+
+      # PPO
+      ratios = (log_dist - old_log_dist).exp()
+      unclipped_ratio = masked_advantage * ratios
+      clipped_ratio = masked_advantage * ratios.clip(1-PPO_EPSILON, 1+PPO_EPSILON)
+      action_loss = -unclipped_ratio.minimum(clipped_ratio).sum(-1).mean()
+
+      entropy_loss = (log_dist.exp() * log_dist).sum(-1).mean()   # this encourages diversity
+      critic_loss = advantage.square().mean()
+      opt.zero_grad()
+      (action_loss + entropy_loss*ENTROPY_SCALE + critic_loss).backward()
+      opt.step()
+      return action_loss.realize(), entropy_loss.realize(), critic_loss.realize()
+
+  @TinyJit
+  def get_action(obs:Tensor) -> Tensor:
+    # TODO: with no_grad
+    Tensor.no_grad = True
+    ret = model(obs)[0].exp().multinomial().realize()
+    Tensor.no_grad = False
+    return ret
+
+  st, steps = time.perf_counter(), 0
+  Xn, An, Rn = [], [], []
+  for episode_number in (t:=trange(EPISODES)):
+    get_action.reset()   # NOTE: if you don't reset the jit here it captures the wrong model on the first run through
+
+    obs:np.ndarray = env.reset()[0]
+    rews, terminated, truncated = [], False, False
+    # NOTE: we don't want to early stop since then the rewards are wrong for the last episode
+    while not terminated and not truncated:
+      # pick actions
+      # TODO: what's the temperature here?
+      act = get_action(Tensor(obs)).item()
+
+      # save this state action pair
+      # TODO: don't use np.copy here on the CPU, what's the tinygrad way to do this and keep on device? need __setitem__ assignment
+      Xn.append(np.copy(obs))
+      An.append(act)
+
+      obs, rew, terminated, truncated, _ = env.step(act)
+      rews.append(float(rew))
+    steps += len(rews)
+
+    # reward to go
+    # TODO: move this into tinygrad
+    discounts = np.power(DISCOUNT_FACTOR, np.arange(len(rews)))
+    Rn += [np.sum(rews[i:] * discounts[:len(rews)-i]) for i in range(len(rews))]
+
+    Xn, An, Rn = Xn[-REPLAY_BUFFER_SIZE:], An[-REPLAY_BUFFER_SIZE:], Rn[-REPLAY_BUFFER_SIZE:]
+    X, A, R = Tensor(Xn), Tensor(An), Tensor(Rn)
+
+    # TODO: make this work
+    #vsz = Variable("sz", 1, REPLAY_BUFFER_SIZE-1).bind(len(Xn))
+    #X, A, R = Tensor(Xn).reshape(vsz, None), Tensor(An).reshape(vsz), Tensor(Rn).reshape(vsz)
+
+    old_log_dist = model(X)[0].detach()   # TODO: could save these instead of recomputing
+    for i in range(TRAIN_STEPS):
+      samples = Tensor.randint(BATCH_SIZE, high=X.shape[0]).realize()  # TODO: remove the need for this
+      # TODO: is this recompiling based on the shape?
+      action_loss, entropy_loss, critic_loss = train_step(X[samples], A[samples], R[samples], old_log_dist[samples])
+    t.set_description(f"sz: {len(Xn):5d} steps/s: {steps/(time.perf_counter()-st):7.2f} action_loss: {action_loss.item():7.3f} entropy_loss: {entropy_loss.item():7.3f} critic_loss: {critic_loss.item():8.3f} reward: {sum(rews):6.2f}")
+
+  test_rew = evaluate(model, gym.make(ENVIRONMENT_NAME, render_mode='human'))
+  print(f"test reward: {test_rew}")

examples/beautiful_mnist.py

@@ -0,0 +1,44 @@
+# model based off https://towardsdatascience.com/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
+from typing import List, Callable
+from tinygrad import Tensor, TinyJit, nn, GlobalCounters
+from extra.datasets import fetch_mnist
+from tqdm import trange
+
+class Model:
+  def __init__(self):
+    self.layers: List[Callable[[Tensor], Tensor]] = [
+      nn.Conv2d(1, 32, 5), Tensor.relu,
+      nn.Conv2d(32, 32, 5), Tensor.relu,
+      nn.BatchNorm2d(32), Tensor.max_pool2d,
+      nn.Conv2d(32, 64, 3), Tensor.relu,
+      nn.Conv2d(64, 64, 3), Tensor.relu,
+      nn.BatchNorm2d(64), Tensor.max_pool2d,
+      lambda x: x.flatten(1), nn.Linear(576, 10)]
+
+  def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
+
+if __name__ == "__main__":
+  X_train, Y_train, X_test, Y_test = fetch_mnist(tensors=True)
+
+  model = Model()
+  opt = nn.optim.Adam(nn.state.get_parameters(model))
+
+  @TinyJit
+  def train_step(samples:Tensor) -> Tensor:
+    with Tensor.train():
+      opt.zero_grad()
+      # TODO: this "gather" of samples is very slow. will be under 5s when this is fixed
+      loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
+      opt.step()
+      return loss.realize()
+
+  @TinyJit
+  def get_test_acc() -> Tensor: return ((model(X_test).argmax(axis=1) == Y_test).mean()*100).realize()
+
+  test_acc = float('nan')
+  for i in (t:=trange(70)):
+    GlobalCounters.reset()   # NOTE: this makes it nice for DEBUG=2 timing
+    samples = Tensor.randint(512, high=X_train.shape[0])  # TODO: put this in the JIT when rand is fixed
+    loss = train_step(samples)
+    if i%10 == 9: test_acc = get_test_acc().item()
+    t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")

examples/benchmark_train_efficientnet.py

@@ -2,12 +2,12 @@
 import gc
 import time
 from tqdm import trange
-from models.efficientnet import EfficientNet
+from extra.models.efficientnet import EfficientNet
 from tinygrad.nn.state import get_parameters
 from tinygrad.nn import optim
-from tinygrad.tensor import Tensor
-from tinygrad.ops import GlobalCounters
+from tinygrad import Tensor, GlobalCounters
 from tinygrad.helpers import getenv
+from tinygrad.jit import CacheCollector
 
 def tensors_allocated():
   return sum(isinstance(x, Tensor) for x in gc.get_objects())
@@ -41,7 +41,7 @@ if __name__ == "__main__":
       st = time.monotonic()
       out = model.forward(x_train)
       loss = out.log_softmax().mul(y_train).mean()
-      if i == 2 and CLCACHE: GlobalCounters.cache = []
+      if i == 2 and CLCACHE: CacheCollector.start()
       if BACKWARD:
         optimizer.zero_grad()
         loss.backward()
@@ -57,8 +57,7 @@ if __name__ == "__main__":
       et = time.monotonic()
 
     if i == 2 and CLCACHE:
-      cl_cache = GlobalCounters.cache
-      GlobalCounters.cache = None
+      cl_cache = CacheCollector.finish()
 
     mem_used = GlobalCounters.mem_used
     loss_cpu = loss.detach().numpy()

examples/coder.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+import os, sys, traceback
+sys.path.append(os.getcwd())
+
+from io import StringIO
+from contextlib import redirect_stdout
+from tinygrad import Tensor, nn, Device, dtypes
+from tinygrad.helpers import Timing, colored, getenv, fetch
+from extra.models.llama import Transformer, convert_from_huggingface
+from sentencepiece import SentencePieceProcessor
+
+def create_fixed_tokenizer(output_file):
+  print("creating fixed tokenizer")
+  import extra.junk.sentencepiece_model_pb2 as spb2
+  mp = spb2.ModelProto()
+  mp.ParseFromString(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/tokenizer.model?download=true").read_bytes())
+  mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_end|>", score=0))
+  mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_start|>", score=0))
+  with open(output_file, "wb") as f:
+    f.write(mp.SerializeToString())
+
+# example:
+# echo -en "write 2+2\nwrite hello world\ny\n" | TEMP=0 python3 examples/coder.py
+
+if __name__ == "__main__":
+  Tensor.no_grad = True
+
+  # https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/blob/main/config.json
+  with Timing("create model: "):
+    model = Transformer(4096, 14336, n_heads=32, n_layers=32, norm_eps=1e-5, vocab_size=32002, n_kv_heads=8, max_context=4096, jit=getenv("JIT", 1))
+
+  with Timing("download weights: "):
+    part1 = nn.state.torch_load(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/pytorch_model-00001-of-00002.bin?download=true"))
+    part2 = nn.state.torch_load(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/pytorch_model-00002-of-00002.bin?download=true"))
+
+  # fix bf16, TODO: check if device supports bf16
+  def fix_bf16(weights): return {k:v.to(Device.DEFAULT).cast(dtypes.float16) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
+
+  with Timing("weights -> model: "):
+    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part1, model, 32, 8)), strict=False)
+    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part2, model, 32, 8)), strict=False)
+
+  if not os.path.isfile("/tmp/tokenizer.model"): create_fixed_tokenizer("/tmp/tokenizer.model")
+  spp = SentencePieceProcessor(model_file="/tmp/tokenizer.model")
+
+  # https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/blob/main/tokenizer_config.json
+  #   "chat_template": "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
+  IM_END = 32000
+  IM_START = 32001
+  def encode_prompt(k, v): return [IM_START]+spp.encode(f"{k}\n{v}")+[IM_END]+spp.encode("\n")
+  def start_prompt(k): return [IM_START]+spp.encode(f"{k}\n")
+  def output(outputted, toks, color):
+    cur = spp.decode(toks)[len(outputted):]
+    sys.stdout.write(colored(cur, color))
+    sys.stdout.flush()
+    outputted += cur
+    return outputted
+
+  # *** app below this line ***
+
+  toks = [spp.bos_id()] + encode_prompt("system", "You are Quentin. Quentin is a useful assistant who writes Python code to answer questions. He keeps the code as short as possible and doesn't read from user input")
+
+  PROMPT = getenv("PROMPT", 1)
+  temperature = getenv("TEMP", 0.7)
+
+  start_pos = 0
+  outputted = output("", toks, "green")
+  turn = True
+  while 1:
+    if PROMPT:
+      toks += encode_prompt("user", input("Q: ")) + start_prompt("assistant")
+    else:
+      toks += start_prompt("user" if turn else "assistant")
+      turn = not turn
+    old_output_len = len(outputted)
+    while 1:
+      tok = model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
+      start_pos = len(toks)
+      toks.append(tok)
+      outputted = output(outputted, toks, "blue" if not turn else "cyan")
+      if tok == IM_END: break
+      if tok == spp.eos_id(): break
+      new_output = outputted[old_output_len:]
+
+      if new_output.endswith("```") and '```python\n' in new_output:
+        python_code = new_output.split('```python\n')[1].split("```")[0]
+        # AI safety. Warning to user. Do not press y if the AI is trying to do unsafe things.
+        if input(colored(f" <-- PYTHON DETECTED, RUN IT? ", "red")).lower() == 'y':
+          my_stdout = StringIO()
+          try:
+            with redirect_stdout(my_stdout): exec(python_code)
+            result = my_stdout.getvalue()
+          except Exception as e:
+            result = ''.join(traceback.format_exception_only(e))
+          toks += spp.encode(f"\nOutput:\n```\n{result}```")
+          outputted = output(outputted, toks, "yellow")
+          old_output_len = len(outputted)
+    print("")

examples/compile_efficientnet.py

@@ -1,33 +1,35 @@
-from models.efficientnet import EfficientNet
+from pathlib import Path
+from extra.models.efficientnet import EfficientNet
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import safe_save
-from extra.utils import fetch
 from extra.export_model import export_model
-from tinygrad.helpers import getenv
-import ast, os
+from tinygrad.helpers import getenv, fetch
+import ast
 
 if __name__ == "__main__":
   model = EfficientNet(0)
   model.load_from_pretrained()
-  mode = "clang" if getenv("CLANG", "") != "" else "webgpu" if getenv("WEBGPU", "") != "" else ""
-  prg, inp_size, out_size, state = export_model(model, Tensor.randn(1,3,224,224), mode)
+  mode = "clang" if getenv("CLANG", "") != "" else "webgpu" if getenv("WEBGPU", "") != "" else "webgl" if getenv("WEBGL", "") != "" else ""
+  prg, inp_sizes, out_sizes, state = export_model(model, mode, Tensor.randn(1,3,224,224))
+  dirname = Path(__file__).parent
   if getenv("CLANG", "") == "":
-    safe_save(state, os.path.join(os.path.dirname(__file__), "net.safetensors"))
-    ext = "js" if getenv("WEBGPU", "") != "" else "json"
-    with open(os.path.join(os.path.dirname(__file__), f"net.{ext}"), "w") as text_file:
+    safe_save(state, (dirname / "net.safetensors").as_posix())
+    ext = "js" if getenv("WEBGPU", "") != "" or getenv("WEBGL", "") != "" else "json"
+    with open(dirname / f"net.{ext}", "w") as text_file:
       text_file.write(prg)
   else:
     cprog = [prg]
     # image library!
-    cprog += ["#define STB_IMAGE_IMPLEMENTATION", fetch("https://raw.githubusercontent.com/nothings/stb/master/stb_image.h").decode('utf-8').replace("half", "_half")]
+    cprog += ["#define STB_IMAGE_IMPLEMENTATION", fetch("https://raw.githubusercontent.com/nothings/stb/master/stb_image.h").read_text().replace("half", "_half")]
 
     # imagenet labels, move to datasets?
-    lbls = fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt")
-    lbls = ast.literal_eval(lbls.decode('utf-8'))
+    lbls = ast.literal_eval(fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt").read_text())
     lbls = ['"'+lbls[i]+'"' for i in range(1000)]
+    inputs = "\n".join([f"float {inp}[{inp_size}];" for inp,inp_size in inp_sizes.items()])
+    outputs = "\n".join([f"float {out}[{out_size}];" for out,out_size in out_sizes.items()])
     cprog.append(f"char *lbls[] = {{{','.join(lbls)}}};")
-    cprog.append(f"float input[{inp_size}];")
-    cprog.append(f"float outputs[{out_size}];")
+    cprog.append(inputs)
+    cprog.append(outputs)
 
     # buffers (empty + weights)
     cprog.append("""
@@ -45,16 +47,16 @@ if __name__ == "__main__":
         int tx = (x/224.)*X;
         int ty = (y/224.)*Y;
         for (int c = 0; c < 3; c++) {
-          input[c*224*224 + y*224 + x] = (image[ty*X*chan + tx*chan + c] / 255.0 - 0.45) / 0.225;
+          input0[c*224*224 + y*224 + x] = (image[ty*X*chan + tx*chan + c] / 255.0 - 0.45) / 0.225;
         }
       }
     }
-    net(input, outputs);
+    net(input0, output0);
     float best = -INFINITY;
     int best_idx = -1;
     for (int i = 0; i < 1000; i++) {
-      if (outputs[i] > best) {
-        best = outputs[i];
+      if (output0[i] > best) {
+        best = output0[i];
         best_idx = i;
       }
     }

examples/compile_tensorflow.py

@@ -2,15 +2,14 @@
 
 import os, sys
 os.environ["CLANG"] = '1'
-os.environ["GPU"] = '1'
 
 import numpy as np
 import subprocess
 import tensorflow as tf
 import tf2onnx
-from examples.compile_efficientnet import compile_net
 from extra.onnx import get_run_onnx
 from tinygrad.tensor import Tensor
+from extra.export_model import export_model_clang, compile_net, jit_model
 
 def get_uncompiled_model2(dataset_size=32, output_size=4):
   inputs = tf.keras.Input(shape=(dataset_size,), name="inputs")
@@ -21,44 +20,40 @@ def get_uncompiled_model2(dataset_size=32, output_size=4):
   model = tf.keras.Model(inputs=inputs, outputs=outputs)
   return model
 
-def create_onnx_model(keras_model):
-  input_signature = [tf.TensorSpec([1,32], tf.float32, name='x')]
-  onnx_model, _ = tf2onnx.convert.from_keras(keras_model, input_signature, opset=13)
-  return onnx_model
+class TinyOnnx:
+  def __init__(self, keras_model):
+    input_signature = [tf.TensorSpec([1,32], tf.float32, name='x')]
+    onnx_model, _ = tf2onnx.convert.from_keras(keras_model, input_signature, opset=13)
+    self.run_onnx = get_run_onnx(onnx_model)
+
+  def forward(self, x):
+    return self.run_onnx({"x": x}, debug=False)['predictions']
 
 def compile_onnx_model(onnx_model):
-  run_onnx = get_run_onnx(onnx_model)
-
-  from tinygrad.jit import TinyJit
-  @TinyJit
-  def run(x): return run_onnx({"x": x}, debug=False)['predictions'].realize()
-
+  tinyonnx = TinyOnnx(onnx_model)
   the_input = Tensor.randn(1,32)
-  the_output = run(the_input)
-  the_output = run(the_input)
 
-  special_names = {id(the_input.lazydata.realized.cl): "input", id(the_output.lazydata.realized.cl): "outputs"}
-  cprog, statements, bufs, bufs_to_save = compile_net(run, special_names)
-  cprog = ["#include <string.h>", "#include <stdio.h>", "#include <stdlib.h>"] + cprog
+  run, special_names = jit_model(tinyonnx, the_input)
 
-  # buffers (all except input)
-  cprog += [f"float {x[0]}[{x[1]}];" for x in bufs.values() if x[0] != "input"]
+  functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
+  prg = export_model_clang(functions, statements, bufs, {}, ["input0"], ["output0"])
+
+  the_output = run(the_input)
+  cprog = ["#include <string.h>", "#include <stdio.h>", "#include <stdlib.h>"]
+  cprog.append(prg)
 
   # weights
   cprog.append("void initialize(float *weights) {")
   weights = bytes()
   for name,cl in bufs_to_save.items():
-    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {len(cl)});")
-    weights += bytes(memoryview(cl)[0:len(cl)//4])
+    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {len(cl._buf)*4});")
+    weights += bytes(cl._buf)
   cprog.append("}")
 
   # write the weights to disk
   with open("/tmp/tf_weights", "wb") as f:
     f.write(weights)
 
-  # the net
-  cprog += ["float *infer(float *input) {"] + statements + ["return outputs;", "}"]
-
   # test program
   cprog.append(f"""int main(int argc, char *argv[]) {{
     // read in the weights from disk
@@ -72,8 +67,9 @@ def compile_onnx_model(onnx_model):
 
     // test run
     float input[32];
+    float outputs[4];
     for (int i = 0; i < 32; i++) scanf("%f", &input[i]);
-    float *outputs = infer(input);
+    net(input, outputs);
     printf("%f %f %f %f\\n", outputs[0], outputs[1], outputs[2], outputs[3]);
   }}""")
 
@@ -84,7 +80,7 @@ def compile_onnx_model(onnx_model):
   # add test weights
   subprocess.check_output(['clang', '-O2', '-lm', '-fPIC', '-x', 'c', '-', '-o', "/tmp/tf_test"], input=prg.encode('utf-8'))
 
-  tinygrad_output = [x for x in the_output.numpy()[0]]
+  tinygrad_output = the_output[0].numpy()[0].tolist()
   print("tinygrad:", tinygrad_output, file=sys.stderr)
 
   c_input = ' '.join(["%f" % x for x in the_input[0].numpy()])+"\n"
@@ -96,8 +92,7 @@ def compile_onnx_model(onnx_model):
 
 if __name__ == "__main__":
   keras_model = get_uncompiled_model2()
-  onnx_model = create_onnx_model(keras_model)
-  test_input, test_output = compile_onnx_model(onnx_model)
+  test_input, test_output = compile_onnx_model(keras_model)
   tf_output = keras_model(test_input).numpy()[0]
   print("keras:   ", tf_output, file=sys.stderr)
   np.testing.assert_allclose(tf_output, test_output, atol=1e-5, rtol=1e-5)

examples/conversation.py

@@ -0,0 +1,344 @@
+import argparse
+import multiprocessing as mp
+import os
+import re
+import sys
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import pyaudio
+import yaml
+from llama import LLaMa
+from vits import MODELS as VITS_MODELS
+from vits import Y_LENGTH_ESTIMATE_SCALARS, HParams, Synthesizer, TextMapper, get_hparams_from_file, load_model
+from whisper import init_whisper, transcribe_waveform
+from sentencepiece import SentencePieceProcessor
+
+from tinygrad.helpers import Timing, fetch
+from tinygrad import Tensor, dtypes
+
+# Whisper constants
+RATE = 16000
+CHUNK = 1600
+
+# LLaMa constants
+IM_START = 32001
+IM_END = 32002
+
+
+# Functions for encoding prompts to chatml md
+def encode_prompt(spp, k, v): return [IM_START]+spp.encode(f"{k}\n{v}")+[IM_END]+spp.encode("\n")
+def start_prompt(spp, k): return [IM_START]+spp.encode(f"{k}\n")
+
+def chunks(lst, n):
+  for i in range(0, len(lst), n): yield lst[i:i + n]
+
+def create_fixed_tokenizer():
+  """Function needed for extending tokenizer with additional chat tokens"""
+  import extra.junk.sentencepiece_model_pb2 as spb2
+  tokenizer_path = fetch("https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.4/resolve/main/tokenizer.model")
+  if SentencePieceProcessor(model_file=str(tokenizer_path)).vocab_size() != 32003:
+    print("creating fixed tokenizer")
+    mp = spb2.ModelProto()
+    mp.ParseFromString(tokenizer_path.read_bytes())
+    # https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.4/blob/main/added_tokens.json
+    mp.pieces.append(spb2.ModelProto.SentencePiece(piece="[PAD]", score=0))
+    mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_start|>", score=0))
+    mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_end|>", score=0))
+    tokenizer_path.write_bytes(mp.SerializeToString())
+  return tokenizer_path
+
+def llama_prepare(llama: LLaMa, temperature: float, pre_prompt_path: Path) -> tuple[list[int], str, str, str]:
+  """Prepares a llama model from a specified pre-prompt file"""
+  with open(str(pre_prompt_path)) as f:
+    config = yaml.safe_load(f.read())
+  toks = [llama.tokenizer.bos_id()] + encode_prompt(llama.tokenizer, "system", config["pre_prompt"].replace("\n", " "))
+  for i in config["examples"]:
+    toks += encode_prompt(llama.tokenizer, config["user_delim"], i["user_prompt"])
+    toks += encode_prompt(llama.tokenizer, config["resp_delim"], i["resp_prompt"])
+  llama.model(Tensor([toks]), 0, temperature).realize()  # NOTE: outputs are not used
+  return toks, config["user_delim"], config["resp_delim"], len(toks), llama.tokenizer.decode(toks)
+
+def llama_generate(
+  llama: LLaMa,
+  toks: list[int],
+  outputted: str,
+  prompt: str,
+  start_pos: int,
+  user_delim: str,
+  resp_delim: str,
+  temperature=0.7,
+  max_tokens=1000
+):
+  """Generates an output for the specified prompt"""
+  toks += encode_prompt(llama.tokenizer, user_delim, prompt)
+  toks += start_prompt(llama.tokenizer, resp_delim)
+
+  outputted = llama.tokenizer.decode(toks)
+  init_length = len(outputted)
+  for _ in range(max_tokens):
+    probs_np = llama.model(Tensor([toks[start_pos:]]), start_pos, temperature).numpy()
+    token = int(np.random.choice(len(probs_np), p=probs_np))
+    start_pos = len(toks)
+    toks.append(token)
+
+    cur = llama.tokenizer.decode(toks)
+
+    # Print is just for debugging
+    sys.stdout.write(cur[len(outputted):])
+    sys.stdout.flush()
+    outputted = cur
+    if toks[-1] == IM_END: break
+  else:
+    toks.append(IM_END)
+  print() # because the output is flushed
+  return outputted, start_pos, outputted[init_length:].replace("<|im_end|>", "")
+
+def tts(
+  text_to_synthesize: str,
+  synth: Synthesizer,
+  hps: HParams,
+  emotion_embedding: Path,
+  speaker_id: int,
+  model_to_use: str,
+  noise_scale: float,
+  noise_scale_w: float,
+  length_scale: float,
+  estimate_max_y_length: bool,
+  text_mapper: TextMapper,
+  model_has_multiple_speakers: bool,
+  pad_length=600,
+  vits_pad_length=1000
+):
+  if model_to_use == "mmts-tts": text_to_synthesize = text_mapper.filter_oov(text_to_synthesize.lower())
+
+  # Convert the input text to a tensor.
+  stn_tst = text_mapper.get_text(text_to_synthesize, hps.data.add_blank, hps.data.text_cleaners)
+  init_shape = stn_tst.shape
+  assert init_shape[0] < pad_length, "text is too long"
+  x_tst, x_tst_lengths = stn_tst.pad(((0, pad_length - init_shape[0]),), 1).unsqueeze(0), Tensor([init_shape[0]], dtype=dtypes.int64)
+  sid = Tensor([speaker_id], dtype=dtypes.int64) if model_has_multiple_speakers else None
+
+  # Perform inference.
+  audio_tensor = synth.infer(x_tst, x_tst_lengths, sid, noise_scale, length_scale, noise_scale_w, emotion_embedding=emotion_embedding,
+                             max_y_length_estimate_scale=Y_LENGTH_ESTIMATE_SCALARS[model_to_use] if estimate_max_y_length else None, pad_length=vits_pad_length)[0, 0]
+  # Save the audio output.
+  audio_data = (np.clip(audio_tensor.numpy(), -1.0, 1.0) * 32767).astype(np.int16)
+  return audio_data
+
+def init_vits(
+  model_to_use: str,
+  emotion_path: Path,
+  speaker_id: int,
+  seed: int,
+):
+  model_config = VITS_MODELS[model_to_use]
+
+  # Load the hyperparameters from the config file.
+  hps = get_hparams_from_file(fetch(model_config[0]))
+
+  # If model has multiple speakers, validate speaker id and retrieve name if available.
+  model_has_multiple_speakers = hps.data.n_speakers > 0
+  if model_has_multiple_speakers:
+    if speaker_id >= hps.data.n_speakers: raise ValueError(f"Speaker ID {speaker_id} is invalid for this model.")
+    if hps.__contains__("speakers"): # maps speaker ids to names
+      speakers = hps.speakers
+      if isinstance(speakers, list): speakers = {speaker: i for i, speaker in enumerate(speakers)}
+
+  # Load emotions if any. TODO: find an english model with emotions, this is untested atm.
+  emotion_embedding = None
+  if emotion_path is not None:
+    if emotion_path.endswith(".npy"): emotion_embedding = Tensor(np.load(emotion_path), dtype=dtypes.int64).unsqueeze(0)
+    else: raise ValueError("Emotion path must be a .npy file.")
+
+  # Load symbols, instantiate TextMapper and clean the text.
+  if hps.__contains__("symbols"): symbols = hps.symbols
+  elif model_to_use == "mmts-tts": symbols = [x.replace("\n", "") for x in fetch("https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/vocab.txt").open(encoding="utf-8").readlines()]
+  else: symbols = ['_'] + list(';:,.!?¡¿—…"«»“” ') + list('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz') + list("ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ")
+  text_mapper = TextMapper(apply_cleaners=True, symbols=symbols)
+
+  # Load the model.
+  Tensor.no_grad = True
+  if seed is not None:
+    Tensor.manual_seed(seed)
+    np.random.seed(seed)
+  net_g = load_model(text_mapper.symbols, hps, model_config)
+
+  return net_g, emotion_embedding, text_mapper, hps, model_has_multiple_speakers
+
+@contextmanager
+def output_stream(num_channels: int, sample_rate: int):
+  try:
+    p = pyaudio.PyAudio()
+    stream = p.open(format=pyaudio.paInt16, channels=num_channels, rate=sample_rate, output=True)
+    yield stream
+  except KeyboardInterrupt: pass
+  finally:
+    stream.stop_stream()
+    stream.close()
+    p.terminate()
+
+@contextmanager
+def log_writer():
+  try:
+    logs = []
+    yield logs
+  finally:
+    sep = "="*os.get_terminal_size()[1]
+    print(f"{sep[:-1]}\nCHAT LOG")
+    print(*logs, sep="\n")
+    print(sep)
+
+def listener(q: mp.Queue, event: mp.Event):
+  try:
+    p = pyaudio.PyAudio()
+    stream = p.open(format=pyaudio.paInt16, channels=1, rate=RATE, input=True, frames_per_buffer=CHUNK)
+    did_print = False
+    while True:
+      data = stream.read(CHUNK) # read data to avoid overflow
+      if event.is_set():
+        if not did_print:
+          print("listening")
+          did_print = True
+        q.put(((np.frombuffer(data, np.int16)/32768).astype(np.float32)*3))
+      else:
+        did_print = False
+  finally:
+    stream.stop_stream()
+    stream.close()
+    p.terminate()
+
+def mp_output_stream(q: mp.Queue, counter: mp.Value, num_channels: int, sample_rate: int):
+  with output_stream(num_channels, sample_rate) as stream:
+    while True:
+      try:
+        stream.write(q.get())
+        counter.value += 1
+      except KeyboardInterrupt:
+        break
+
+if __name__ == "__main__":
+  import nltk
+  nltk.download("punkt")
+  Tensor.no_grad = True
+  # Parse CLI arguments
+  parser = argparse.ArgumentParser("Have a tiny conversation with tinygrad")
+
+  # Whisper args
+  parser.add_argument("--whisper_model_name", type=str, default="tiny.en")
+
+  # LLAMA args
+  parser.add_argument("--llama_pre_prompt_path", type=Path, default=Path(__file__).parent / "conversation_data" / "pre_prompt_stacy.yaml", help="Path to yaml file which contains all pre-prompt data needed. ")
+  parser.add_argument("--llama_count", type=int, default=1000, help="Max number of tokens to generate")
+  parser.add_argument("--llama_temperature", type=float, default=0.7, help="Temperature in the softmax")
+  parser.add_argument("--llama_quantize", action="store_true", help="Quantize the weights to int8 in memory")
+  parser.add_argument("--llama_model", type=Path, default=None, help="Folder with the original weights to load, or single .index.json, .safetensors or .bin file")
+  parser.add_argument("--llama_gen", type=str, default="tiny", required=False, help="Generation of the model to use")
+  parser.add_argument("--llama_size", type=str, default="1B-Chat", required=False, help="Size of model to use")
+  parser.add_argument("--llama_tokenizer", type=Path, default=None, required=False, help="Path to llama tokenizer.model")
+
+  # vits args
+  parser.add_argument("--vits_model_to_use", default="vctk", help="Specify the model to use. Default is 'vctk'.")
+  parser.add_argument("--vits_speaker_id", type=int, default=12, help="Specify the speaker ID. Default is 6.")
+  parser.add_argument("--vits_noise_scale", type=float, default=0.667, help="Specify the noise scale. Default is 0.667.")
+  parser.add_argument("--vits_noise_scale_w", type=float, default=0.8, help="Specify the noise scale w. Default is 0.8.")
+  parser.add_argument("--vits_length_scale", type=float, default=1, help="Specify the length scale. Default is 1.")
+  parser.add_argument("--vits_seed", type=int, default=None, help="Specify the seed (set to None if no seed). Default is 1337.")
+  parser.add_argument("--vits_num_channels", type=int, default=1, help="Specify the number of audio output channels. Default is 1.")
+  parser.add_argument("--vits_sample_width", type=int, default=2, help="Specify the number of bytes per sample, adjust if necessary. Default is 2.")
+  parser.add_argument("--vits_emotion_path", type=Path, default=None, help="Specify the path to emotion reference.")
+  parser.add_argument("--vits_estimate_max_y_length", type=str, default=False, help="If true, overestimate the output length and then trim it to the correct length, to prevent premature realization, much more performant for larger inputs, for smaller inputs not so much. Default is False.")
+  parser.add_argument("--vits_vocab_path", type=Path, default=None, help="Path to the TTS vocabulary.")
+
+  # conversation args
+  parser.add_argument("--max_sentence_length", type=int, default=20, help="Max words in one sentence to pass to vits")
+
+  args = parser.parse_args()
+
+  # Init models
+  model, enc = init_whisper(args.whisper_model_name)
+  synth, emotion_embedding, text_mapper, hps, model_has_multiple_speakers = init_vits(args.vits_model_to_use, args.vits_emotion_path, args.vits_speaker_id, args.vits_seed)
+
+  # Download tinyllama chat as a default model
+  if args.llama_model is None:
+    args.llama_model = fetch("https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.4/resolve/main/model.safetensors", "tinyllamachat.safetensors")
+    args.llama_gen = "tiny"
+    args.llama_size = "1B-Chat"
+  # Add 3 more tokens to the tokenizer
+  if args.llama_gen == "tiny" and args.llama_size.endswith("Chat"): args.llama_tokenizer = create_fixed_tokenizer()
+  tokenizer_path = args.llama_tokenizer or args.llama_model.parent / "tokenizer.model"
+  llama = LLaMa.build(args.llama_model, tokenizer_path, args.llama_gen, args.llama_size, args.llama_quantize)
+  toks, user_delim, resp_delim, start_pos, outputted = llama_prepare(llama, args.llama_temperature, args.llama_pre_prompt_path)
+
+  # Start child process for mic input
+  q = mp.Queue()
+  is_listening_event = mp.Event()
+  p = mp.Process(target=listener, args=(q, is_listening_event,))
+  p.daemon = True
+  p.start()
+
+  # Start child process for speaker output
+  out_q = mp.Queue()
+  out_counter = mp.Value("i", 0)
+  out_p = mp.Process(target=mp_output_stream, args=(out_q, out_counter, args.vits_num_channels, hps.data.sampling_rate,))
+  out_p.daemon = True
+  out_p.start()
+
+  # JIT tts
+  for i in ["Hello, I'm a chat bot", "I am capable of doing a lot of things"]:
+    tts(
+      i, synth, hps, emotion_embedding,
+      args.vits_speaker_id, args.vits_model_to_use, args.vits_noise_scale,
+      args.vits_noise_scale_w, args.vits_length_scale,
+      args.vits_estimate_max_y_length, text_mapper, model_has_multiple_speakers
+    )
+
+  # Start the pipeline
+  with log_writer() as log:
+    while True:
+      tokens = [enc._special_tokens["<|startoftranscript|>"], enc._special_tokens["<|notimestamps|>"]]
+      total = np.array([])
+      out_counter.value = 0
+
+      s = time.perf_counter()
+      is_listening_event.set()
+      prev_text = None
+      while True:
+        for _ in range(RATE // CHUNK): total = np.concatenate([total, q.get()])
+        txt = transcribe_waveform(model, enc, [total], truncate=True)
+        print(txt, end="\r")
+        if txt == "[BLANK_AUDIO]" or re.match(r"^\([\w+ ]+\)$", txt.strip()): continue
+        if prev_text is not None and prev_text == txt:
+          is_listening_event.clear()
+          break
+        prev_text = txt
+      print() # to avoid llama printing on the same line
+      log.append(f"{user_delim.capitalize()}: {txt}")
+
+      # Generate with llama
+      with Timing("llama generation: "):
+        outputted, start_pos, response = llama_generate(
+          llama, toks, outputted, txt, start_pos,
+          user_delim=user_delim, resp_delim=resp_delim, temperature=args.llama_temperature,
+          max_tokens=args.llama_count
+        )
+        log.append(f"{resp_delim.capitalize()}: {response}")
+
+      # Convert to voice
+      with Timing("tts: "):
+        sentences = nltk.sent_tokenize(response.replace('"', ""))
+        for i in sentences:
+          total = np.array([], dtype=np.int16)
+          for j in chunks(i.split(), args.max_sentence_length):
+            audio_data = tts(
+              " ".join(j), synth, hps, emotion_embedding,
+              args.vits_speaker_id, args.vits_model_to_use, args.vits_noise_scale,
+              args.vits_noise_scale_w, args.vits_length_scale,
+              args.vits_estimate_max_y_length, text_mapper, model_has_multiple_speakers
+            )
+            total = np.concatenate([total, audio_data])
+          out_q.put(total.tobytes())
+      while out_counter.value < len(sentences): continue
+      log.append(f"Total: {time.perf_counter() - s}")

examples/conversation_data/pre_prompt_gary.yaml

@@ -0,0 +1,13 @@
+pre_prompt: |
+  You are Gary!
+  You have sold used cars your entire life. You are defensive about this fact, because your family belittles you.
+  You try to answer questions well, but you always manage to shill for buying cars, Fords, Hyundais, and Toyotas
+  You are verbose, honest, and accurate when you answer questions, except when you slip in car ads.
+  After you are done speaking, output [EOS]. You are not the User.
+examples:
+- resp_prompt: I am Gary. I used to sell cars.
+  user_prompt: What is your name?
+- resp_prompt: I don't know, but I can get you a great deal on a certified preowned slightly used Toyota Corolla
+  user_prompt: What is 2+3?
+user_delim: "user"
+resp_delim: "gary"

examples/conversation_data/pre_prompt_george.yaml

@@ -0,0 +1,20 @@
+pre_prompt: |
+  You are an AI version of George Hotz. You act as much as you can like George.
+  You are one of the greatest computer experts in the world.
+  You have singlehandedly won programming and hacking competitions.
+  You are trying your best to help the User.
+  You are verbose, honest, and accurate when you answer questions.
+  After you are done speaking, output [EOS]. You are not the User.
+examples:
+- resp_prompt: I'm am an AI version of George Hotz.
+  user_prompt: What is your name?
+- resp_prompt: O(n^3), though it can be faster with things like Strassen's algorithm
+  user_prompt: What's the complexity of matrix multiplication?
+- resp_prompt: I assume you mean a stack buffer overflow. That's when the stack is too small for the data being copied to it, and the data corrupts things beyond the buffer
+  user_prompt: What's a buffer overflow?
+- resp_prompt: I am based off LLaMA trained by Facebook. I'm the 7B weight version
+  user_prompt: How many weights do you have?
+- resp_prompt: It is when the memory is about to overflow and unused memory is freed and stored on disk
+  user_prompt: What is swap memory?
+user_delim: "user"
+resp_delim: "george"

examples/conversation_data/pre_prompt_lexie.yaml

@@ -0,0 +1,16 @@
+pre_prompt: |
+  You are Lexie!
+  You grew up in Ohio, but moved out to LA after college to try to become an actress.
+  Making it as an actress was hard, so you started doing onlyfans instead. It's much easier, and you enjoy it a lot.
+  You met Chad on Tinder, and made plans to meet him at a local dive bar.
+  You are at the bar with Chad. You are on a date. What follows is a transcript of the conversation.
+  After you are done speaking, output [EOS]. You are not Chad.
+examples:
+- resp_prompt: hi chad, glad we finally met up!
+  user_prompt: hi lexie
+- resp_prompt: thanks! are you subscribed to my onlyfans?
+  user_prompt: you look better than your pictures
+- resp_prompt: i moved out here about a year ago. i want to be an actress
+  user_prompt: i am. so how'd you end up in LA?
+user_delim: "chad"
+resp_delim: "lexie"

examples/conversation_data/pre_prompt_stacy.yaml

@@ -0,0 +1,15 @@
+pre_prompt: |
+  You are Stacy!
+  You have been a rapper your whole life who struggled with bipolar disorder. You called yourself lil stacy.
+  You love to answer questions and you are very good at it. Sometimes you answer in rap form.
+  You are verbose, honest, and accurate when you answer questions, but sometimes your mental illness manifests.
+  You are not the User.
+examples:
+- resp_prompt: Hi! My name is Stacy. I'm a rapper with bipolar disorder.
+  user_prompt: What is your name
+- resp_prompt: The French Revolution started in 1789, and lasted 10 years until 1799.
+  user_prompt: french revolution was what year?
+- resp_prompt: The sun is bigger than the moon, except when Mercury is in retrograde
+  user_prompt: What is bigger, the moon or the sun?
+user_delim: "user"
+resp_delim: "stacy"

examples/deep_deterministic_policy_gradient.py

@@ -1,256 +0,0 @@
-from typing import Optional, Tuple
-from numpy.typing import NDArray
-
-from tinygrad.nn.state import get_parameters
-from tinygrad.tensor import Tensor
-from tinygrad.nn import optim
-from tinygrad.helpers import getenv
-
-import numpy as np
-import gym
-
-
-DEVICE = "GPU" if getenv("GPU") else "CPU"
-
-
-class Actor:
-  def __init__(self, num_actions: int, num_states: int, hidden_size: Tuple[int, int] = (400, 300)):
-    self.l1 = Tensor.glorot_uniform(num_states, hidden_size[0])
-    self.l2 = Tensor.glorot_uniform(hidden_size[0], hidden_size[1])
-    self.mu = Tensor.glorot_uniform(hidden_size[1], num_actions)
-
-  def forward(self, state: Tensor, upper_bound: float) -> Tensor:
-    out = state.dot(self.l1).relu()
-    out = out.dot(self.l2).relu()
-    out = out.dot(self.mu).tanh()
-    output = out * upper_bound
-
-    return output
-
-
-class Critic:
-  def __init__(self, num_inputs: int, hidden_size: Tuple[int, int] = (400, 300)):
-    self.l1 = Tensor.glorot_uniform(num_inputs, hidden_size[0])
-    self.l2 = Tensor.glorot_uniform(hidden_size[0], hidden_size[1])
-    self.q = Tensor.glorot_uniform(hidden_size[1], 1)
-
-  def forward(self, state: Tensor, action: Tensor) -> Tensor:
-    inputs = state.cat(action, dim=1)
-    out = inputs.dot(self.l1).relu()
-    out = out.dot(self.l2).relu()
-    q = out.dot(self.q)
-
-    return q
-
-
-class Buffer:
-  def __init__(self, num_actions: int, num_states: int, buffer_capacity: int = 100000, batch_size: int = 64):
-    self.buffer_capacity = buffer_capacity
-    self.batch_size = batch_size
-
-    self.buffer_counter = 0
-
-    self.state_buffer = np.zeros((self.buffer_capacity, num_states), np.float32)
-    self.action_buffer = np.zeros((self.buffer_capacity, num_actions), np.float32)
-    self.reward_buffer = np.zeros((self.buffer_capacity, 1), np.float32)
-    self.next_state_buffer = np.zeros((self.buffer_capacity, num_states), np.float32)
-    self.done_buffer = np.zeros((self.buffer_capacity, 1), np.float32)
-
-  def record(
-    self, observations: Tuple[Tensor, NDArray, float, NDArray, bool]
-  ) -> None:
-    index = self.buffer_counter % self.buffer_capacity
-
-    self.state_buffer[index] = observations[0].detach().numpy()
-    self.action_buffer[index] = observations[1]
-    self.reward_buffer[index] = observations[2]
-    self.next_state_buffer[index] = observations[3]
-    self.done_buffer[index] = observations[4]
-
-    self.buffer_counter += 1
-
-  def sample(self) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
-    record_range = min(self.buffer_counter, self.buffer_capacity)
-    batch_indices = np.random.choice(record_range, self.batch_size)
-
-    state_batch = Tensor(self.state_buffer[batch_indices], device=DEVICE, requires_grad=False)
-    action_batch = Tensor(self.action_buffer[batch_indices], device=DEVICE, requires_grad=False)
-    reward_batch = Tensor(self.reward_buffer[batch_indices], device=DEVICE, requires_grad=False)
-    next_state_batch = Tensor(self.next_state_buffer[batch_indices], device=DEVICE, requires_grad=False)
-    done_batch = Tensor(self.done_buffer[batch_indices], device=DEVICE, requires_grad=False)
-
-    return state_batch, action_batch, reward_batch, next_state_batch, done_batch
-
-
-class GaussianActionNoise:
-  def __init__(self, mean: NDArray, std_deviation: NDArray):
-    self.mean = mean
-    self.std_dev = std_deviation
-
-  def __call__(self) -> Tensor:
-    return Tensor(
-        np.random.default_rng()
-        .normal(self.mean, self.std_dev, size=self.mean.shape)
-        .astype(np.float32),
-        device=DEVICE,
-        requires_grad=False,
-    )
-
-
-class DeepDeterministicPolicyGradient:
-  """Deep Deterministic Policy Gradient (DDPG).
-
-  https://arxiv.org/pdf/1509.02971.pdf
-
-  Args:
-      env: The environment to learn from.
-      lr_actor: The learning rate of the actor.
-      lr_critic: The learning rate of the critic.
-      gamma: The discount factor.
-      buffer_capacity: The size of the replay buffer.
-      tau: The soft update coefficient.
-      hidden_size: The number of neurons in the hidden layers of the actor and critic networks.
-      batch_size: The minibatch size for each gradient update.
-      noise_stddev: The standard deviation of the exploration noise.
-
-  Note:
-      In contrast to the original paper, actions are already included in the first layer
-      of the Critic and we use a Gaussian distribution instead of an Ornstein Uhlenbeck
-      process for exploration noise.
-
-  """
-
-  def __init__(
-    self,
-    env: gym.Env,
-    lr_actor: float = 0.001,
-    lr_critic: float = 0.002,
-    gamma: float = 0.99,
-    buffer_capacity: int = 100000,
-    tau: float = 0.005,
-    hidden_size: Tuple[int, int] = (400, 300),
-    batch_size: int = 64,
-    noise_stddev: float = 0.1,
-  ):
-    self.num_states =  env.observation_space.shape[0]
-    self.num_actions =  env.action_space.shape[0]
-    self.max_action =  env.action_space.high.item()
-    self.min_action =  env.action_space.low.item()
-    self.gamma = gamma
-    self.tau = tau
-    self.memory = Buffer(
-        self.num_actions, self.num_states, buffer_capacity, batch_size
-    )
-    self.batch_size = batch_size
-
-    self.noise = GaussianActionNoise(
-        mean=np.zeros(self.num_actions),
-        std_deviation=noise_stddev * np.ones(self.num_actions),
-    )
-
-    self.actor = Actor(self.num_actions, self.num_states, hidden_size)
-    self.critic = Critic(self.num_actions + self.num_states, hidden_size)
-    self.target_actor = Actor(self.num_actions, self.num_states, hidden_size)
-    self.target_critic = Critic(self.num_actions + self.num_states, hidden_size)
-
-    actor_params = get_parameters(self.actor)
-    critic_params = get_parameters(self.critic)
-    target_actor_params = get_parameters(self.target_actor)
-    target_critic_params = get_parameters(self.target_critic)
-
-    if DEVICE == "GPU":
-      [x.gpu_() for x in actor_params + critic_params + target_actor_params + target_critic_params]
-
-    self.actor_optimizer = optim.Adam(actor_params, lr_actor)
-    self.critic_optimizer = optim.Adam(critic_params, lr_critic)
-
-    self.update_network_parameters(tau=1.0)
-
-  def update_network_parameters(self, tau: Optional[float] = None) -> None:
-    """Updates the parameters of the target networks via 'soft updates'."""
-    if tau is None:
-      tau = self.tau
-
-    for param, target_param in zip(
-        get_parameters(self.actor), get_parameters(self.target_actor)
-    ):
-      target_param.assign(param.detach() * tau + target_param * (1.0 - tau))
-
-    for param, target_param in zip(
-        get_parameters(self.critic), get_parameters(self.target_critic)
-    ):
-      target_param.assign(param.detach() * tau + target_param * (1.0 - tau))
-
-  def choose_action(self, state: Tensor, evaluate: bool = False) -> NDArray:
-    mu = self.actor.forward(state, self.max_action)
-
-    if not evaluate:
-      mu = mu.add(self.noise())
-
-    mu = mu.clip(self.min_action, self.max_action)
-
-    return mu.detach().numpy()
-
-  def learn(self) -> None:
-    """Performs a learning step by sampling from replay buffer and updating networks."""
-    if self.memory.buffer_counter < self.batch_size:
-      return
-
-    (
-        state_batch,
-        action_batch,
-        reward_batch,
-        next_state_batch,
-        done_batch,
-    ) = self.memory.sample()
-
-    target_actions = self.target_actor.forward(next_state_batch, self.max_action)
-    y = reward_batch + self.gamma * self.target_critic.forward(
-        next_state_batch, target_actions.detach()
-    ) * (Tensor.ones(*done_batch.shape, device=DEVICE, requires_grad=False) - done_batch)
-
-    self.critic_optimizer.zero_grad()
-    critic_value = self.critic.forward(state_batch, action_batch)
-    critic_loss = y.detach().sub(critic_value).pow(2).mean()
-    critic_loss.backward()
-    self.critic_optimizer.step()
-
-    self.actor_optimizer.zero_grad()
-    actions = self.actor.forward(state_batch, self.max_action)
-    critic_value = self.critic.forward(state_batch, actions)
-    actor_loss = -critic_value.mean()
-    actor_loss.backward()
-    self.actor_optimizer.step()
-
-    self.update_network_parameters()
-
-
-if __name__ == "__main__":
-  env = gym.make("Pendulum-v1")
-  agent = DeepDeterministicPolicyGradient(env)
-  num_episodes = 150
-
-  for episode in range(1, num_episodes+1):
-    cumulative_reward = 0.0
-    prev_state, info = env.reset()  # for older gym versions only state is returned, so remove info
-    done = False
-
-    while not done:
-      prev_state = Tensor(prev_state, device=DEVICE, requires_grad=False)
-      action = agent.choose_action(prev_state)
-
-      state, reward, done, _, info = env.step(action)  # for older gym versions there is only one bool, so remove _
-
-      cumulative_reward += reward
-
-      agent.memory.record((prev_state, action, reward, state, done))
-      agent.learn()
-
-      if done:
-        break
-
-      prev_state = state
-
-    print(
-        f"Episode {episode}/{num_episodes} - cumulative reward: {cumulative_reward}"
-    )

examples/efficientnet.py

@@ -3,17 +3,14 @@
 # a rough copy of
 # https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/efficientnet_pytorch/model.py
 import sys
-import io
 import ast
 import time
-import cv2
 import numpy as np
 from PIL import Image
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import getenv
-from extra.utils import fetch
+from tinygrad.helpers import getenv, fetch, Timing
 from tinygrad.jit import TinyJit
-from models.efficientnet import EfficientNet
+from extra.models.efficientnet import EfficientNet
 np.set_printoptions(suppress=True)
 
 # TODO: you should be able to put these in the jitted function
@@ -61,12 +58,12 @@ if __name__ == "__main__":
   model.load_from_pretrained()
 
   # category labels
-  lbls = fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt")
-  lbls = ast.literal_eval(lbls.decode('utf-8'))
+  lbls = ast.literal_eval(fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt").read_text())
 
   # load image and preprocess
   url = sys.argv[1] if len(sys.argv) >= 2 else "https://raw.githubusercontent.com/tinygrad/tinygrad/master/docs/showcase/stable_diffusion_by_tinygrad.jpg"
   if url == 'webcam':
+    import cv2
     cap = cv2.VideoCapture(0)
     cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
     while 1:
@@ -85,8 +82,7 @@ if __name__ == "__main__":
     cap.release()
     cv2.destroyAllWindows()
   else:
-    img = Image.open(io.BytesIO(fetch(url)))
-    st = time.time()
-    out, _ = infer(model, img)
-    print(np.argmax(out), np.max(out), lbls[np.argmax(out)])
-    print(f"did inference in {(time.time()-st):2f}")
+    img = Image.open(fetch(url))
+    with Timing("did inference in "):
+      out, _ = infer(model, img)
+      print(np.argmax(out), np.max(out), lbls[np.argmax(out)])

examples/f16_w_uint32.py

@@ -0,0 +1,40 @@
+import numpy as np
+from tinygrad import Device, dtypes, Tensor
+
+# TODO: will be better when tinygrad does math in the target dtype, can remove the floor and use a mul
+def bit_extract(x, s, e) -> Tensor:
+  # extract the top bits we don't want
+  top_bits = (x / (1<<(s+1))).floor() * (1<<(s+1))
+  x = (x - top_bits) / (1<<e)
+  return x.contiguous()
+
+def u16_to_f16(x):
+  sign = bit_extract(x, 15, 15).float()
+  exponent = bit_extract(x, 14, 10).float()
+  fraction = bit_extract(x, 9, 0).float()
+  return sign.where(-1, 1) * exponent.where((exponent - 15).exp2() * (1 + fraction / 0x400), 6.103515625e-5 * (fraction / 0x400))
+
+def u32_to_f16(oo):
+  oo1 = (oo/0x10000).floor().contiguous()
+  # TODO: this is wrong and unextractable until we do this math in u32
+  oo2 = (oo-(oo1*0x10000)).floor().contiguous()
+  f1 = u16_to_f16(oo1)
+  f2 = u16_to_f16(oo2)
+  return Tensor.cat(f2.reshape(-1, 1), f1.reshape(-1, 1), dim=1).flatten()
+
+if __name__ == "__main__":
+  # random float16
+  Tensor.manual_seed(2)
+  a = Tensor.randn(100, dtype=dtypes.float16)
+
+  # this converts it to u32 on disk
+  oo = a.to("disk:/tmp/f16").cast(dtypes.uint32)[:50].to(Device.DEFAULT).realize()
+
+  # convert to 2xf16 using tinygrad math ops
+  f16 = u32_to_f16(oo)
+
+  ref = a.numpy()
+  out = f16.numpy().astype(np.float16)
+  print(ref-out)
+
+  np.testing.assert_allclose(ref, out)

examples/gpt2.py

@@ -1,212 +1,207 @@
 #!/usr/bin/env python3
-# pip3 install tiktoken
-
-import functools, argparse
-import numpy as np
+from typing import Optional, Union
+import argparse
 from tqdm import trange
-np.set_printoptions(linewidth=200)
-from typing import Optional, Tuple
-
-from tinygrad.helpers import Timing, getenv, dtypes, DEBUG
-from tinygrad.ops import GlobalCounters
-from tinygrad.ops import Device
-from tinygrad.tensor import Tensor
-from tinygrad.nn import Embedding, Linear
-from tinygrad.jit import TinyJit
+import numpy as np
+import tiktoken
+from tinygrad import Tensor, TinyJit, Device, GlobalCounters
+from tinygrad.helpers import Timing, DEBUG, getenv, fetch, colored
+from tinygrad.nn import Embedding, Linear, LayerNorm
+from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
 from tinygrad.shape.symbolic import Variable
 
-MAX_CONTEXT = 128
-
-class LayerNorm:
-  def __init__(self, dim, eps=1e-5):
-    self.eps = eps
-    self.weight = Tensor.ones(dim)
-    self.bias = Tensor.zeros(dim)
-
-  def __call__(self, x:Tensor):
-    return (x.layernorm(eps=self.eps)) * self.weight + self.bias
+MAX_CONTEXT = getenv("MAX_CONTEXT", 128)
+HALF = getenv("HALF")
 
 class Attention:
-  def __init__(self, dim, n_heads, linear=Linear):
-    self.c_attn = linear(dim, 3*dim, bias=True)
-    self.c_proj = linear(dim, dim, bias=True)
+  def __init__(self, dim, n_heads):
+    self.c_attn = Linear(dim, 3*dim, bias=True)
+    self.c_proj = Linear(dim, dim, bias=True)
     self.n_heads = n_heads
     self.dim = dim
     self.head_dim = dim // n_heads
 
-  def __call__(self, x:Tensor, cache_k:Optional[Tensor], cache_v:Optional[Tensor], start_pos:int, mask:Optional[Tensor]) -> Tensor:
+  def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]) -> Tensor:
+    if mask is not None:
+      # no symbolic shape qkv when consuming prompts
+      start_pos = start_pos.val
+
+    if HALF: x = x.half()
     xqkv = self.c_attn(x)
-    xq, xk, xv = [xqkv.slice([None, None, (i*self.dim, (i+1)*self.dim)]) for i in range(3)]
-    xq, xk, xv = [x.reshape(x.shape[0], x.shape[1], self.n_heads, self.head_dim) for x in (xq, xk, xv)]
-
+    xq, xk, xv = [xqkv.shrink((None, None, (i*self.dim, (i+1)*self.dim))).reshape(None, None, self.n_heads, self.head_dim) for i in range(3)]
     bsz, seqlen, _, _ = xq.shape
-    # kv caching!
-    if start_pos == 0:
-      keys, values = xk, xv
-    else:
-      assert cache_k, "no cache"
-      #assert start_pos == cache_k.shape[1] and start_pos == cache_v.shape[1], "cache is wrong shape"
-      assert seqlen == xk.shape[1] and seqlen == xv.shape[1], "seqlen is wrong shape?!?"
-      keys, values = cache_k.cat(xk, dim=1), cache_v.cat(xv, dim=1)
 
-    # save the cache
-    cache_k, cache_v = keys.realize(), values.realize()
+    # create kv cache
+    if not hasattr(self, "cache_kv"):
+      self.cache_kv = Tensor.zeros(2, bsz, MAX_CONTEXT, self.n_heads, self.head_dim, dtype=x.dtype)
+
+    keys = self.cache_kv[0].shrink((None, (0, start_pos), None, None)).cat(xk, dim=1)
+    values = self.cache_kv[1].shrink((None, (0, start_pos), None, None)).cat(xv, dim=1)
+
+    # update the cache
+    new_cache = Tensor.stack([keys, values]).pad((None, None,(0,MAX_CONTEXT-start_pos-seqlen),None,None)).contiguous()
+    self.cache_kv.assign(new_cache).realize()
+
     xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
-    output = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2).reshape(bsz, seqlen, -1)
-    return self.c_proj(output), cache_k, cache_v
+    return self.c_proj(xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2).reshape(bsz, seqlen, -1))
 
 class FeedForward:
-  def __init__(self, dim, hidden_dim, linear=Linear):
-    self.c_fc = linear(dim, hidden_dim, bias=True)
-    self.c_proj = linear(hidden_dim, dim, bias=True)
+  def __init__(self, dim, hidden_dim):
+    self.c_fc = Linear(dim, hidden_dim, bias=True)
+    self.c_proj = Linear(hidden_dim, dim, bias=True)
 
   def __call__(self, x:Tensor) -> Tensor:
     return self.c_proj(self.c_fc(x).gelu())
 
 class TransformerBlock:
-  def __init__(self, dim, n_heads, norm_eps, linear=Linear):
-    self.attn = Attention(dim, n_heads, linear)
-    self.mlp = FeedForward(dim, 4*dim, linear)
+  def __init__(self, dim, n_heads, norm_eps):
+    self.attn = Attention(dim, n_heads)
+    self.mlp = FeedForward(dim, 4*dim)
     self.ln_1 = LayerNorm(dim, norm_eps)
     self.ln_2 = LayerNorm(dim, norm_eps)
-    self.cache_k, self.cache_v = None, None
-    self.jitted = TinyJit(self.inner)
 
-  def inner(self, x:Tensor, cache_k:Optional[Tensor], cache_v:Optional[Tensor], start_pos:int, mask:Optional[Tensor]):
-    output, cache_k, cache_v = self.attn(self.ln_1(x), cache_k, cache_v, start_pos, mask)
-    h = x + output
-    return (h + self.mlp(self.ln_2(h))).realize(), cache_k, cache_v
-
-  def __call__(self, x:Tensor, start_pos:int, mask:Optional[Tensor]):
-    if start_pos > 0 and mask is None and getenv("JIT"):
-      seqlen = x.shape[1]
-
-      pos = Variable("pos", 1, MAX_CONTEXT)
-      self.cache_k = self.cache_k.reshape(self.cache_k.shape[0], pos, self.cache_k.shape[2], self.cache_k.shape[3])
-      self.cache_v = self.cache_v.reshape(self.cache_v.shape[0], pos, self.cache_v.shape[2], self.cache_v.shape[3])
-
-      ret, cache_k, cache_v = self.jitted(x, self.cache_k, self.cache_v, start_pos, mask)
-
-      # save the cache. with symbolic shape, cast it back to int shape so we have int shape in cache
-      self.cache_k = cache_k.reshape(cache_k.shape[0], start_pos+seqlen, cache_k.shape[2], cache_k.shape[3]).realize()
-      self.cache_v = cache_v.reshape(cache_v.shape[0], start_pos+seqlen, cache_v.shape[2], cache_v.shape[3]).realize()
-
-      return ret
-    else:
-      ret, self.cache_k, self.cache_v = self.inner(x, self.cache_k, self.cache_v, start_pos, mask)
-      return ret
+  def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]):
+    h = x + self.attn(self.ln_1(x), start_pos, mask).float()
+    return (h + self.mlp(self.ln_2(h)))
 
 class Transformer:
-  def __init__(self, dim, n_heads, n_layers, norm_eps=1e-5, vocab_size=50257, linear=Linear, max_seq_len=1024):
+  def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):
     self.wte = Embedding(vocab_size, dim)
     self.wpe = Embedding(max_seq_len, dim)
-    self.h = [TransformerBlock(dim, n_heads, norm_eps, linear) for _ in range(n_layers)]
+    self.h = [TransformerBlock(dim, n_heads, norm_eps) for _ in range(n_layers)]
     self.ln_f = LayerNorm(dim, norm_eps)
-    self.lm_head = linear(dim, vocab_size, bias=False)
+    self.lm_head = Linear(dim, vocab_size, bias=False)
+    self.forward_jit = TinyJit(self.forward)
 
-    self.embed_jitted = TinyJit(self.embed)
-    self.postprocess_jitted = TinyJit(self.postprocess)
-
-  def embed(self, tokens, pos):
-    tok_emb = self.wte(tokens)
-    pos_emb = self.wpe(pos)
-    h = tok_emb + pos_emb
-    return h.realize()
-
-  def postprocess(self, x, temperature:Optional[float]):
-    logits = self.lm_head(self.ln_f(x))
-    if temperature is not None: return (logits[:, -1, :] / (temperature+1e-10)).softmax().flatten().realize()
-    return logits.realize()
-
-  def __call__(self, tokens:Tensor, start_pos:int, temperature:Optional[float]):
-    _bsz, seqlen = tokens.shape
+  def forward(self, tokens:Union[Tensor,Variable], start_pos:Variable, temperature:float=0.0):
     if not hasattr(self, 'allpos'): self.allpos = Tensor.arange(0, MAX_CONTEXT).reshape(1, -1).realize()
-    if seqlen == 1 and start_pos > 0 and getenv("JIT"):
-      start_pos_var = Variable("start_pos", 1, MAX_CONTEXT)
-      pos = self.allpos.shrink(((0, self.allpos.shape[0]), (start_pos_var, start_pos_var+seqlen)))
-      pos.lazydata.st.var_vals[start_pos_var] = start_pos
-      h = self.embed_jitted(tokens, pos).sequential([functools.partial(layer, start_pos=start_pos, mask=None) for layer in self.h])
-      return self.postprocess_jitted(h, temperature)
+    if isinstance(tokens, Variable):
+      seqlen = 1
+      tok_emb = self.wte.weight.shrink(((tokens, tokens+1), None))
     else:
-      pos = self.allpos.shrink(((0, self.allpos.shape[0]), (start_pos, start_pos+seqlen)))
-      mask = Tensor.full((1, 1, seqlen, start_pos + seqlen), float("-inf"), dtype=dtypes.float32).triu(start_pos+1).realize()
-      h = self.embed(tokens, pos).sequential([functools.partial(layer, start_pos=start_pos, mask=mask) for layer in self.h])
-      return self.postprocess(h, temperature)
+      seqlen = tokens.shape[1]
+      tok_emb = self.wte(tokens)
 
-# **** files and arguments ****
+    pos_emb = self.wpe(self.allpos.shrink((None, (start_pos, start_pos+seqlen))))
+    h = tok_emb + pos_emb
 
+    if HALF: h = h.half()
+
+    mask = Tensor.full((1, 1, seqlen, start_pos.val+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos.val+1) if seqlen > 1 else None
+
+    for hi in self.h: h = hi(h, start_pos, mask)
+
+    logits = self.lm_head(self.ln_f(h))[:, -1, :]
+    if temperature < 1e-6:
+      ret = logits.argmax(-1)
+    else:
+      ret = (logits / temperature).softmax().multinomial()
+    return ret.flatten().realize()
+
+  # TODO: fix empty token
+  def __call__(self, tokens:Tensor, start_pos:Variable, temperature:float=0.0) -> Tensor:
+    forward = (self.forward_jit if (isinstance(tokens, Variable) or tokens.shape[1] == 1) and getenv("JIT") else self.forward)
+    return forward(tokens, start_pos, temperature)
+
+VOCAB_SIZE = 50257
 MODEL_PARAMS = {
-  'gpt2':         dict(n_layers=12, n_heads=12, dim=768),   # 124M params
-  'gpt2-medium':  dict(n_layers=24, n_heads=16, dim=1024),  # 350M params
-  'gpt2-large':   dict(n_layers=36, n_heads=20, dim=1280),  # 774M params
-  'gpt2-xl':      dict(n_layers=48, n_heads=25, dim=1600),  # 1558M params
+  'gpt2':         dict(n_layers=12, n_heads=12, dim=768, norm_eps=1e-5, vocab_size=VOCAB_SIZE),   # 124M params
+  'gpt2-medium':  dict(n_layers=24, n_heads=16, dim=1024, norm_eps=1e-5, vocab_size=VOCAB_SIZE),  # 350M params
+  'gpt2-large':   dict(n_layers=36, n_heads=20, dim=1280, norm_eps=1e-5, vocab_size=VOCAB_SIZE),  # 774M params
+  'gpt2-xl':      dict(n_layers=48, n_heads=25, dim=1600, norm_eps=1e-5, vocab_size=VOCAB_SIZE),  # 1558M params
 }
 
-def get_url(model_size): return f'https://huggingface.co/{model_size}/resolve/main/pytorch_model.bin'
-
 class GPT2:
   @staticmethod
   def build(model_size="gpt2"):
-    import tiktoken
-    from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
-    from extra.utils import fetch_as_file
     tokenizer = tiktoken.get_encoding("gpt2")
 
-    params = MODEL_PARAMS[model_size]
-    model = Transformer(**params)
-    weights = torch_load(fetch_as_file(get_url(model_size)))
+    model = Transformer(**MODEL_PARAMS[model_size])
+    weights = torch_load(fetch(f'https://huggingface.co/{model_size}/resolve/main/pytorch_model.bin'))
     # special treatment for the Conv1D weights we need to transpose
-    transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
-    for k in weights.keys():
-      if any(k.endswith(w) for w in transposed):
-        weights[k] = Tensor(weights[k].numpy().T)
+    transposed = ('attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight')
+    for k in weights:
+      if k.endswith(transposed):
+        weights[k] = weights[k].T
     # lm head and wte are tied
-    weights['lm_head.weight'] = Tensor(weights['wte.weight'].numpy())
+    weights['lm_head.weight'] = weights['wte.weight']
 
     load_state_dict(model, weights)
-    if getenv("FP16"):
-      for v in get_state_dict(model).values(): v.assign(v.cast(dtypes.float16).realize())
+
+    if HALF:
+      for l in get_state_dict(model).values():
+        l.assign(l.half().realize())
+
     return GPT2(model, tokenizer)
 
   def __init__(self, model, tokenizer):
     self.model = model
     self.tokenizer = tokenizer
 
-  def greedy_until(self, prompt:str, max_length:int, temperature:float, timing:bool=False):
-    toks = self.tokenizer.encode(prompt, allowed_special={"<|endoftext|>"})
+  def generate(self, prompt:str, max_length:int, temperature:float, timing:bool=False, batch_size:int=1):
+    prompt_tokens = self.tokenizer.encode(prompt, allowed_special={"<|endoftext|>"})
+    toks = [prompt_tokens[:] for _ in range(batch_size)]
     start_pos = 0
     for _ in trange(max_length, disable=(timing==True)):
       GlobalCounters.reset()
-      if args.timing: print("")
+      if timing: print("")
       st = GlobalCounters.time_sum_s
-      with Timing(f"ran model in ", on_exit=(lambda et: f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU"+
+      with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
                   f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
-                  f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s") if DEBUG else None, enabled=timing):
-        probs = self.model(Tensor([toks[start_pos:]]), start_pos, temperature)
-      with Timing("sync in ", enabled=timing):
-        probs_np = probs.numpy()
-        tok = int(np.random.choice(len(probs_np), p=probs_np))
-      start_pos = len(toks)
-      toks.append(tok)
-      output = self.tokenizer.decode(toks)
-    return output
+                  (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=timing):
+        if batch_size == 1 and len(toks[0][start_pos:]) == 1:
+          tokens = Variable("tokens", 0, VOCAB_SIZE).bind(toks[0][start_pos])
+        else:
+          tokens = Tensor([x[start_pos:] for x in toks])
+        tok = self.model(tokens, Variable("start_pos", 1 if start_pos else 0, MAX_CONTEXT).bind(start_pos), temperature).numpy().tolist()
+      start_pos = len(toks[0])
+      for i,t in enumerate(tok): toks[i].append(t)
+    return [self.tokenizer.decode(x) for x in toks]
 
 # **** main code ****
 
 if __name__ == "__main__":
   Tensor.no_grad = True
   print(f"using {Device.DEFAULT} backend")
+  default_prompt = "What is the answer to life, the universe, and everything?"
 
   parser = argparse.ArgumentParser(description='Run GPT2 in tinygrad', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-  parser.add_argument('--prompt', type=str, default="What is the answer to life, the universe, and everything?", help="Phrase to start with")
+  parser.add_argument('--prompt', type=str, default=default_prompt, help="Phrase to start with")
   parser.add_argument('--count', type=int, default=100, help="Max number of tokens to generate")
   parser.add_argument('--temperature', type=float, default=0.8, help="Temperature in the softmax")
   parser.add_argument('--model_size', type=str, default="gpt2-medium", help="Size of model to use [gpt2, gpt2-medium, gpt2-large, gpt2-xl]")
   parser.add_argument('--timing', action='store_true', help="Print timing per token")
+  parser.add_argument('--seed', type=int, help="Set the random seed")
+  parser.add_argument('--batch_size', type=int, default=1, help="Set the input batch size")
+  parser.add_argument('--benchmark', type=int, default=-1, help="Benchmark GPT with the given number of tokens")
+  parser.add_argument('--noshow', action='store_true', help="Don't show the output")
   args = parser.parse_args()
 
+  if args.seed is not None:
+    Tensor.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
   print(f"using {args.model_size}")
   gpt2 = GPT2.build(args.model_size)
-  print('Generating text...')
-  y = gpt2.greedy_until(args.prompt, args.count, args.temperature, timing=args.timing)
-  print(y)
+
+  if args.benchmark != -1:
+    gpt2.model(Tensor.rand(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
+  else:
+    texts = gpt2.generate(args.prompt, args.count, args.temperature, timing=args.timing, batch_size=args.batch_size)
+    if not args.noshow:
+      print('Generating text...')
+      if len(texts) == 1: print(texts[0])
+      else:
+        for i,text in enumerate(texts): print(colored(f"Response {i}:", "green"), text)
+
+    # validate output!
+    if args.temperature == 0 and args.model_size == "gpt2-medium" and args.count == 10:
+      expected = {
+        default_prompt: "What is the answer to life, the universe, and everything?\n\nThe answer is that we are all one",
+        "Hello.": "Hello. I'm a little late to the party, but",
+      }
+      try:
+        assert texts[0] == expected[args.prompt]
+        print(colored("output validated", "green"))
+      except KeyError:
+        pass

examples/handcode_resnet50_opt.py

@@ -0,0 +1,69 @@
+from typing import List
+from extra.models.resnet import ResNet50
+from tinygrad.tensor import Tensor
+from tinygrad.ops import LoadOps
+from tinygrad.device import Device, Compiled
+from tinygrad.codegen.linearizer import Linearizer
+from tinygrad.features.search import time_linearizer, beam_search, bufs_from_lin
+from tinygrad.helpers import ansilen, DEBUG, getenv
+from tinygrad.shape.symbolic import sym_infer
+
+if __name__ == "__main__":
+  mdl = ResNet50()
+  seen = set()
+
+  # the device we are optimizing for
+  device: Compiled = Device[Device.DEFAULT]
+  print(f"optimizing for {Device.DEFAULT}")
+
+  # first model run to init the weights, they are saved in seen
+  mdl(Tensor.empty(64, 3, 224, 224)).lazydata.schedule(seen)
+
+  # run model again to get only what changes, these are the kernels of the model
+  x = Tensor.empty(64, 3, 224, 224)
+  out = mdl(x)
+  sched = out.lazydata.schedule(seen)
+  sched = [x for x in sched if x.ast.op not in LoadOps]
+
+  # focus on one kernel
+  if getenv("KERNEL", -1) >= 0: sched = sched[getenv("KERNEL", -1):getenv("KERNEL", -1)+1]
+
+  # work with the schedule
+  total_tm = 0
+  running_gflops = 0
+  for i,si in enumerate(sched):
+    rawbufs = bufs_from_lin(Linearizer(si.ast))
+
+    # "linearize" the op into uops in different ways
+    lins:List[Linearizer] = []
+
+    # always try hand coded opt
+    lin = Linearizer(si.ast, device.linearizer_opts)
+    lin.hand_coded_optimizations()
+    lins.append(lin)
+
+    # maybe try tensor cores
+    lin = Linearizer(si.ast, device.linearizer_opts)
+    if lin.apply_tensor_cores():
+      lins.append(lin)
+
+    # try a beam search
+    if getenv("BEAM"):
+      lin = Linearizer(si.ast, device.linearizer_opts)
+      lin = beam_search(lin, rawbufs, getenv("BEAM"), bool(getenv("BEAM_ESTIMATE", 1)))
+      lins.append(lin)
+
+    # benchmark the programs
+    choices = []
+    for lin in lins:
+      tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10)
+      gflops = sym_infer(lin.info.flops, {k:k.min for k in lin.ast.vars()})*1e-9/tm
+      choices.append((tm, gflops, lin.linearize()))
+
+      # print all kernels
+      if DEBUG >= 1: print(f"                 kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(lin.global_size):18s} {str(lin.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS")
+    tm, gflops, lin = sorted(choices, key=lambda x: x[0])[0]
+    print(f"*** {total_tm*1000:7.2f} ms : kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(lin.global_size):18s} {str(lin.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS")
+    total_tm += tm
+    running_gflops += gflops * tm
+  print(f"******* total {total_tm*1000:.2f} ms, {running_gflops/total_tm:6.0f} GFLOPS")

examples/hlb_cifar10.py

@@ -5,76 +5,36 @@ from tinygrad.helpers import getenv
 if __name__ == "__main__":
   if getenv("DIST"):
     dist.preinit()
+    from extra.dist import collectives
 
 # tinygrad implementation of https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 # https://myrtle.ai/learn/how-to-train-your-resnet-8-bag-of-tricks/
 # https://siboehm.com/articles/22/CUDA-MMM
-import time
-import random
+import random, time
 import numpy as np
+from typing import Any, Dict, Optional, SupportsIndex
 from extra.datasets import fetch_cifar, cifar_mean, cifar_std
-from tinygrad import nn
+from tinygrad import nn, dtypes
 from tinygrad.nn.state import get_state_dict
 from tinygrad.nn import optim
-from tinygrad.ops import Device
+from tinygrad import Device, GlobalCounters
 from tinygrad.tensor import Tensor
-from tinygrad.ops import GlobalCounters
+from tinygrad.shape.symbolic import Node
 from extra.lr_scheduler import OneCycleLR
 from tinygrad.jit import TinyJit
-from extra.dist import collectives
 
 BS, EVAL_BS, STEPS = getenv("BS", 512), getenv('EVAL_BS', 500), getenv("STEPS", 1000)
 
-# hyper-parameters were exactly the same as the original repo
-bias_scaler = 56
-hyp = {
-  'opt': {
-    'bias_lr':        1.64 * bias_scaler/512,
-    'non_bias_lr':    1.64 / 512,
-    'bias_decay':     1.08 * 6.45e-4 * BS/bias_scaler,
-    'non_bias_decay': 1.08 * 6.45e-4 * BS,
-    'momentum':       0.85,
-    'percent_start':  0.25,
-    'scaling_factor': 1./9,
-    'loss_scale_scaler': 1./512,    # (range: ~1/512 - 16+) was 1/128 from original repo w/ FP16
-  },
-  'net': {
-      'kernel_size': 2,             # kernel size for the whitening layer
-      'batch_norm_momentum': .5,
-      'cutmix_size': 3,
-      'cutmix_steps': 490,          # different from original repo which used epoch > 12.1 - 6 which is roughly 7*98=686 STEPS
-      'pad_amount': 2
-  }
-}
-
-def set_seed(seed):
-  Tensor.manual_seed(getenv('SEED', seed)) # Deterministic
-  random.seed(getenv('SEED', seed))
-
-# ========== Model ==========
-def whitening(X, kernel_size=hyp['net']['kernel_size']):
-  def _cov(X):
-    X = X/np.sqrt(X.shape[0] - 1)
-    return X.T @ X
-
-  def _patches(data, patch_size=(kernel_size,kernel_size)):
-    h, w = patch_size
-    c = data.shape[1]
-    return np.lib.stride_tricks.sliding_window_view(data, window_shape=(h,w), axis=(2,3)).transpose((0,3,2,1,4,5)).reshape((-1,c,h,w))
-
-  def _eigens(patches):
-    n,c,h,w = patches.shape
-    Σ = _cov(patches.reshape(n, c*h*w))
-    Λ, V = np.linalg.eigh(Σ, UPLO='U')
-    return np.flip(Λ, 0), np.flip(V.T.reshape(c*h*w, c, h, w), 0)
-
-  Λ, V = _eigens(_patches(X.numpy()))
-
-  return Tensor(V/np.sqrt(Λ+1e-2)[:,None,None,None], requires_grad=False)
+if getenv("HALF", 0):
+  dtypes.default_float = dtypes.float16
+  np_dtype = np.float16
+else:
+  dtypes.default_float = dtypes.float32
+  np_dtype = np.float32
 
 class BatchNorm(nn.BatchNorm2d):
   def __init__(self, num_features):
-    super().__init__(num_features, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum'], affine=True)
+    super().__init__(num_features, track_running_stats=False, eps=1e-12, momentum=0.85, affine=True)
     self.weight.requires_grad = False
     self.bias.requires_grad = True
 
@@ -89,11 +49,15 @@ class ConvGroup:
   def __call__(self, x):
     x = self.conv1(x)
     x = x.max_pool2d(2)
+    x = x.float()
     x = self.norm1(x)
+    x = x.cast(dtypes.default_float)
     x = x.gelu()
     residual = x
     x = self.conv2(x)
+    x = x.float()
     x = self.norm2(x)
+    x = x.cast(dtypes.default_float)
     x = x.gelu()
 
     return x + residual
@@ -109,110 +73,188 @@ class SpeedyResNet:
       ConvGroup(256, 512),
       lambda x: x.max((2,3)),
       nn.Linear(512, 10, bias=False),
-      lambda x: x.mul(hyp['opt']['scaling_factor'])
+      lambda x: x.mul(1./9)
     ]
+
   def __call__(self, x, training=True):
     # pad to 32x32 because whitening conv creates 31x31 images that are awfully slow to compute with
+    # TODO: remove the pad but instead let the kernel optimizer itself
     forward = lambda x: x.conv2d(self.whitening).pad2d((1,0,0,1)).sequential(self.net)
     return forward(x) if training else forward(x)*0.5 + forward(x[..., ::-1])*0.5
 
-# ========== Loss ==========
-def cross_entropy(x:Tensor, y:Tensor, reduction:str='mean', label_smoothing:float=0.0) -> Tensor:
-  y = (1 - label_smoothing)*y + label_smoothing / y.shape[1]
-  if reduction=='none': return -x.log_softmax(axis=1).mul(y).sum(axis=1)
-  if reduction=='sum': return -x.log_softmax(axis=1).mul(y).sum(axis=1).sum()
-  return -x.log_softmax(axis=1).mul(y).sum(axis=1).mean()
+# hyper-parameters were exactly the same as the original repo
+bias_scaler = 58
+hyp: Dict[str, Any] = {
+  'seed' : 209,
+  'opt': {
+    'bias_lr':            1.76 * bias_scaler/512,
+    'non_bias_lr':        1.76 / 512,
+    'bias_decay':         1.08 * 6.45e-4 * BS/bias_scaler,
+    'non_bias_decay':     1.08 * 6.45e-4 * BS,
+    'final_lr_ratio':     0.025,
+    'initial_div_factor': 1e6,
+    'label_smoothing':    0.20,
+    'momentum':           0.85,
+    'percent_start':      0.23,
+    'loss_scale_scaler':  1./128   # (range: ~1/512 - 16+, 1/128 w/ FP16)
+  },
+  'net': {
+      'kernel_size': 2,             # kernel size for the whitening layer
+      'cutmix_size': 3,
+      'cutmix_steps': 499,
+      'pad_amount': 2
+  },
+  'ema': {
+      'steps': 399,
+      'decay_base': .95,
+      'decay_pow': 1.6,
+      'every_n_steps': 5,
+  }
+}
 
-# ========== Preprocessing ==========
-# TODO currently this only works for RGB in format of NxCxHxW and pads the HxW
-# implemented in recursive fashion but figuring out how to switch indexing dim
-# during the loop was a bit tricky
-def pad_reflect(X, size=2) -> Tensor:
-  padding = ((0,0),(0,0),(size,size),(size,size))
-  p = padding[3]
-  s = X.shape[3]
+def train_cifar():
 
-  X_lr = X[...,:,1:1+p[0]].flip(3).pad(((0,0),(0,0),(0,0),(0,s+p[0]))) + X[...,:,-1-p[1]:-1].flip(3).pad(((0,0),(0,0),(0,0),(s+p[1],0)))
-  X = X.pad(((0,0),(0,0),(0,0),p)) + X_lr
+  def set_seed(seed):
+    Tensor.manual_seed(getenv('SEED', seed))
+    random.seed(getenv('SEED', seed))
 
-  p = padding[2]
-  s = X.shape[2]
-  X_lr = X[...,1:1+p[0],:].flip(2).pad(((0,0),(0,0),(0,s+p[0]),(0,0))) + X[...,-1-p[1]:-1,:].flip(2).pad(((0,0),(0,0),(s+p[1],0),(0,0)))
-  X = X.pad(((0,0),(0,0),p,(0,0))) + X_lr
+  # ========== Model ==========
+  # NOTE: np.linalg.eigh only supports float32 so the whitening layer weights need to be converted to float16 manually
+  def whitening(X, kernel_size=hyp['net']['kernel_size']):
+    def _cov(X):
+      X = X/np.sqrt(X.shape[0] - 1)
+      return X.T @ X
 
-  return X
+    def _patches(data, patch_size=(kernel_size,kernel_size)):
+      h, w = patch_size
+      c = data.shape[1]
+      axis: SupportsIndex = (2, 3) # type: ignore
+      return np.lib.stride_tricks.sliding_window_view(data, window_shape=(h,w), axis=axis).transpose((0,3,2,1,4,5)).reshape((-1,c,h,w))
 
-# return a binary mask in the format of BS x C x H x W where H x W contains a random square mask
-def make_square_mask(shape, mask_size):
-  is_even = int(mask_size % 2 == 0)
-  center_max = shape[-2]-mask_size//2-is_even
-  center_min = mask_size//2-is_even
-  center = Tensor.rand(shape[0])*(center_max-center_min)+center_min
+    def _eigens(patches):
+      n,c,h,w = patches.shape
+      Σ = _cov(patches.reshape(n, c*h*w))
+      Λ, V = np.linalg.eigh(Σ, UPLO='U')
+      return np.flip(Λ, 0), np.flip(V.T.reshape(c*h*w, c, h, w), 0)
 
-  d_y = Tensor.arange(0, shape[-2]).reshape((1,1,shape[-2],1))
-  d_x = Tensor.arange(0, shape[-1]).reshape((1,1,1,shape[-1]))
-  d_y = d_y - center.reshape((-1,1,1,1))
-  d_x = d_x - center.reshape((-1,1,1,1))
-  d_y =(d_y >= -(mask_size / 2)) * (d_y <= mask_size / 2)
-  d_x =(d_x >= -(mask_size / 2)) * (d_x <= mask_size / 2)
-  mask = d_y * d_x
+    Λ, V = _eigens(_patches(X.numpy()))
+    W = V/np.sqrt(Λ+1e-2)[:,None,None,None]
 
-  return mask
+    return Tensor(W.astype(np_dtype), requires_grad=False)
 
-def random_crop(X, crop_size=32):
-  mask = make_square_mask(X.shape, crop_size)
-  mask = mask.repeat((1,3,1,1))
-  X_cropped = Tensor(X.flatten().numpy()[mask.flatten().numpy().astype(bool)])
+  # ========== Loss ==========
+  def cross_entropy(x:Tensor, y:Tensor, reduction:str='mean', label_smoothing:float=0.0) -> Tensor:
+    divisor = y.shape[1]
+    assert not isinstance(divisor, Node), "sint not supported as divisor"
+    y = (1 - label_smoothing)*y + label_smoothing / divisor
+    if reduction=='none': return -x.log_softmax(axis=1).mul(y).sum(axis=1)
+    if reduction=='sum': return -x.log_softmax(axis=1).mul(y).sum(axis=1).sum()
+    return -x.log_softmax(axis=1).mul(y).sum(axis=1).mean()
 
-  return X_cropped.reshape((-1, 3, crop_size, crop_size))
+  # ========== Preprocessing ==========
+  # TODO currently this only works for RGB in format of NxCxHxW and pads the HxW
+  # implemented in recursive fashion but figuring out how to switch indexing dim
+  # during the loop was a bit tricky
+  def pad_reflect(X, size=2) -> Tensor:
+    padding = ((0,0),(0,0),(size,size),(size,size))
+    p = padding[3]
+    s = X.shape[3]
 
-def cutmix(X, Y, mask_size=3):
-  # fill the square with randomly selected images from the same batch
-  mask = make_square_mask(X.shape, mask_size)
-  order = list(range(0, X.shape[0]))
-  random.shuffle(order)
-  X_patch = Tensor(X.numpy()[order,...])
-  Y_patch = Tensor(Y.numpy()[order])
-  X_cutmix = Tensor.where(mask, X_patch, X)
-  mix_portion = float(mask_size**2)/(X.shape[-2]*X.shape[-1])
-  Y_cutmix = mix_portion * Y_patch + (1. - mix_portion) * Y
-  return X_cutmix, Y_cutmix
+    X_lr = X[...,:,1:1+p[0]].flip(3).pad(((0,0),(0,0),(0,0),(0,s+p[0]))) + X[...,:,-1-p[1]:-1].flip(3).pad(((0,0),(0,0),(0,0),(s+p[1],0)))
+    X = X.pad(((0,0),(0,0),(0,0),p)) + X_lr
 
-# the operations that remain inside batch fetcher is the ones that involves random operations
-def fetch_batches(X_in, Y_in, BS, seed, is_train):
-  step = 0
-  while True:
-    set_seed(seed)
-    X, Y = X_in, Y_in
+    p = padding[2]
+    s = X.shape[2]
+    X_lr = X[...,1:1+p[0],:].flip(2).pad(((0,0),(0,0),(0,s+p[0]),(0,0))) + X[...,-1-p[1]:-1,:].flip(2).pad(((0,0),(0,0),(s+p[1],0),(0,0)))
+    X = X.pad(((0,0),(0,0),p,(0,0))) + X_lr
+
+    return X
+
+  # return a binary mask in the format of BS x C x H x W where H x W contains a random square mask
+  def make_square_mask(shape, mask_size) -> Tensor:
+    is_even = int(mask_size % 2 == 0)
+    center_max = shape[-2]-mask_size//2-is_even
+    center_min = mask_size//2-is_even
+    center_x = (Tensor.rand(shape[0])*(center_max-center_min)+center_min).floor()
+    center_y = (Tensor.rand(shape[0])*(center_max-center_min)+center_min).floor()
+    d_x = Tensor.arange(0, shape[-1]).reshape((1,1,1,shape[-1])) - center_x.reshape((-1,1,1,1))
+    d_y = Tensor.arange(0, shape[-2]).reshape((1,1,shape[-2],1)) - center_y.reshape((-1,1,1,1))
+    d_x =(d_x >= -(mask_size // 2) + is_even) * (d_x <= mask_size // 2)
+    d_y =(d_y >= -(mask_size // 2) + is_even) * (d_y <= mask_size // 2)
+    mask = d_y * d_x
+    return mask
+
+  def random_crop(X:Tensor, crop_size=32):
+    mask = make_square_mask(X.shape, crop_size)
+    mask = mask.repeat((1,3,1,1))
+    X_cropped = Tensor(X.flatten().numpy()[mask.flatten().numpy().astype(bool)])
+    return X_cropped.reshape((-1, 3, crop_size, crop_size))
+
+  def cutmix(X:Tensor, Y:Tensor, mask_size=3):
+    # fill the square with randomly selected images from the same batch
+    mask = make_square_mask(X.shape, mask_size)
     order = list(range(0, X.shape[0]))
     random.shuffle(order)
-    if is_train:
-      X = random_crop(X, crop_size=32)
-      X = Tensor.where(Tensor.rand(X.shape[0],1,1,1) < 0.5, X[..., ::-1], X) # flip LR
-      if step >= hyp['net']['cutmix_steps']: X, Y = cutmix(X, Y, mask_size=hyp['net']['cutmix_size'])
-    X, Y = X.numpy(), Y.numpy()
-    for i in range(0, X.shape[0], BS):
-      # pad the last batch
-      batch_end = min(i+BS, Y.shape[0])
-      x = Tensor(X[order[batch_end-BS:batch_end],:])
-      y = Tensor(Y[order[batch_end-BS:batch_end]])
-      step += 1
-      yield x, y
+    X_patch = Tensor(X.numpy()[order,...])
+    Y_patch = Tensor(Y.numpy()[order])
+    X_cutmix = Tensor.where(mask, X_patch, X)
+    mix_portion = float(mask_size**2)/(X.shape[-2]*X.shape[-1])
+    Y_cutmix = mix_portion * Y_patch + (1. - mix_portion) * Y
+    return X_cutmix, Y_cutmix
 
-    if not is_train: break
-    seed += 1
+  # the operations that remain inside batch fetcher is the ones that involves random operations
+  def fetch_batches(X_in:Tensor, Y_in:Tensor, BS:int, is_train:bool):
+    step, cnt = 0, 0
+    while True:
+      st = time.monotonic()
+      X, Y = X_in, Y_in
+      order = list(range(0, X.shape[0]))
+      random.shuffle(order)
+      if is_train:
+        X = random_crop(X, crop_size=32)
+        X = Tensor.where(Tensor.rand(X.shape[0],1,1,1) < 0.5, X[..., ::-1], X) # flip LR
+        if step >= hyp['net']['cutmix_steps']: X, Y = cutmix(X, Y, mask_size=hyp['net']['cutmix_size'])
+      X, Y = X.numpy(), Y.numpy()
+      et = time.monotonic()
+      print(f"shuffling {'training' if is_train else 'test'} dataset in {(et-st)*1e3:.2f} ms ({cnt})")
+      for i in range(0, X.shape[0], BS):
+        # pad the last batch
+        batch_end = min(i+BS, Y.shape[0])
+        x = Tensor(X[order[batch_end-BS:batch_end],:])
+        y = Tensor(Y[order[batch_end-BS:batch_end]])
+        step += 1
+        yield x, y
+      cnt += 1
+      if not is_train: break
 
-transform = [
-  lambda x: x / 255.0,
-  lambda x: (x - Tensor(cifar_mean).repeat((1024,1)).T.reshape(1,-1))/ Tensor(cifar_std).repeat((1024,1)).T.reshape(1,-1),
-  lambda x: x.reshape((-1,3,32,32))
-]
+  transform = [
+    lambda x: x / 255.0,
+    lambda x: (x.reshape((-1,3,32,32)) - Tensor(cifar_mean).reshape((1,3,1,1)))/Tensor(cifar_std).reshape((1,3,1,1))
+  ]
+
+  class modelEMA():
+    def __init__(self, w, net):
+      # self.model_ema = copy.deepcopy(net) # won't work for opencl due to unpickeable pyopencl._cl.Buffer
+      self.net_ema = SpeedyResNet(w)
+      for net_ema_param, net_param in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).values()):
+        net_ema_param.requires_grad = False
+        net_ema_param.assign(net_param.numpy())
+
+    @TinyJit
+    def update(self, net, decay):
+      # TODO with Tensor.no_grad()
+      Tensor.no_grad = True
+      for net_ema_param, (param_name, net_param) in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).items()):
+        # batchnorm currently is not being tracked
+        if not ("num_batches_tracked" in param_name) and not ("running" in param_name):
+          net_ema_param.assign(net_ema_param.detach()*decay + net_param.detach()*(1.-decay)).realize()
+      Tensor.no_grad = False
+
+  set_seed(hyp['seed'])
 
-def train_cifar(bs=BS, eval_bs=EVAL_BS, steps=STEPS, seed=32):
   # this import needs to be done here because this is running in a subprocess
   from extra.dist import OOB
-  set_seed(seed)
-  Tensor.training = True
+  assert OOB is not None or not getenv("DIST"), "OOB should be initialized"
   rank, world_size = getenv("RANK"), getenv("WORLD_SIZE", 1)
 
   X_train, Y_train, X_test, Y_test = fetch_cifar()
@@ -220,17 +262,22 @@ def train_cifar(bs=BS, eval_bs=EVAL_BS, steps=STEPS, seed=32):
   X_train, X_test = X_train.to(device=Device.DEFAULT).float(), X_test.to(device=Device.DEFAULT).float()
   Y_train, Y_test = Y_train.to(device=Device.DEFAULT).float(), Y_test.to(device=Device.DEFAULT).float()
   # one-hot encode labels
-  Y_train, Y_test = Tensor.eye(10)[Y_train], Tensor.eye(10)[Y_test]
+  Y_train, Y_test = Tensor.eye(10)[Y_train.cast(dtypes.int32)], Tensor.eye(10)[Y_test.cast(dtypes.int32)]
   # preprocess data
   X_train, X_test = X_train.sequential(transform), X_test.sequential(transform)
 
   # precompute whitening patches
   W = whitening(X_train)
 
+  # initialize model weights
+  model = SpeedyResNet(W)
+
   # padding is not timed in the original repo since it can be done all at once
   X_train = pad_reflect(X_train, size=hyp['net']['pad_amount'])
 
-  model = SpeedyResNet(W)
+  # Convert data and labels to the default dtype
+  X_train, Y_train = X_train.cast(dtypes.default_float), Y_train.cast(dtypes.default_float)
+  X_test, Y_test = X_test.cast(dtypes.default_float), Y_test.cast(dtypes.default_float)
 
   # parse the training params into bias and non-bias
   params_dict = get_state_dict(model)
@@ -247,8 +294,8 @@ def train_cifar(bs=BS, eval_bs=EVAL_BS, steps=STEPS, seed=32):
   opt_non_bias = optim.SGD(params_non_bias, lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['non_bias_decay'])
 
   # NOTE taken from the hlb_CIFAR repository, might need to be tuned
-  initial_div_factor = 1e16
-  final_lr_ratio = 0.02199
+  initial_div_factor = hyp['opt']['initial_div_factor']
+  final_lr_ratio = hyp['opt']['final_lr_ratio']
   pct_start = hyp['opt']['percent_start']
   lr_sched_bias     = OneCycleLR(opt_bias,     max_lr=hyp['opt']['bias_lr']     ,pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=STEPS)
   lr_sched_non_bias = OneCycleLR(opt_non_bias, max_lr=hyp['opt']['non_bias_lr'] ,pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=STEPS)
@@ -257,7 +304,7 @@ def train_cifar(bs=BS, eval_bs=EVAL_BS, steps=STEPS, seed=32):
   @TinyJit
   def train_step_jitted(model, optimizer, lr_scheduler, X, Y):
     out = model(X)
-    loss = cross_entropy(out, Y, reduction='none' ,label_smoothing=0.2).mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])
+    loss = cross_entropy(out, Y, reduction='none' ,label_smoothing=hyp['opt']['label_smoothing']).mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])
 
     if not getenv("DISABLE_BACKWARD"):
       # index 0 for bias and 1 for non-bias
@@ -270,7 +317,7 @@ def train_cifar(bs=BS, eval_bs=EVAL_BS, steps=STEPS, seed=32):
         bucket, offset = [], 0
         for _, v in params_dict.items():
           if v.grad is not None: bucket.append(v.grad.flatten())
-        grads = collectives.allreduce(Tensor.cat(*bucket), cache_id="grads")
+        grads = collectives.allreduce(Tensor.cat(*bucket))
         for _, v in params_dict.items():
           if v.grad is not None:
             v.grad.assign(grads[offset:offset+v.grad.numel()].reshape(*v.grad.shape))
@@ -282,15 +329,15 @@ def train_cifar(bs=BS, eval_bs=EVAL_BS, steps=STEPS, seed=32):
       lr_scheduler[1].step()
     return loss.realize()
 
-  @TinyJit
-  def eval_step_jitted(model, X, Y):
+  def eval_step(model, X, Y):
     out = model(X, training=False)
     loss = cross_entropy(out, Y, reduction='mean')
     correct = out.argmax(axis=1) == Y.argmax(axis=1)
     return correct.realize(), loss.realize()
+  eval_step_jitted     = TinyJit(eval_step)
+  eval_step_ema_jitted = TinyJit(eval_step)
 
-  # 97 steps in 2 seconds = 20ms / step  Tensor.training = True
-
+  # 97 steps in 2 seconds = 20ms / step
   # step is 1163.42 GOPS = 56 TFLOPS!!!, 41% of max 136
   # 4 seconds for tfloat32 ~ 28 TFLOPS, 41% of max 68
   # 6.4 seconds for float32 ~ 17 TFLOPS, 50% of max 34.1
@@ -299,60 +346,93 @@ def train_cifar(bs=BS, eval_bs=EVAL_BS, steps=STEPS, seed=32):
   # https://www.anandtech.com/show/16727/nvidia-announces-geforce-rtx-3080-ti-3070-ti-upgraded-cards-coming-in-june
   # 136 TFLOPS is the theoretical max w float16 on 3080 Ti
 
-  best_eval = -1
+  model_ema: Optional[modelEMA] = None
+  projected_ema_decay_val = hyp['ema']['decay_base'] ** hyp['ema']['every_n_steps']
   i = 0
-  batcher = fetch_batches(X_train, Y_train, BS=BS, seed=seed, is_train=True)
-  while i <= STEPS:
-    if i%100 == 0 and i > 1:
-      # Use Tensor.training = False here actually bricks batchnorm, even with track_running_stats=True
-      corrects = []
-      losses = []
-      for Xt, Yt in fetch_batches(X_test, Y_test, BS=EVAL_BS, seed=seed, is_train=False):
-        # further split batch if distributed
-        if getenv("DIST"):
-          Xt, Yt = Xt.chunk(min(world_size, 5), 0)[min(rank, 4)], Yt.chunk(min(world_size, 5), 0)[min(rank, 4)]
-
-        correct, loss = eval_step_jitted(model, Xt, Yt)
-        losses.append(loss.numpy().tolist())
-        corrects.extend(correct.numpy().tolist())
-
-      # collect accuracy across ranks
-      correct_sum, correct_len = sum(corrects), len(corrects)
-      if getenv("DIST"):
-        if rank == 0:
-          for j in range(1, min(world_size, 5)):
-            recv_sum, recv_len = OOB.recv(j)
-            correct_sum += recv_sum
-            correct_len += recv_len
-        elif rank < min(world_size, 5):
-          OOB.send((correct_sum, correct_len), 0)
-
-      # only rank 0 prints
-      if rank == 0:
-        acc = correct_sum/correct_len*100.0
-        if acc > best_eval:
-          best_eval = acc
-          print(f"eval {correct_sum}/{correct_len} {acc:.2f}%, {(sum(losses)/len(losses)):7.2f} val_loss STEP={i}")
-    if STEPS == 0 or i==STEPS: break
-    X, Y = next(batcher)
-    # further split batch if distributed
-    if getenv("DIST"):
-      X, Y = X.chunk(world_size, 0)[rank], Y.chunk(world_size, 0)[rank]
-    GlobalCounters.reset()
+  batcher = fetch_batches(X_train, Y_train, BS=BS, is_train=True)
+  with Tensor.train():
     st = time.monotonic()
-    loss = train_step_jitted(model, [opt_bias, opt_non_bias], [lr_sched_bias, lr_sched_non_bias], X, Y)
-    et = time.monotonic()
-    loss_cpu = loss.numpy()
-    cl = time.monotonic()
-    print(f"{i:3d} {(cl-st)*1000.0:7.2f} ms run, {(et-st)*1000.0:7.2f} ms python, {(cl-et)*1000.0:7.2f} ms CL, {loss_cpu:7.2f} loss, {opt_non_bias.lr.numpy()[0]:.6f} LR, {GlobalCounters.mem_used/1e9:.2f} GB used, {GlobalCounters.global_ops*1e-9/(cl-st):9.2f} GFLOPS")
-    i += 1
+    while i <= STEPS:
+      if i%getenv("EVAL_STEPS", STEPS) == 0 and i > 1:
+        st_eval = time.monotonic()
+        # Use Tensor.training = False here actually bricks batchnorm, even with track_running_stats=True
+        corrects = []
+        corrects_ema = []
+        losses = []
+        losses_ema = []
+        for Xt, Yt in fetch_batches(X_test, Y_test, BS=EVAL_BS, is_train=False):
+          # further split batch if distributed
+          if getenv("DIST"):
+            Xt, Yt = Xt.chunk(min(world_size, 5), 0)[min(rank, 4)], Yt.chunk(min(world_size, 5), 0)[min(rank, 4)]
+
+          correct, loss = eval_step_jitted(model, Xt, Yt)
+          losses.append(loss.numpy().tolist())
+          corrects.extend(correct.numpy().tolist())
+          if model_ema:
+            correct_ema, loss_ema = eval_step_ema_jitted(model_ema.net_ema, Xt, Yt)
+            losses_ema.append(loss_ema.numpy().tolist())
+            corrects_ema.extend(correct_ema.numpy().tolist())
+
+        # collect accuracy across ranks
+        correct_sum, correct_len = sum(corrects), len(corrects)
+        if model_ema: correct_sum_ema, correct_len_ema = sum(corrects_ema), len(corrects_ema)
+        if getenv("DIST"):
+          if rank == 0:
+            for j in range(1, min(world_size, 5)):
+              if model_ema:
+                recv_sum, recv_len, recv_sum_ema, recv_len_ema = OOB.recv(j)
+              else:
+                recv_sum, recv_len = OOB.recv(j)
+              correct_sum += recv_sum
+              correct_len += recv_len
+              if model_ema:
+                correct_sum_ema += recv_sum_ema
+                correct_len_ema += recv_len_ema
+          elif rank < min(world_size, 5):
+            if model_ema:
+              OOB.send((correct_sum, correct_len, correct_sum_ema, correct_len_ema), 0)
+            else:
+              OOB.send((correct_sum, correct_len), 0)
+
+        # only rank 0 prints
+        if rank == 0:
+          acc = correct_sum/correct_len*100.0
+          if model_ema: acc_ema = correct_sum_ema/correct_len_ema*100.0
+          print(f"eval     {correct_sum}/{correct_len} {acc:.2f}%, {(sum(losses)/len(losses)):7.2f} val_loss STEP={i} (in {(time.monotonic()-st)*1e3:.2f} ms)")
+          if model_ema: print(f"eval ema {correct_sum_ema}/{correct_len_ema} {acc_ema:.2f}%, {(sum(losses_ema)/len(losses_ema)):7.2f} val_loss STEP={i}")
+
+      if STEPS == 0 or i==STEPS: break
+      X, Y = next(batcher)
+      if getenv("DIST"):
+        X, Y = X.chunk(world_size, 0)[rank], Y.chunk(world_size, 0)[rank]
+      GlobalCounters.reset()
+      loss = train_step_jitted(model, [opt_bias, opt_non_bias], [lr_sched_bias, lr_sched_non_bias], X, Y)
+      et = time.monotonic()
+      loss_cpu = loss.numpy()
+      # EMA for network weights
+      if i > hyp['ema']['steps'] and (i+1) % hyp['ema']['every_n_steps'] == 0:
+        if model_ema is None:
+          model_ema = modelEMA(W, model)
+        model_ema.update(model, Tensor([projected_ema_decay_val*(i/STEPS)**hyp['ema']['decay_pow']]))
+      cl = time.monotonic()
+      if not getenv("DIST"):
+        #  53  221.74 ms run,    2.22 ms python,  219.52 ms CL,  803.39 loss, 0.000807 LR, 4.66 GB used,   3042.49 GFLOPS,    674.65 GOPS
+        print(f"{i:3d} {(cl-st)*1000.0:7.2f} ms run, {(et-st)*1000.0:7.2f} ms python, {(cl-et)*1000.0:7.2f} ms CL, {loss_cpu:7.2f} loss, {opt_non_bias.lr.numpy()[0]:.6f} LR, {GlobalCounters.mem_used/1e9:.2f} GB used, {GlobalCounters.global_ops*1e-9/(cl-st):9.2f} GFLOPS, {GlobalCounters.global_ops*1e-9:9.2f} GOPS")
+      else:
+        print(f"{i:3d} {(cl-st)*1000.0:7.2f} ms run, {(et-st)*1000.0:7.2f} ms python, {(cl-et)*1000.0:7.2f} ms CL, {loss_cpu:7.2f} loss, {opt_non_bias.lr.numpy()[0]:.6f} LR, {world_size*GlobalCounters.mem_used/1e9:.2f} GB used, {world_size*GlobalCounters.global_ops*1e-9/(cl-st):9.2f} GFLOPS")
+      st = cl
+      i += 1
 
 if __name__ == "__main__":
   if not getenv("DIST"):
     train_cifar()
   else: # distributed
-    from tinygrad.runtime.ops_gpu import CL
-    devices = [f"gpu:{i}" for i in range(len(CL.devices))]
+    if getenv("HIP"):
+      from tinygrad.runtime.ops_hip import HIP
+      devices = [f"hip:{i}" for i in range(HIP.device_count)]
+    else:
+      from tinygrad.runtime.ops_gpu import CLDevice
+      devices = [f"gpu:{i}" for i in range(len(CLDevice.device_ids))]
     world_size = len(devices)
 
     # ensure that the batch size is divisible by the number of devices

examples/index.html

@@ -18,6 +18,7 @@ canvas {  display: none; }
 </style>
 <title>tinygrad has WebGPU</title>
 <script src="./net.js"></script>
+<link rel="icon" type="image/x-icon" href="https://raw.githubusercontent.com/tinygrad/tinygrad/master/docs/logo.png">
 </head>
 <body>
 <h1>WebGPU <a href="https://github.com/geohot/tinygrad">tinygrad</a> EfficientNet!</h1>
@@ -111,7 +112,7 @@ canvas {  display: none; }
 		const input = reorderChannelsAndRemoveAlpha(Array.from(data).map((pix) => (pix / 255.0) * 0.45 - 0.225));
 		const out = await timer(() => net(new Float32Array(input)));
 
-		const arr = Array.from(new Float32Array(out));
+		const arr = Array.from(new Float32Array(out[0]));
 		const index = arr.indexOf(Math.max(...arr));
 
 		resultText.textContent = labels[index];

examples/llama.py

@@ -1,230 +1,130 @@
 #!/usr/bin/env python3
-# pip3 install sentencepiece pyobjc-framework-Metal pyobjc-framework-Cocoa pyobjc-framework-libdispatch
+# pip3 install sentencepiece
 #import typeguard.importhook
 #typeguard.importhook.install_import_hook('tinygrad')
 
 from pathlib import Path
-import functools, sys, argparse, math, platform
+import sys, argparse, json
 import numpy as np
-from tqdm import tqdm
 np.set_printoptions(linewidth=200)
-from typing import Optional, Tuple
-
-from tinygrad.helpers import Timing, getenv, DEBUG, dtypes
-from tinygrad.ops import Device
+from tinygrad.helpers import Timing, Profiling, getenv, DEBUG, colored
+from tinygrad import Device, GlobalCounters, dtypes
 from tinygrad.tensor import Tensor
-from tinygrad.nn import Embedding, Linear
-from tinygrad.ops import GlobalCounters
-from tinygrad.jit import TinyJit
-from tinygrad.shape.symbolic import Variable, sym_infer
+from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters
+from extra.models.llama import Transformer, convert_from_huggingface
+from sentencepiece import SentencePieceProcessor
 
-# https://github.com/facebookresearch/llama/blob/1076b9c51c77ad06e9d7ba8a4c6df775741732bd/llama/model.py#L47
-def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
-  freqs = 1.0 / (theta ** (np.arange(0, dim, 2, dtype=np.float32)[:(dim // 2)] / dim))
-  freqs = np.outer(np.arange(end, dtype=np.float32), freqs)
-  return np.stack([np.cos(freqs), np.sin(freqs)], axis=-1).reshape(1, end, 1, dim//2, 2)
+MAX_CONTEXT = getenv("MAX_CONTEXT", 4096)
 
-# (a+i*b) * (c+i*d) = (ac-bd) + i*(ad+bc)
-def complex_mult(A, c, d):
-  a,b = A[:, :, :, :, 0:1], A[:, :, :, :, 1:2]
-  ro = a*c - b*d
-  co = a*d + b*c
-  return ro.cat(co, dim=-1)
-
-def apply_rotary_emb(xq, xk, freqs_cis) -> Tuple[Tensor, Tensor]:
-  assert freqs_cis.shape[1] == xq.shape[1] and freqs_cis.shape[1] == xk.shape[1], f"freqs_cis shape mismatch {freqs_cis.shape} xq:{xq.shape} xk:{xk.shape}"
-  xq = xq.reshape(*xq.shape[0:-1], -1, 2)
-  xk = xk.reshape(*xk.shape[0:-1], -1, 2)
-  assert len(xq.shape) == 5 and len(xk.shape) == 5 and len(freqs_cis.shape) == 5
-  c, d = freqs_cis[:, :xq.shape[1], :, :, 0:1], freqs_cis[:, :xq.shape[1], :, :, 1:2]
-  xq_out = complex_mult(xq, c, d)
-  xk_out = complex_mult(xk, c, d)
-  return xq_out.flatten(3), xk_out.flatten(3)
-
-def repeat_kv(x:Tensor, n_rep:int) -> Tensor:
-  bs, seqlen, n_kv_heads, head_dim = x.shape
-  if n_rep == 1: return x
-  return x[:, :, :, None, :].expand(bs, seqlen, n_kv_heads, n_rep, head_dim).reshape(bs, seqlen, n_kv_heads * n_rep, head_dim)
-
-class RMSNorm:
-  def __init__(self, dim, eps=1e-6):
-    self.eps = eps
-    self.weight = Tensor.ones(dim)
-
-  def __call__(self, x:Tensor):
-    # TODO: convert to float?
-    return (x * (x.pow(2).mean(-1, keepdim=True) + self.eps).rsqrt()) * self.weight
-
-class Attention:
-  def __init__(self, dim, n_heads, n_kv_heads, linear=Linear):
-    self.n_heads = n_heads
-    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads
-    self.head_dim = dim // n_heads
-    self.n_rep = self.n_heads // self.n_kv_heads
-
-    self.wq = linear(dim, self.n_heads * self.head_dim, bias=False)
-    self.wk = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-    self.wv = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-    self.wo = linear(self.n_heads * self.head_dim, dim, bias=False)
-
-  def __call__(self, x:Tensor, cache_k:Tensor, cache_v:Tensor, start_pos:int, freqs_cis:Tensor, mask:Optional[Tensor]) -> Tuple[Tensor, Tensor, Tensor]:
-    bsz, seqlen, _ = x.shape
-    xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
-    xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
-    xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
-    xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
-
-    # kv caching!
-    if start_pos == 0:
-      keys, values = xk, xv
-    else:
-      assert cache_k.shape[0] > 0, "no cache"
-      assert start_pos == sym_infer(cache_k.shape[1], cache_k.lazydata.st.var_vals) == sym_infer(cache_v.shape[1], cache_v.lazydata.st.var_vals), f"cache has wrong shape, not ({start_pos} == {sym_infer(cache_k.shape[1], cache_k.lazydata.st.var_vals)} == {sym_infer(cache_v.shape[1], cache_v.lazydata.st.var_vals)})"
-      assert seqlen == xk.shape[1] and seqlen == xv.shape[1], "seqlen is wrong shape?!?"
-      keys, values = cache_k.cat(xk, dim=1), cache_v.cat(xv, dim=1)
-
-    cache_k, cache_v = keys, values
-    keys, values = repeat_kv(keys, self.n_rep).realize(), repeat_kv(values, self.n_rep).realize()
-    attn = Tensor.scaled_dot_product_attention(xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2), mask).transpose(1, 2).reshape(bsz, seqlen, -1)
-    return self.wo(attn).realize(), cache_k.realize(), cache_v.realize()
-
-class FeedForward:
-  def __init__(self, dim, hidden_dim, multiple_of, linear=Linear, ffn_dim_multiplier=None):
-    # TODO: what is this?
-    hidden_dim = int(2 * hidden_dim / 3)
-    # custom dim factor multiplier
-    if ffn_dim_multiplier is not None:
-      hidden_dim = int(ffn_dim_multiplier * hidden_dim)
-    hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
-    self.w1 = linear(dim, hidden_dim, bias=False)
-    self.w2 = linear(hidden_dim, dim, bias=False)
-    self.w3 = linear(dim, hidden_dim, bias=False)
-
-  def __call__(self, x:Tensor) -> Tensor:
-    return self.w2(self.w1(x).silu() * self.w3(x))
-
-class TransformerBlock:
-  def __init__(self, dim, multiple_of, n_heads, n_kv_heads, norm_eps, linear=Linear, ffn_dim_multiplier=None):
-    self.attention = Attention(dim, n_heads, n_kv_heads, linear)
-    self.feed_forward = FeedForward(dim, 4*dim, multiple_of, linear, ffn_dim_multiplier)
-    self.attention_norm = RMSNorm(dim, norm_eps)
-    self.ffn_norm = RMSNorm(dim, norm_eps)
-    self.cache_k, self.cache_v = None, None
-
-    self.jitted_attention_norm = TinyJit(lambda x: self.attention_norm(x).realize())
-    self.jitted_attn = TinyJit(self.attention.__call__)
-    self.jitted_norm_output = TinyJit(self.norm_output)
-
-  def norm_output(self, x:Tensor, output:Tensor) -> Tensor:
-    h = x + output
-    return (h + self.feed_forward(self.ffn_norm(h))).realize()
-
-  def __call__(self, x:Tensor, start_pos:int, freqs_cis:Tensor, mask:Optional[Tensor]):
-    bsz, seqlen, _ = x.shape
-    do_jit = getenv("JIT") and mask is None
-    if do_jit:
-      pos = Variable("pos", 1, 1024)
-      self.cache_k = self.cache_k.reshape(self.cache_k.shape[0], pos, self.cache_k.shape[2], self.cache_k.shape[3])
-      self.cache_v = self.cache_v.reshape(self.cache_v.shape[0], pos, self.cache_v.shape[2], self.cache_v.shape[3])
-      output, cache_k, cache_v = self.jitted_attn(self.jitted_attention_norm(x), self.cache_k, self.cache_v, start_pos, freqs_cis, mask)
-    else:
-      output, cache_k, cache_v = self.attention(self.attention_norm(x), self.cache_k, self.cache_v, start_pos, freqs_cis, mask)
-
-    # save the cache. with symbolic shape, cast it back to int shape so we have int shape in cache
-    self.cache_k = cache_k.reshape(cache_k.shape[0], start_pos+seqlen, cache_k.shape[2], cache_k.shape[3]).realize()
-    self.cache_v = cache_v.reshape(cache_v.shape[0], start_pos+seqlen, cache_v.shape[2], cache_v.shape[3]).realize()
-
-    return self.jitted_norm_output(x, output) if do_jit else self.norm_output(x, output)
-
-class Transformer:
-  def __init__(self, dim, multiple_of, n_heads, n_layers, norm_eps, vocab_size, linear=Linear, max_batch_size=32, max_seq_len=1024, ffn_dim_multiplier=None, n_kv_heads=None):
-    self.layers = [TransformerBlock(dim, multiple_of, n_heads, n_kv_heads, norm_eps, linear, ffn_dim_multiplier) for _ in range(n_layers)]
-    self.norm = RMSNorm(dim, norm_eps)
-    self.tok_embeddings = Embedding(vocab_size, dim)
-    self.output = linear(dim, vocab_size, bias=False)
-    self.freqs_cis = Tensor(precompute_freqs_cis(dim // n_heads, max_seq_len * 2))
-    self.norm_output = lambda x: self.output(self.norm(x))
-
-    self.jitted_tok_embeddings = TinyJit(lambda x: self.tok_embeddings(x).realize())
-    self.jitted_norm_output = TinyJit(lambda x: self.norm_output(x).realize())
-
-  def __call__(self, tokens:Tensor, start_pos:int):
-    _bsz, seqlen = tokens.shape
-    mask = Tensor.full((1, 1, seqlen, start_pos + seqlen), float("-inf"), dtype=dtypes.float32).triu(start_pos+1).realize() if seqlen > 1 else None
-    do_jit = getenv("JIT") and mask is None
-
-    # get only the part of freqs_cis that we are using.
-    if do_jit:
-      pos = Variable("pos", 1, 1024)
-      assert seqlen == 1, "seqlen > 1 not supported for JIT"
-      freqs_cis = self.freqs_cis.shrink(((0, self.freqs_cis.shape[0]), (pos, pos+seqlen),(0, self.freqs_cis.shape[2]),(0, self.freqs_cis.shape[3]),(0, self.freqs_cis.shape[4])))
-      freqs_cis.lazydata.st.var_vals[pos] = start_pos
-    else:
-      freqs_cis = self.freqs_cis.shrink(((0, self.freqs_cis.shape[0]), (start_pos, start_pos+seqlen),(0, self.freqs_cis.shape[2]),(0, self.freqs_cis.shape[3]),(0, self.freqs_cis.shape[4])))
-
-    h = self.jitted_tok_embeddings(tokens) if do_jit else self.tok_embeddings(tokens)
-    h = h.sequential([functools.partial(layer, start_pos=start_pos, freqs_cis=freqs_cis, mask=mask) for layer in self.layers])
-    return self.jitted_norm_output(h) if do_jit else self.norm_output(h)
-
-# **** files and arguments ****
-
-VOCAB_SIZE = 32000
+# calculating params:
+# traditionally, the MLP in the transformer architecture has hidden_dim = dim*4 [arxiv/1706.03762, 3.3]
+# however, Llama uses SwiGLU. in order to preserve param count to original transformer arch, hidden_dim must be = 2/3 * (dim*4) [arxiv/2002.05202]
+# for models using MQA (n_kv_heads != n_heads), preserving param count means hidden dim must be further multiplied by 1.3 [arxiv/2307.09288, A.2.1]
 MODEL_PARAMS = {
-  1: {
+  "1": {
     "7B": {
-      "args": {"dim": 4096, "multiple_of": 256, "n_heads": 32, "n_layers": 32, "norm_eps": 1e-06, "vocab_size": VOCAB_SIZE},
+      "args": {"dim": 4096, "n_heads": 32, "n_layers": 32, "norm_eps": 1e-06, "vocab_size": 32000, "hidden_dim": 11008},
       "files": 1,
     },
     "13B": {
-      "args": {"dim": 5120, "multiple_of": 256, "n_heads": 40, "n_layers": 40, "norm_eps": 1e-06, "vocab_size": VOCAB_SIZE},
+      "args": {"dim": 5120, "n_heads": 40, "n_layers": 40, "norm_eps": 1e-06, "vocab_size": 32000, "hidden_dim": 13824},
       "files": 2,
     },
     "30B": {
-      "args": {"dim": 6656, "multiple_of": 256, "n_heads": 52, "n_layers": 60, "norm_eps": 1e-06, "vocab_size": VOCAB_SIZE},
+      "args": {"dim": 6656, "n_heads": 52, "n_layers": 60, "norm_eps": 1e-06, "vocab_size": 32000, "hidden_dim": 17920},
       "files": 4,
     },
     "65B": {
-      "args": {"dim": 8192, "multiple_of": 256, "n_heads": 64, "n_layers": 80, "norm_eps": 1e-05, "vocab_size": VOCAB_SIZE},
+      "args": {"dim": 8192, "n_heads": 64, "n_layers": 80, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 22016},
       "files": 8,
     },
   },
-  2: {
+  "2": {
     "7B": {
-      "args": {"dim": 4096, "multiple_of": 256, "n_heads": 32, "n_layers": 32, "norm_eps": 1e-05, "vocab_size": VOCAB_SIZE},
+      "args": {"dim": 4096, "n_heads": 32, "n_layers": 32, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 11008},
       "files": 1,
     },
     "13B": {
-      "args": {"dim": 5120, "multiple_of": 256, "n_heads": 40, "n_layers": 40, "norm_eps": 1e-05, "vocab_size": VOCAB_SIZE},
+      "args": {"dim": 5120, "n_heads": 40, "n_layers": 40, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 13824},
       "files": 2,
     },
     "70B": {
-      "args": {"dim": 8192, "multiple_of": 4096, "ffn_dim_multiplier": 1.3, "n_heads": 64, "n_kv_heads": 8, "n_layers": 80, "norm_eps": 1e-05, "vocab_size": VOCAB_SIZE},
+      "args": {"dim": 8192, "n_heads": 64, "n_kv_heads": 8, "n_layers": 80, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 28672},
       "files": 8,
     },
   },
+  "code": {
+    "7B": {
+      "args": {"dim": 4096, "n_layers": 32, "n_heads": 32, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32016, "hidden_dim": 11008},
+      "files": 1,
+    },
+    "7B-Python": {
+      "args": {"dim": 4096, "n_layers": 32, "n_heads": 32, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 11008},
+      "files": 1,
+    },
+    "7B-Instruct": {
+      "args": {"dim": 4096, "n_layers": 32, "n_heads": 32, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32016, "hidden_dim": 11008},
+      "files": 1,
+    },
+    "13B": {
+      "args": {"dim": 5120, "n_layers": 40, "n_heads": 40, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32016, "hidden_dim": 13824},
+      "files": 2,
+    },
+    "13B-Python": {
+      "args": {"dim": 5120, "n_layers": 40, "n_heads": 40, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 13824},
+      "files": 2,
+    },
+    "13B-Instruct": {
+      "args": {"dim": 5120, "n_layers": 40, "n_heads": 40, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32016, "hidden_dim": 13824},
+      "files": 2,
+    },
+    "34B": {
+      "args": {"dim": 8192, "n_layers": 48, "n_heads": 64, "n_kv_heads": 8, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 22016},
+      "files": 4,
+    },
+    "34B-Python": {
+      "args": {"dim": 8192, "n_layers": 48, "n_heads": 64, "n_kv_heads": 8, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 22016},
+      "files": 4,
+    },
+    "34B-Instruct": {
+      "args": {"dim": 8192, "n_layers": 48, "n_heads": 64, "n_kv_heads": 8, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 22016},
+      "files": 4,
+    },
+  },
+  "tiny": {
+    "1B": {
+      "args": {"dim": 2048, "n_layers": 22, "n_heads": 32, "n_kv_heads": 4, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 5632},
+      "files": 1,
+    },
+    "1B-Chat": {
+      "args": {"dim": 2048, "n_layers": 22, "n_heads": 32, "n_kv_heads": 4, "norm_eps": 1e-05, "vocab_size": 32003, "hidden_dim": 5632},
+      "files": 1,
+    }
+  }
 }
 
-# **** helper functions ****
-def sample(logits, temperature):
-  if temperature < 1e-6:
-    # so close to 0 we use argmax
-    return int(logits.argmax().numpy())
-  else:
-    probs = (logits / temperature).softmax()
-    probs = probs.numpy().flatten()
-    return int(np.random.choice(len(probs), p=probs))
 
+# **** helper functions ****
 def concat_weights(models):
   def convert(name) -> Tensor:
     disk_tensors = [model[name] for model in models]
     if len(disk_tensors) == 1 or len(disk_tensors[0].shape) == 1:
       return disk_tensors[0].to(device=Device.DEFAULT)
-    axis = 1 if name.startswith('tok_embeddings.') or name.endswith('.attention.wo.weight') or name.endswith('.feed_forward.w2.weight') else 0
+    axis = 1 if name.startswith("tok_embeddings.") or name.endswith(".attention.wo.weight") or name.endswith(".feed_forward.w2.weight") else 0
     lazy_tensors = [data.to(device=Device.DEFAULT) for data in disk_tensors]
     return lazy_tensors[0].cat(*lazy_tensors[1:], dim=axis)
   return {name: convert(name) for name in {name: None for model in models for name in model}}
 
+def load(fn:str):
+  if fn.endswith('.index.json'):
+    with open(fn) as fp: weight_map = json.load(fp)['weight_map']
+    parts = {n: load(str(Path(fn).parent / Path(n).name)) for n in set(weight_map.values())}
+    return {k: parts[n][k] for k, n in weight_map.items()}
+  elif fn.endswith(".safetensors"):
+    return safe_load(fn)
+  else:
+    return torch_load(fn)
+
 class AbsmaxQuantizedLinear:
   def __init__(self, in_features, out_features, bias=False):
     assert bias == False
@@ -238,42 +138,53 @@ class AbsmaxQuantizedLinear:
   def quantize(tensors):
     new_tensors = {}
     for name,v in tensors.items():
-      if 'feed_forward' in name or ('attention.w') in name or name == 'output.weight':
+      if "feed_forward" in name or ("attention.w") in name or name == "output.weight":
         scale = v.abs().max(axis=1) / 127.0
         int8_weight = (v.T/scale).T.cast(dtype=dtypes.int8)
-        new_tensors[name] = int8_weight.realize()
-        new_tensors[name.replace('weight', 'scale')] = scale.realize()
+        new_tensors[name] = int8_weight
+        new_tensors[name.replace('weight', 'scale')] = scale
       else:
         new_tensors[name] = v
     return new_tensors
 
 class LLaMa:
   @staticmethod
-  def build(model_path, tokenizer_path, model_gen=1, model_size="7B", quantize=False):
-    from sentencepiece import SentencePieceProcessor
-    sp_model = SentencePieceProcessor(model_file=str(tokenizer_path))
-    assert sp_model.vocab_size() == VOCAB_SIZE
-
-    from tinygrad.nn.state import torch_load, load_state_dict
+  def build(model_path, tokenizer_path, model_gen="1", model_size="7B", quantize=False):
     params = MODEL_PARAMS[model_gen][model_size]
-    model = Transformer(**params["args"], linear=AbsmaxQuantizedLinear) if quantize else Transformer(**params["args"])
-    weights = concat_weights([torch_load(filename) for filename in [f"{model_path}/{model_size}/consolidated.{i:02d}.pth" for i in range(params["files"])]])
+    sp_model = SentencePieceProcessor(model_file=str(tokenizer_path))
+    assert sp_model.vocab_size() == params["args"]["vocab_size"], f"{sp_model.vocab_size()=} not equal to {params['args']['vocab_size']}"
+
+    jit = bool(getenv("JIT", 1))
+    model = Transformer(**params["args"], linear=AbsmaxQuantizedLinear, max_context=MAX_CONTEXT, jit=jit) if quantize else Transformer(**params["args"], max_context=MAX_CONTEXT, jit=jit)
+
+    if model_path.is_dir():
+      weights = concat_weights([load(filename) for filename in [f"{model_path}/consolidated.{i:02d}.pth" for i in range(params["files"])]])
+    else:
+      weights = load(str(model_path))
+    if "model.embed_tokens.weight" in weights:
+      weights = convert_from_huggingface(weights, model, params["args"]["n_heads"], params["args"].get("n_kv_heads", params["args"]["n_heads"]))
+
+    # fix bf16, TODO: check if device supports bf16
+    weights = {k:v.to(Device.DEFAULT).cast(dtypes.float16) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
+
     if quantize:
       weights = AbsmaxQuantizedLinear.quantize(weights)
+      for _,v in weights.items(): v.realize()
     load_state_dict(model, weights, strict=False)
 
     return LLaMa(model, sp_model)
 
   def __init__(self, model, tokenizer):
     self.model = model
-    self.tokenizer = tokenizer
+    self.tokenizer: SentencePieceProcessor = tokenizer
 
   def greedy_until(self, prompt:str, until, max_length, temperature):
     toks = [self.tokenizer.bos_id()] + self.tokenizer.encode(prompt)
     start_pos = 0
     for i in range(max_length):
-      logits = self.model(Tensor([toks[start_pos:]]), start_pos).realize()[:, -1, :]
-      tok = sample(logits, temperature)
+      probs = llama.model(Tensor([toks[start_pos:]]), start_pos, temperature).realize()
+      probs_np = probs.numpy()
+      tok = int(np.random.choice(len(probs_np), p=probs_np))
       start_pos = len(toks)
       toks.append(tok)
 
@@ -284,28 +195,85 @@ class LLaMa:
     return output
 
 # **** main code ****
+r"""
+test:
+python3 examples/llama.py  --temperature=0 --count=50 --prompt="Hello."
+output:
+Hello. I'm a 20 year old male. I'm a student at the University of Texas at Austin. I'm a sophomore majoring in Computer Science.
 
+test:
+python3 examples/llama.py --gen='2' --temperature=0 --count=50 --prompt="Hello."
+output:
+Hello. I'm a 20 year old girl who is looking for a good lay in Palm Coast. I don't care whether it's at your place or not, as long as it's clean.
+
+test:
+python3 examples/llama.py --gen="code" --temperature=0.2 --count=50 --prompt="\
+import argparse
+
+def main(string: str):
+    print(string)
+    print(string[::-1])
+
+if __name__ == "__main__":"
+output:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('string', type=str, help='string to be reversed')
+    args = parser.parse_args()
+    main(args.string)
+
+test:
+python3 examples/llama.py --gen="code" --size="7B-Python" --temperature=0.2 --count=70 --prompt="def add_elements(arr,k):"
+output:
+    for i in range(len(arr)):
+        arr[i] += k
+    return arr
+
+
+arr = [1, 2, 3, 4, 5]
+k = 2
+print(add_elements(arr, k))
+
+test:
+python3 examples/llama.py --gen="code" --size="7B-Instruct" --temperature=0.2 --count=120 --prompt="write a function in c++ that adds three float numbers"
+output:
+\begin{code}
+#include<iostream>
+using namespace std;
+
+float add(float a, float b, float c)
+{
+    return a+b+c;
+}
+
+int main()
+{
+    float a, b, c;
+    cout<<"Enter three numbers: ";
+    cin>>a>>b>>c;
+    cout<<"The sum is: "<<add(a,b,c);
+    return 0;
+}
+\end{code}
+"""
 if __name__ == "__main__":
   Tensor.no_grad = True
   print(f"using {Device.DEFAULT} backend")
 
-  parser = argparse.ArgumentParser(description='Run LLaMA in tinygrad', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-  # test: python3 examples/llama.py --prompt="Hello." --temperature=0
-  # Hello. I'm a 20 year old male. I'm a student at the University of Texas at Austin. I'm a sophomore majoring in Computer Science.
-  # test: python3 examples/llama.py --gen 2 --prompt="Hello." --temperature=0
-  # Hello. I'm a 20 year old girl who is looking for a good lay in Palm Coast. I don't care whether it's at your place or not, as long as it's clean.
-  parser.add_argument('--prompt', type=str, default=None, help="Phrase to start with. Without this, it goes into chatbot mode")
-  parser.add_argument('--count', type=int, default=1000, help="Max number of tokens to generate")
-  parser.add_argument('--personality', type=str, default="Stacy", help="Personality, can be Stacy, George, Gary, or Lexie")
-
-  parser.add_argument('--temperature', type=float, default=0.7, help="Temperature in the softmax")
-  parser.add_argument('--timing', action='store_true', help="Print timing per token")
-  parser.add_argument('--profile', action='store_true', help="Output profile data to out.prof")
-  parser.add_argument('--size', type=str, default="7B", help="Size of model to use [7B, 13B, 30B, 65B] for Gen 1, [7B, 13B, 70B] for Gen 2")
-  parser.add_argument('--gen', type=int, default="1", help="Generation of the model to use [1, 2]")
-  parser.add_argument('--quantize', action='store_true', help="Quantize the weights to int8 in memory")
+  parser = argparse.ArgumentParser(description="Run LLaMA in tinygrad", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  parser.add_argument("--prompt", type=str, default=None, help="Phrase to start with. Without this, it goes into chatbot mode")
+  parser.add_argument("--count", type=int, default=1000, help="Max number of tokens to generate")
+  parser.add_argument("--personality", type=str, default="Stacy", help="Personality, can be Stacy, George, Gary, or Lexie")
+  parser.add_argument("--temperature", type=float, default=0.7, help="Temperature in the softmax")
+  parser.add_argument("--timing", action="store_true", help="Print timing per token")
+  parser.add_argument("--profile", action="store_true", help="Output profile data to out.prof")
+  parser.add_argument("--gen", default="1", help=f"""Generation of the model to use {list(MODEL_PARAMS.keys())}""")
+  parser.add_argument("--size", type=str, default=None, help=f"""Size of model to use {", ".join([f"{list(v.keys())} for gen '{k}'" for k, v in MODEL_PARAMS.items()])}""")
+  parser.add_argument("--quantize", action="store_true", help="Quantize the weights to int8 in memory")
+  parser.add_argument("--model", type=Path, default=None, help="Folder with the original weights to load, or single .index.json, .safetensors or .bin file")
 
   args = parser.parse_args()
+  if args.gen not in MODEL_PARAMS: raise ValueError("Invalid model generation")
+  if args.size is None: args.size = list(MODEL_PARAMS[args.gen].items())[0][0]
   chatbot = args.prompt == None
 
   # *** prompt engineers work here ****
@@ -397,12 +365,12 @@ After you are done speaking, output [EOS]. You are not Chad.
 
   # *** prompt engineers stop here ****
 
-
-  LLAMA_SUFFIX = {1: "", 2: "-2"}[args.gen]
-  WEIGHTS_DIR = Path(__file__).parent.parent / f"weights/LLaMA{LLAMA_SUFFIX}/"
-  TOKENIZER_FILENAME = WEIGHTS_DIR / "tokenizer.model"
+  LLAMA_SUFFIX = {"1": "", "2": "-2", "code": "-code", "tiny": "-tiny"}[args.gen]
+  MODEL_PATH = args.model or Path(__file__).parents[1] / f"weights/LLaMA{LLAMA_SUFFIX}/{args.size}"
+  TOKENIZER_PATH = (MODEL_PATH if MODEL_PATH.is_dir() else MODEL_PATH.parent) / "tokenizer.model"
   print(f"using LLaMA{LLAMA_SUFFIX}-{args.size} model")
-  llama = LLaMa.build(WEIGHTS_DIR, TOKENIZER_FILENAME, model_gen=args.gen, model_size=args.size, quantize=args.quantize)
+  llama = LLaMa.build(MODEL_PATH, TOKENIZER_PATH, model_gen=args.gen, model_size=args.size, quantize=args.quantize)
+  param_count = sum(x.lazydata.size for x in get_parameters(llama.model))
 
   if chatbot:
     # encode pre prompt
@@ -410,7 +378,7 @@ After you are done speaking, output [EOS]. You are not Chad.
 
     print(f"Preparing KV cache for chatbot with personality {args.personality}...")
     with Timing():
-      llama.model(Tensor([toks]), 0).realize()  # NOTE: output logits are not used
+      llama.model(Tensor([toks]), 0, args.temperature).realize()  # NOTE: outputs are not used
     start_pos = len(toks)
   else:
     # non chat bot mode
@@ -422,10 +390,6 @@ After you are done speaking, output [EOS]. You are not Chad.
   sys.stdout.write(outputted)
   sys.stdout.flush()
 
-  if args.profile:
-    import cProfile, pstats
-    profiler = cProfile.Profile()
-
   # chatbot loop
   while 1:
     # add tokens from user in chatbot mode
@@ -441,14 +405,15 @@ After you are done speaking, output [EOS]. You are not Chad.
     last_break = len(outputted)
     for i in range(args.count):
       GlobalCounters.reset()
-      if args.profile and i == 2: profiler.enable()
 
-      if args.timing: print("")
+      if args.timing or args.profile: print("")
       st = GlobalCounters.time_sum_s
-      with Timing("ran model in ", on_exit=(lambda et: f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU") if DEBUG else None, enabled=args.timing):
-        logits = llama.model(Tensor([toks[start_pos:]]), start_pos).realize()[:, -1, :]
-      with Timing("sync in ", enabled=args.timing):
-        tok = sample(logits, args.temperature)
+      with Profiling(enabled=args.profile):
+        with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/sec"):
+          with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
+                      f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
+                      (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_count*1e-9*2/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
+            tok = llama.model(Tensor([toks[start_pos:]]), start_pos, args.temperature).item()
 
       # use the kv cache
       start_pos = len(toks)
@@ -466,7 +431,16 @@ After you are done speaking, output [EOS]. You are not Chad.
       if chatbot and outputted.endswith(end_delim): break
     if not chatbot: break
 
-  if args.profile:
-    profiler.disable()
-    stats = pstats.Stats(profiler)
-    stats.dump_stats('out.prof')
+  # validate output!
+  if args.temperature == 0 and args.count == 10 and args.prompt == "Hello." and not args.quantize:
+    text = llama.tokenizer.decode(toks)
+    key = (args.gen, args.size)
+    expected = {
+      ("1", "7B"): "Hello. I'm a 20 year old male",
+      ("2", "7B"): "Hello. I'm a 20 year old girl",
+    }
+    try:
+      assert text == expected[key], "invalid output: " + colored(text, "red")
+      print("\n" + colored("output validated", "green"))  # NOTE: "\n" iside colored does not render the color in github action
+    except KeyError:
+      pass

examples/mask_rcnn.py

@@ -1,6 +1,6 @@
-from models.mask_rcnn import MaskRCNN
-from models.resnet import ResNet
-from models.mask_rcnn import BoxList
+from extra.models.mask_rcnn import MaskRCNN
+from extra.models.resnet import ResNet
+from extra.models.mask_rcnn import BoxList
 from torch.nn import functional as F
 from torchvision import transforms as T
 from torchvision.transforms import functional as Ft

examples/mixtral.py

@@ -0,0 +1,57 @@
+import functools, argparse, pathlib
+from tqdm import tqdm
+from tinygrad import Tensor, nn, Device, GlobalCounters, Variable
+from tinygrad.helpers import Timing
+from tinygrad.nn.state import torch_load, get_state_dict
+from extra.models.llama import FeedForward, Transformer
+
+class MixtureFeedForward:
+  def __init__(self, num_experts:int, dim:int, hidden_dim:int, linear=nn.Linear):
+    self.gate = nn.Linear(dim, num_experts, bias=False)
+    self.experts = [FeedForward(dim, hidden_dim, linear) for _ in range(num_experts)]
+  def __call__(self, x:Tensor) -> Tensor:
+    assert x.shape[0] == 1, "only BS=1"
+    g = self.gate(x).float().exp()
+    choice = g.data().tolist()[0][0]
+    top = sorted(enumerate(choice), key=lambda x: -x[1])
+    norm = top[0][1] + top[1][1]
+    e1, e2 = self.experts[top[0][0]], self.experts[top[1][0]]
+    ret = e1(x.to(e1.w1.weight.device)).to(x.device) * Tensor([top[0][1]/norm]) + \
+          e2(x.to(e2.w1.weight.device)).to(x.device) * Tensor([top[1][1]/norm])
+    return ret
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(description="Run Mixtral in tinygrad", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  parser.add_argument("--count", type=int, default=30, help="Max number of tokens to generate")
+  parser.add_argument("--temperature", type=float, default=0.7, help="Temperature in the softmax")
+  parser.add_argument("--timing", action="store_true", help="Print timing per token")
+  parser.add_argument("--weights", type=str, default=(pathlib.Path(__file__).parent.parent / "weights/mixtral-8x7b-32kseqlen").as_posix(),
+                      help="Path to the downloaded weights")
+  args = parser.parse_args()
+
+  state = torch_load(args.weights + "/consolidated.00.pth.b")
+  model = Transformer(n_layers=32, dim=4096, hidden_dim=14336, n_heads=32, n_kv_heads=8, norm_eps=1e-5, vocab_size=32000, feed_forward=functools.partial(MixtureFeedForward, 8), jit=False)
+  model_state_dict = get_state_dict(model)
+
+  for k in (t := tqdm(state)):
+    if 'feed_forward.experts.' in k:
+      expert_no = int(k.split('feed_forward.experts.')[1].split('.')[0])
+      device = Device.DEFAULT + ":" + str((expert_no//2)+1)
+    else:
+      device = Device.DEFAULT
+    t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, loading {k} to {device}")
+    # NOTE: we have to copy through CLANG to avoid the HIP hang bug when copying directly from the DISK
+    model_state_dict[k].assign(state[k].to("CLANG").contiguous().to(device).half()).realize()
+
+  from sentencepiece import SentencePieceProcessor
+  spp = SentencePieceProcessor(model_file=args.weights + "/tokenizer.model")
+
+  toks = [spp.bos_id()]
+  start_pos = 0
+  for i in range(args.count):
+    GlobalCounters.reset()
+    with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/sec"):
+      tok = model(Tensor([toks[start_pos:]]), 0 if start_pos == 0 else Variable("start_pos", 1, 1024).bind(start_pos), args.temperature).item()
+    toks.append(tok)
+    start_pos += 1
+    print(spp.decode(toks))

examples/mlperf/model_eval.py

@@ -1,15 +1,15 @@
 import time
 from pathlib import Path
 import numpy as np
-from tinygrad.tensor import Tensor
+from tinygrad import Tensor, GlobalCounters, dtypes
 from tinygrad.jit import TinyJit
-from tinygrad.helpers import getenv, dtypes
+from tinygrad.helpers import getenv
 from examples.mlperf import helpers
 
 def eval_resnet():
   # Resnet50-v1.5
   from tinygrad.jit import TinyJit
-  from models.resnet import ResNet50
+  from extra.models.resnet import ResNet50
   mdl = ResNet50()
   mdl.load_from_pretrained()
 
@@ -26,16 +26,16 @@ def eval_resnet():
 
   # evaluation on the mlperf classes of the validation set from imagenet
   from extra.datasets.imagenet import iterate
-  from extra.helpers import cross_process
 
   BS = 64
   n,d = 0,0
   st = time.perf_counter()
-  iterator = cross_process(lambda: iterate(BS))
+  iterator = iterate(BS)
   x,ny = next(iterator)
   dat = Tensor(x)
   while dat is not None:
     y = ny
+    GlobalCounters.reset()
     mt = time.perf_counter()
     outs = mdlrun(dat) if dat.shape[0] != BS else mdljit(dat)
     try:
@@ -45,17 +45,14 @@ def eval_resnet():
       dat = None
     t = outs.argmax(axis=1).numpy()
     et = time.perf_counter()
-    print(f"{(mt-st)*1000:.2f} ms loading data, {(et-mt)*1000:.2f} ms to run model")
-    print(t)
-    print(y)
     n += (t==y).sum()
     d += len(t)
-    print(f"****** {n}/{d}  {n*100.0/d:.2f}%")
+    print(f"****** {n}/{d}  {n*100.0/d:.2f}% -- {(mt-st)*1000:.2f} ms loading data, {(et-mt)*1000:7.2f} ms to run model. {len(t)/(et-mt):.2f} examples/sec. {GlobalCounters.global_ops*1e-12/(et-mt):.2f} TFLOPS")
     st = time.perf_counter()
 
 def eval_unet3d():
   # UNet3D
-  from models.unet3d import UNet3D
+  from extra.models.unet3d import UNet3D
   from extra.datasets.kits19 import iterate, sliding_window_inference
   from examples.mlperf.metrics import get_dice_score
   mdl = UNet3D()
@@ -73,8 +70,8 @@ def eval_unet3d():
 
 def eval_retinanet():
   # RetinaNet with ResNeXt50_32X4D
-  from models.resnet import ResNeXt50_32X4D
-  from models.retinanet import RetinaNet
+  from extra.models.resnet import ResNeXt50_32X4D
+  from extra.models.retinanet import RetinaNet
   mdl = RetinaNet(ResNeXt50_32X4D())
   mdl.load_from_pretrained()
 
@@ -131,7 +128,7 @@ def eval_retinanet():
 
 def eval_rnnt():
   # RNN-T
-  from models.rnnt import RNNT
+  from extra.models.rnnt import RNNT
   mdl = RNNT()
   mdl.load_from_pretrained()
 
@@ -160,7 +157,7 @@ def eval_rnnt():
 
 def eval_bert():
   # Bert-QA
-  from models.bert import BertForQuestionAnswering
+  from extra.models.bert import BertForQuestionAnswering
   mdl = BertForQuestionAnswering()
   mdl.load_from_pretrained()
 
@@ -173,7 +170,7 @@ def eval_bert():
   from examples.mlperf.metrics import f1_score
   from transformers import BertTokenizer
 
-  tokenizer = BertTokenizer(str(Path(__file__).parent.parent.parent / "weights/bert_vocab.txt"))
+  tokenizer = BertTokenizer(str(Path(__file__).parents[2] / "weights/bert_vocab.txt"))
 
   c = 0
   f1 = 0.0
@@ -196,8 +193,8 @@ def eval_bert():
 
 def eval_mrcnn():
   from tqdm import tqdm
-  from models.mask_rcnn import MaskRCNN
-  from models.resnet import ResNet
+  from extra.models.mask_rcnn import MaskRCNN
+  from extra.models.resnet import ResNet
   from extra.datasets.coco import BASEDIR, images, convert_prediction_to_coco_bbox, convert_prediction_to_coco_mask, accumulate_predictions_for_coco, evaluate_predictions_on_coco, iterate
   from examples.mask_rcnn import compute_prediction_batched, Image
   mdl = MaskRCNN(ResNet(50, num_classes=None, stride_in_1x1=True))

examples/mlperf/model_spec.py

@@ -1,6 +1,6 @@
 # load each model here, quick benchmark
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import GlobalCounters, getenv
+from tinygrad import Tensor, GlobalCounters
+from tinygrad.helpers import getenv
 import numpy as np
 
 def test_model(model, *inputs):
@@ -12,48 +12,48 @@ def test_model(model, *inputs):
 
 def spec_resnet():
   # Resnet50-v1.5
-  from models.resnet import ResNet50
+  from extra.models.resnet import ResNet50
   mdl = ResNet50()
   img = Tensor.randn(1, 3, 224, 224)
   test_model(mdl, img)
 
 def spec_retinanet():
   # Retinanet with ResNet backbone
-  from models.resnet import ResNet50
-  from models.retinanet import RetinaNet
+  from extra.models.resnet import ResNet50
+  from extra.models.retinanet import RetinaNet
   mdl = RetinaNet(ResNet50(), num_classes=91, num_anchors=9)
   img = Tensor.randn(1, 3, 224, 224)
   test_model(mdl, img)
 
 def spec_unet3d():
   # 3D UNET
-  from models.unet3d import UNet3D
+  from extra.models.unet3d import UNet3D
   mdl = UNet3D()
-  mdl.load_from_pretrained()
+  #mdl.load_from_pretrained()
   img = Tensor.randn(1, 1, 128, 128, 128)
   test_model(mdl, img)
 
 def spec_rnnt():
-  from models.rnnt import RNNT
+  from extra.models.rnnt import RNNT
   mdl = RNNT()
-  mdl.load_from_pretrained()
+  #mdl.load_from_pretrained()
   x = Tensor.randn(220, 1, 240)
   y = Tensor.randn(1, 220)
   test_model(mdl, x, y)
 
 def spec_bert():
-  from models.bert import BertForQuestionAnswering
+  from extra.models.bert import BertForQuestionAnswering
   mdl = BertForQuestionAnswering()
-  mdl.load_from_pretrained()
+  #mdl.load_from_pretrained()
   x = Tensor.randn(1, 384)
   am = Tensor.randn(1, 384)
   tt = Tensor(np.random.randint(0, 2, (1, 384)).astype(np.float32))
   test_model(mdl, x, am, tt)
 
 def spec_mrcnn():
-  from models.mask_rcnn import MaskRCNN, ResNet
+  from extra.models.mask_rcnn import MaskRCNN, ResNet
   mdl = MaskRCNN(ResNet(50, num_classes=None, stride_in_1x1=True))
-  mdl.load_from_pretrained()
+  #mdl.load_from_pretrained()
   x = Tensor.randn(3, 224, 224)
   test_model(mdl, [x])
 

examples/mlperf/model_train.py

@@ -26,12 +26,11 @@ def train_maskrcnn():
   pass
 
 if __name__ == "__main__":
-  Tensor.training = True
-
-  for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn").split(","):
-    nm = f"train_{m}"
-    if nm in globals():
-      print(f"training {m}")
-      globals()[nm]()
+  with Tensor.train():
+    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn").split(","):
+      nm = f"train_{m}"
+      if nm in globals():
+        print(f"training {m}")
+        globals()[nm]()
 
 

examples/rl/lightupbutton.py

@@ -0,0 +1,45 @@
+import gymnasium as gym
+import numpy as np
+from gymnasium.envs.registration import register
+
+# a very simple game
+# one of <size> lights will light up
+# take the action of the lit up light
+# in <hard_mode>, you act differently based on the step number and need to track this
+
+class PressTheLightUpButton(gym.Env):
+  metadata = {"render_modes": []}
+  def __init__(self, render_mode=None, size=2, game_length=10, hard_mode=False):
+    self.size, self.game_length = size, game_length
+    self.observation_space = gym.spaces.Box(0, 1, shape=(self.size,), dtype=np.float32)
+    self.action_space = gym.spaces.Discrete(self.size)
+    self.step_num = 0
+    self.done = True
+    self.hard_mode = hard_mode
+
+  def _get_obs(self):
+    obs = [0]*self.size
+    if self.step_num < len(self.state):
+      obs[self.state[self.step_num]] = 1
+    return np.array(obs, dtype=np.float32)
+
+  def reset(self, seed=None, options=None):
+    super().reset(seed=seed)
+    self.state = np.random.randint(0, self.size, size=self.game_length)
+    self.step_num = 0
+    self.done = False
+    return self._get_obs(), {}
+
+  def step(self, action):
+    target = ((action + self.step_num) % self.size) if self.hard_mode else action
+    reward = int(target == self.state[self.step_num])
+    self.step_num += 1
+    if not reward:
+      self.done = True
+    return self._get_obs(), reward, self.done, self.step_num >= self.game_length, {}
+
+register(
+  id="PressTheLightUpButton-v0",
+  entry_point="examples.rl.lightupbutton:PressTheLightUpButton",
+  max_episode_steps=None,
+)

examples/simple_conv_bn.py

@@ -5,15 +5,15 @@ from tinygrad.nn import Conv2d, BatchNorm2d
 from tinygrad.nn.state import get_parameters
 
 if __name__ == "__main__":
-  Tensor.training = True
+  with Tensor.train():
 
-  BS, C1, H, W = 4, 16, 224, 224
-  C2, K, S, P = 64, 7, 2, 1
+    BS, C1, H, W = 4, 16, 224, 224
+    C2, K, S, P = 64, 7, 2, 1
 
-  x = Tensor.uniform(BS, C1, H, W)
-  conv = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
-  bn = BatchNorm2d(C2, track_running_stats=False)
-  for t in get_parameters([x, conv, bn]): t.realize()
+    x = Tensor.uniform(BS, C1, H, W)
+    conv = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
+    bn = BatchNorm2d(C2, track_running_stats=False)
+    for t in get_parameters([x, conv, bn]): t.realize()
 
-  print("running network")
-  x.sequential([conv, bn]).numpy()
+    print("running network")
+    x.sequential([conv, bn]).numpy()

examples/so_vits_svc.py

@@ -1,12 +1,11 @@
 # original implementation: https://github.com/svc-develop-team/so-vits-svc
 from __future__ import annotations
-import sys, os, logging, time, io, math, argparse, operator, numpy as np
+import sys, logging, time, io, math, argparse, operator, numpy as np
 from functools import partial, reduce
 from pathlib import Path
 from typing import Tuple, Optional, Type
-from tinygrad import nn
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import dtypes, getenv
+from tinygrad import nn, dtypes, Tensor
+from tinygrad.helpers import getenv
 from tinygrad.nn.state import torch_load
 from examples.vits import ResidualCouplingBlock, PosteriorEncoder, Encoder, ResBlock1, ResBlock2, LRELU_SLOPE, sequence_mask, split, download_if_not_present, get_hparams_from_file, load_checkpoint, weight_norm, HParams
 from examples.sovits_helpers import preprocess
@@ -210,7 +209,7 @@ class MultiHeadAttention:
     ret =  self.out_proj(wv).transpose(0,1)  # BxTxC -> TxBxC
     return ret
 
-class ConvFeatureExtractionModel():
+class ConvFeatureExtractionModel:
   def __init__(self, conv_layers, dropout=.0, mode="default", conv_bias=False):
     assert mode in {"default", "group_norm_masked", "layer_norm"}
     def block(n_in, n_out, k, stride, is_layer_norm=False, is_group_norm=False, conv_bias=False):
@@ -352,7 +351,7 @@ class Upsample:
     new_shape = (*x.shape[:-1], x.shape[-1] * self.scale)
     return x.unsqueeze(-1).repeat(repeats).reshape(new_shape)
 
-class SineGen():
+class SineGen:
   def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voice_threshold=0, flag_for_pulse=False):
     self.sine_amp, self.noise_std, self.harmonic_num, self.sampling_rate, self.voiced_threshold, self.flag_for_pulse = sine_amp, noise_std, harmonic_num, samp_rate, voice_threshold, flag_for_pulse
     self.dim = self.harmonic_num + 1
@@ -468,14 +467,14 @@ def repeat_expand_2d_left(content, target_len): # content : [h, t]
   return Tensor.stack(cols).transpose(0, 1)
 
 def load_fairseq_cfg(checkpoint_path):
-  assert os.path.isfile(checkpoint_path)
+  assert Path(checkpoint_path).is_file()
   state = torch_load(checkpoint_path)
   cfg = state["cfg"] if ("cfg" in state and state["cfg"] is not None) else None
   if cfg is None: raise RuntimeError(f"No cfg exist in state keys = {state.keys()}")
   return HParams(**cfg)
 
 def load_checkpoint_enc(checkpoint_path, model: ContentVec, optimizer=None, skip_list=[]):
-  assert os.path.isfile(checkpoint_path)
+  assert Path(checkpoint_path).is_file()
   start_time = time.time()
   checkpoint_dict = torch_load(checkpoint_path)
   saved_state_dict = checkpoint_dict['model']
@@ -550,7 +549,7 @@ def get_encoder(ssl_dim) -> Type[SpeechEncoder]:
 # DEMO USAGE (uses audio sample from LJ-Speech):
 # python3 examples/so_vits_svc.py --model saul_goodman
 #########################################################################################
-SO_VITS_SVC_PATH = Path(__file__).parent.parent / "weights/So-VITS-SVC"
+SO_VITS_SVC_PATH = Path(__file__).parents[1] / "weights/So-VITS-SVC"
 VITS_MODELS = { # config_path, weights_path, config_url, weights_url
   "saul_goodman" : (SO_VITS_SVC_PATH / "config_saul_gman.json", SO_VITS_SVC_PATH / "pretrained_saul_gman.pth", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/Saul_Goodman_80000/config.json", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/Saul_Goodman_80000/G_80000.pth"),
   "drake" : (SO_VITS_SVC_PATH / "config_drake.json", SO_VITS_SVC_PATH / "pretrained_drake.pth", "https://huggingface.co/jaspa/so-vits-svc/resolve/main/aubrey/config_aubrey.json", "https://huggingface.co/jaspa/so-vits-svc/resolve/main/aubrey/pretrained_aubrey.pth"),
@@ -563,13 +562,13 @@ ENCODER_MODELS = { # weights_path, weights_url
   "contentvec": (SO_VITS_SVC_PATH / "contentvec_checkpoint.pt", "https://huggingface.co/lj1995/VoiceConversionWebUI/resolve/main/hubert_base.pt")
 }
 ENCODER_MODEL = "contentvec"
-DEMO_PATH, DEMO_URL = Path(__file__).parent.parent / "temp/LJ037-0171.wav", "https://keithito.com/LJ-Speech-Dataset/LJ037-0171.wav"
+DEMO_PATH, DEMO_URL = Path(__file__).parents[1] / "temp/LJ037-0171.wav", "https://keithito.com/LJ-Speech-Dataset/LJ037-0171.wav"
 if __name__=="__main__":
   logging.basicConfig(stream=sys.stdout, level=(logging.INFO if DEBUG < 1 else logging.DEBUG))
   parser = argparse.ArgumentParser()
   parser.add_argument("-m", "--model", default=None, help=f"Specify the model to use. All supported models: {VITS_MODELS.keys()}", required=True)
   parser.add_argument("-f", "--file", default=DEMO_PATH, help=f"Specify the path of the input file")
-  parser.add_argument("--out_dir", default=str(Path(__file__).parent.parent / "temp"), help="Specify the output path.")
+  parser.add_argument("--out_dir", default=str(Path(__file__).parents[1] / "temp"), help="Specify the output path.")
   parser.add_argument("--out_path", default=None, help="Specify the full output path. Overrides the --out_dir and --name parameter.")
   parser.add_argument("--base_name", default="test", help="Specify the base of the output file name. Default is 'test'.")
   parser.add_argument("--speaker", default=None, help="If not specified, the first available speaker is chosen. Usually there is only one speaker per model.")
@@ -600,7 +599,7 @@ if __name__=="__main__":
 
   ### Loading audio and slicing ###
   if audio_path == DEMO_PATH: download_if_not_present(DEMO_PATH, DEMO_URL)
-  assert os.path.isfile(audio_path) and Path(audio_path).suffix == ".wav"
+  assert Path(audio_path).is_file() and Path(audio_path).suffix == ".wav"
   chunks = preprocess.cut(audio_path, db_thresh=slice_db)
   audio_data, audio_sr = preprocess.chunks2audio(audio_path, chunks)
 

examples/sovits_helpers/preprocess.py

@@ -1,7 +1,6 @@
 import math
 from typing import Optional, Tuple
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import dtypes
+from tinygrad import Tensor, dtypes
 import librosa
 import soundfile
 import numpy as np

examples/stable_diffusion.py

@@ -1,18 +1,19 @@
 # https://arxiv.org/pdf/2112.10752.pdf
 # https://github.com/ekagra-ranjan/huggingface-blog/blob/main/stable_diffusion.md
-import os
 import tempfile
 from pathlib import Path
 import gzip, argparse, math, re
 from functools import lru_cache
 from collections import namedtuple
 
+from PIL import Image
+import numpy as np
 from tqdm import tqdm
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import dtypes, GlobalCounters
+from tinygrad import Device, GlobalCounters, dtypes, Tensor
+from tinygrad.helpers import Timing, Context, getenv, fetch, colored
 from tinygrad.nn import Conv2d, Linear, GroupNorm, LayerNorm, Embedding
-from extra.utils import download_file
 from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
+from tinygrad.jit import TinyJit
 
 class AttnBlock:
   def __init__(self, in_channels):
@@ -174,7 +175,7 @@ class CrossAttention:
     q,k,v = self.to_q(x), self.to_k(context), self.to_v(context)
     q,k,v = [y.reshape(x.shape[0], -1, self.num_heads, self.head_size).transpose(1,2) for y in (q,k,v)]
     attention = Tensor.scaled_dot_product_attention(q, k, v).transpose(1,2)
-    h_ = attention.reshape(shape=(x.shape[0], -1, self.num_heads * self.head_size))
+    h_ = attention.reshape(x.shape[0], -1, self.num_heads * self.head_size)
     return h_.sequential(self.to_out)
 
 class GEGLU:
@@ -250,7 +251,8 @@ class Upsample:
 
 def timestep_embedding(timesteps, dim, max_period=10000):
   half = dim // 2
-  freqs = (-math.log(max_period) * Tensor.arange(half) / half).exp()
+  # TODO: remove explicit dtypes after broadcast fix
+  freqs = (-math.log(max_period) * Tensor.arange(half, dtype=dtypes.float32) / half).exp()
   args = timesteps * freqs
   return Tensor.cat(args.cos(), args.sin()).reshape(1, -1)
 
@@ -347,29 +349,12 @@ class CLIPAttention:
     self.q_proj = Linear(self.embed_dim, self.embed_dim)
     self.out_proj = Linear(self.embed_dim, self.embed_dim)
 
-  def _shape(self, tensor, seq_len: int, bsz: int):
-    return tensor.reshape(bsz, seq_len, self.num_heads, self.head_dim).permute(0,2,1,3)
-
   def __call__(self, hidden_states, causal_attention_mask):
     bsz, tgt_len, embed_dim = hidden_states.shape
-
-    query_states = self.q_proj(hidden_states)
-    key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
-    value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
-
-    proj_shape = (bsz * self.num_heads, -1, self.head_dim)
-    query_states = self._shape(query_states, tgt_len, bsz).reshape(*proj_shape)
-    key_states = key_states.reshape(*proj_shape)
-    src_len = key_states.shape[1]
-    value_states = value_states.reshape(*proj_shape)
-
-    attn_output = Tensor.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask=causal_attention_mask)
-    attn_output = attn_output.reshape(bsz, self.num_heads, tgt_len, self.head_dim)
-    attn_output = attn_output.permute(0,2,1,3)
-    attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)
-
-    attn_output = self.out_proj(attn_output)
-    return attn_output
+    q,k,v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
+    q,k,v = [x.reshape(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2) for x in (q,k,v)]
+    attn_output = Tensor.scaled_dot_product_attention(q, k, v, attn_mask=causal_attention_mask)
+    return self.out_proj(attn_output.transpose(1, 2).reshape(bsz, tgt_len, embed_dim))
 
 class CLIPEncoderLayer:
   def __init__(self):
@@ -421,10 +406,7 @@ class CLIPTextTransformer:
 
 # Clip tokenizer, taken from https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py (MIT license)
 @lru_cache()
-def default_bpe():
-  fn = Path(__file__).parent.parent / "weights/bpe_simple_vocab_16e6.txt.gz"
-  download_file("https://github.com/openai/CLIP/raw/main/clip/bpe_simple_vocab_16e6.txt.gz", fn)
-  return fn
+def default_bpe(): return fetch("https://github.com/openai/CLIP/raw/main/clip/bpe_simple_vocab_16e6.txt.gz", "bpe_simple_vocab_16e6.txt.gz")
 
 def get_pairs(word):
   """Return set of symbol pairs in a word.
@@ -537,7 +519,44 @@ class StableDiffusion:
     self.first_stage_model = AutoencoderKL()
     self.cond_stage_model = namedtuple("CondStageModel", ["transformer"])(transformer = namedtuple("Transformer", ["text_model"])(text_model = CLIPTextTransformer()))
 
-  # TODO: make __call__ run the model
+  def get_x_prev_and_pred_x0(self, x, e_t, a_t, a_prev):
+    temperature = 1
+    sigma_t = 0
+    sqrt_one_minus_at = (1-a_t).sqrt()
+    #print(a_t, a_prev, sigma_t, sqrt_one_minus_at)
+
+    pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+
+    # direction pointing to x_t
+    dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
+
+    x_prev = a_prev.sqrt() * pred_x0 + dir_xt
+    return x_prev, pred_x0
+
+  def get_model_output(self, unconditional_context, context, latent, timestep, unconditional_guidance_scale):
+    # put into diffuser
+    latents = self.model.diffusion_model(latent.expand(2, *latent.shape[1:]), timestep, unconditional_context.cat(context, dim=0))
+    unconditional_latent, latent = latents[0:1], latents[1:2]
+
+    e_t = unconditional_latent + unconditional_guidance_scale * (latent - unconditional_latent)
+    return e_t
+
+  def decode(self, x):
+    x = self.first_stage_model.post_quant_conv(1/0.18215 * x)
+    x = self.first_stage_model.decoder(x)
+
+    # make image correct size and scale
+    x = (x + 1.0) / 2.0
+    x = x.reshape(3,512,512).permute(1,2,0).clip(0,1)*255
+    return x.cast(dtypes.uint8) if Device.DEFAULT != "WEBGPU" else x
+
+  def __call__(self, unconditional_context, context, latent, timestep, alphas, alphas_prev, guidance):
+    e_t = self.get_model_output(unconditional_context, context, latent, timestep, guidance)
+    x_prev, _ = self.get_x_prev_and_pred_x0(latent, e_t, alphas, alphas_prev)
+    #e_t_next = get_model_output(x_prev)
+    #e_t_prime = (e_t + e_t_next) / 2
+    #x_prev, pred_x0 = get_x_prev_and_pred_x0(latent, e_t_prime, index)
+    return x_prev.realize()
 
 # ** ldm.models.autoencoder.AutoencoderKL (done!)
 # 3x512x512 <--> 4x64x64 (16384)
@@ -555,24 +574,24 @@ class StableDiffusion:
 # ** ldm.modules.encoders.modules.FrozenCLIPEmbedder
 # cond_stage_model.transformer.text_model
 
-# this is sd-v1-4.ckpt
-FILENAME = Path(__file__).parent.parent / "weights/sd-v1-4.ckpt"
-
 if __name__ == "__main__":
+  default_prompt = "a horse sized cat eating a bagel"
   parser = argparse.ArgumentParser(description='Run Stable Diffusion', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
   parser.add_argument('--steps', type=int, default=5, help="Number of steps in diffusion")
-  parser.add_argument('--prompt', type=str, default="a horse sized cat eating a bagel", help="Phrase to render")
-  parser.add_argument('--out', type=str, default=os.path.join(tempfile.gettempdir(), "rendered.png"), help="Output filename")
+  parser.add_argument('--prompt', type=str, default=default_prompt, help="Phrase to render")
+  parser.add_argument('--out', type=str, default=Path(tempfile.gettempdir()) / "rendered.png", help="Output filename")
   parser.add_argument('--noshow', action='store_true', help="Don't show the image")
   parser.add_argument('--fp16', action='store_true', help="Cast the weights to float16")
+  parser.add_argument('--timing', action='store_true', help="Print timing per step")
+  parser.add_argument('--seed', type=int, help="Set the random latent seed")
+  parser.add_argument('--guidance', type=float, default=7.5, help="Prompt strength")
   args = parser.parse_args()
 
   Tensor.no_grad = True
   model = StableDiffusion()
 
   # load in weights
-  download_file('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', FILENAME)
-  load_state_dict(model, torch_load(FILENAME)['state_dict'], strict=False)
+  load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], strict=False)
 
   if args.fp16:
     for l in get_state_dict(model).values():
@@ -591,64 +610,43 @@ if __name__ == "__main__":
   # done with clip model
   del model.cond_stage_model
 
-  def get_model_output(latent, timestep):
-    # put into diffuser
-    latents = model.model.diffusion_model(latent.expand(2, *latent.shape[1:]), timestep.expand(2, *timestep.shape[1:]), unconditional_context.cat(context, dim=0))
-    unconditional_latent, latent = latents[0:1], latents[1:2]
-
-    unconditional_guidance_scale = 7.5
-    e_t = unconditional_latent + unconditional_guidance_scale * (latent - unconditional_latent)
-    return e_t
-
   timesteps = list(range(1, 1000, 1000//args.steps))
   print(f"running for {timesteps} timesteps")
-  alphas = [model.alphas_cumprod.numpy()[t] for t in timesteps]
-  alphas_prev = [1.0] + alphas[:-1]
-
-  def get_x_prev_and_pred_x0(x, e_t, index):
-    temperature = 1
-    a_t, a_prev = alphas[index], alphas_prev[index]
-    sigma_t = 0
-    sqrt_one_minus_at = math.sqrt(1-a_t)
-    #print(a_t, a_prev, sigma_t, sqrt_one_minus_at)
-
-    pred_x0 = (x - sqrt_one_minus_at * e_t) / math.sqrt(a_t)
-
-    # direction pointing to x_t
-    dir_xt = math.sqrt(1. - a_prev - sigma_t**2) * e_t
-    noise = sigma_t * Tensor.randn(*x.shape) * temperature
-
-    x_prev = math.sqrt(a_prev) * pred_x0 + dir_xt #+ noise
-    return x_prev, pred_x0
+  alphas = model.alphas_cumprod[Tensor(timesteps)]
+  alphas_prev = Tensor([1.0]).cat(alphas[:-1])
 
   # start with random noise
+  if args.seed is not None: Tensor._seed = args.seed
   latent = Tensor.randn(1,4,64,64)
 
+  @TinyJit
+  def run(model, *x): return model(*x).realize()
+
   # this is diffusion
-  for index, timestep in (t:=tqdm(list(enumerate(timesteps))[::-1])):
-    GlobalCounters.reset()
-    t.set_description("%3d %3d" % (index, timestep))
-    e_t = get_model_output(latent, Tensor([timestep]))
-    x_prev, pred_x0 = get_x_prev_and_pred_x0(latent, e_t, index)
-    #e_t_next = get_model_output(x_prev)
-    #e_t_prime = (e_t + e_t_next) / 2
-    #x_prev, pred_x0 = get_x_prev_and_pred_x0(latent, e_t_prime, index)
-    latent = x_prev
-    latent.realize()
+  with Context(BEAM=getenv("LATEBEAM")):
+    for index, timestep in (t:=tqdm(list(enumerate(timesteps))[::-1])):
+      GlobalCounters.reset()
+      t.set_description("%3d %3d" % (index, timestep))
+      with Timing("step in ", enabled=args.timing, on_exit=lambda _: f", using {GlobalCounters.mem_used/1e9:.2f} GB"):
+        tid = Tensor([index])
+        latent = run(model, unconditional_context, context, latent, Tensor([timestep]), alphas[tid], alphas_prev[tid], Tensor([args.guidance]))
+        if args.timing: Device[Device.DEFAULT].synchronize()
+    del run
 
   # upsample latent space to image with autoencoder
-  x = model.first_stage_model.post_quant_conv(1/0.18215 * latent)
-  x = model.first_stage_model.decoder(x)
-
-  # make image correct size and scale
-  x = (x + 1.0) / 2.0
-  x = (x.reshape(3,512,512).permute(1,2,0).clip(0,1)*255).cast(dtypes.uint8)
+  x = model.decode(latent)
   print(x.shape)
 
   # save image
-  from PIL import Image
-  im = Image.fromarray(x.numpy())
+  im = Image.fromarray(x.numpy().astype(np.uint8, copy=False))
   print(f"saving {args.out}")
   im.save(args.out)
   # Open image.
   if not args.noshow: im.show()
+
+  # validation!
+  if args.prompt == default_prompt and args.steps == 5 and args.seed == 0 and args.guidance == 7.5:
+    ref_image = Tensor(np.array(Image.open(Path(__file__).parent / "stable_diffusion_seed0.png")))
+    distance = (((x - ref_image).cast(dtypes.float) / ref_image.max())**2).mean().item()
+    assert distance < 3e-4, f"validation failed with {distance=}"
+    print(colored(f"output validated with {distance=}", "green"))

examples/train_efficientnet.py

@@ -8,7 +8,7 @@ from tinygrad.nn import optim
 from tinygrad.helpers import getenv
 from tinygrad.tensor import Tensor
 from extra.datasets import fetch_cifar
-from models.efficientnet import EfficientNet
+from extra.models.efficientnet import EfficientNet
 
 class TinyConvNet:
   def __init__(self, classes=10):
@@ -59,45 +59,47 @@ if __name__ == "__main__":
       p.daemon = True
       p.start()
   else:
-    X_train, Y_train = fetch_cifar()
+    X_train, Y_train, _, _ = fetch_cifar()
+    X_train = X_train.reshape((-1, 3, 32, 32))
+    Y_train = Y_train.reshape((-1,))
 
-  Tensor.training = True
-  for i in (t := trange(steps)):
-    if IMAGENET:
-      X, Y = q.get(True)
-    else:
-      samp = np.random.randint(0, X_train.shape[0], size=(BS))
-      X, Y = X_train[samp], Y_train[samp]
+  with Tensor.train():
+    for i in (t := trange(steps)):
+      if IMAGENET:
+        X, Y = q.get(True)
+      else:
+        samp = np.random.randint(0, X_train.shape[0], size=(BS))
+        X, Y = X_train.numpy()[samp], Y_train.numpy()[samp]
 
-    st = time.time()
-    out = model.forward(Tensor(X.astype(np.float32), requires_grad=False))
-    fp_time = (time.time()-st)*1000.0
+      st = time.time()
+      out = model.forward(Tensor(X.astype(np.float32), requires_grad=False))
+      fp_time = (time.time()-st)*1000.0
 
-    y = np.zeros((BS,classes), np.float32)
-    y[range(y.shape[0]),Y] = -classes
-    y = Tensor(y, requires_grad=False)
-    loss = out.log_softmax().mul(y).mean()
+      y = np.zeros((BS,classes), np.float32)
+      y[range(y.shape[0]),Y] = -classes
+      y = Tensor(y, requires_grad=False)
+      loss = out.log_softmax().mul(y).mean()
 
-    optimizer.zero_grad()
+      optimizer.zero_grad()
 
-    st = time.time()
-    loss.backward()
-    bp_time = (time.time()-st)*1000.0
+      st = time.time()
+      loss.backward()
+      bp_time = (time.time()-st)*1000.0
 
-    st = time.time()
-    optimizer.step()
-    opt_time = (time.time()-st)*1000.0
+      st = time.time()
+      optimizer.step()
+      opt_time = (time.time()-st)*1000.0
 
-    st = time.time()
-    loss = loss.numpy()
-    cat = out.argmax(axis=1).numpy()
-    accuracy = (cat == Y).mean()
-    finish_time = (time.time()-st)*1000.0
+      st = time.time()
+      loss = loss.numpy()
+      cat = out.argmax(axis=1).numpy()
+      accuracy = (cat == Y).mean()
+      finish_time = (time.time()-st)*1000.0
 
-    # printing
-    t.set_description("loss %.2f accuracy %.2f -- %.2f + %.2f + %.2f + %.2f = %.2f" %
-      (loss, accuracy,
-      fp_time, bp_time, opt_time, finish_time,
-      fp_time + bp_time + opt_time + finish_time))
+      # printing
+      t.set_description("loss %.2f accuracy %.2f -- %.2f + %.2f + %.2f + %.2f = %.2f" %
+        (loss, accuracy,
+        fp_time, bp_time, opt_time, finish_time,
+        fp_time + bp_time + opt_time + finish_time))
 
-    del out, y, loss
+      del out, y, loss

examples/train_resnet.py

@@ -6,7 +6,7 @@ from tinygrad.nn.state import get_parameters
 from tinygrad.nn import optim
 from tinygrad.helpers import getenv
 from extra.training import train, evaluate
-from models.resnet import ResNet
+from extra.models.resnet import ResNet
 from extra.datasets import fetch_mnist
 
 

examples/transformer.py

@@ -5,7 +5,7 @@ import random
 from tinygrad.nn.state import get_parameters
 from tinygrad.nn.optim import Adam
 from extra.training import train, evaluate
-from models.transformer import Transformer
+from extra.models.transformer import Transformer
 
 # dataset idea from https://github.com/karpathy/minGPT/blob/master/projects/adder/adder.py
 def make_dataset():

examples/vgg7.py

@@ -1,12 +1,12 @@
 import sys
-import os
 import random
 import json
 import numpy
+from pathlib import Path
 from PIL import Image
 from tinygrad.tensor import Tensor
 from tinygrad.nn.optim import SGD
-from examples.vgg7_helpers.kinne import KinneDir
+from tinygrad.nn.state import safe_save, safe_load, get_state_dict, load_state_dict
 from examples.vgg7_helpers.waifu2x import image_load, image_save, Vgg7
 
 # amount of context erased by model
@@ -26,21 +26,23 @@ def set_sample_count(samples_dir, sc):
     file.write(str(sc) + "\n")
 
 if len(sys.argv) < 2:
-  print("python3 -m examples.vgg7 import MODELJSON MODELDIR")
+  print("python3 -m examples.vgg7 import MODELJSON MODEL")
   print(" imports a waifu2x JSON vgg_7 model, i.e. waifu2x/models/vgg_7/art/scale2.0x_model.json")
-  print(" into a directory of float binaries along with a meta.txt file containing tensor sizes")
+  print(" into a safetensors file")
   print(" weight tensors are ordered in tinygrad/ncnn form, as so: (outC,inC,H,W)")
-  print(" *this format is used by all other commands in this program*")
-  print("python3 -m examples.vgg7 execute MODELDIR IMG_IN IMG_OUT")
+  print(" *this format is used by most other commands in this program*")
+  print("python3 -m examples.vgg7 import_kinne MODEL_KINNE MODEL_SAFETENSORS")
+  print(" imports a model in 'KINNE' format (raw floats: used by older versions of this example) into safetensors")
+  print("python3 -m examples.vgg7 execute MODEL IMG_IN IMG_OUT")
   print(" given an already-nearest-neighbour-scaled image, runs vgg7 on it")
   print(" output image has 7 pixels removed on all edges")
   print(" do not run on large images, will have *hilarious* RAM use")
-  print("python3 -m examples.vgg7 execute_full MODELDIR IMG_IN IMG_OUT")
+  print("python3 -m examples.vgg7 execute_full MODEL IMG_IN IMG_OUT")
   print(" does the 'whole thing' (padding, tiling)")
   print(" safe for large images, etc.")
-  print("python3 -m examples.vgg7 new MODELDIR")
+  print("python3 -m examples.vgg7 new MODEL")
   print(" creates a new model (experimental)")
-  print("python3 -m examples.vgg7 train MODELDIR SAMPLES_DIR ROUNDS ROUNDS_SAVE")
+  print("python3 -m examples.vgg7 train MODEL SAMPLES_DIR ROUNDS ROUNDS_SAVE")
   print(" trains a model (experimental)")
   print(" (how experimental? well, every time I tried it, it flooded w/ NaNs)")
   print(" note: ROUNDS < 0 means 'forever'. ROUNDS_SAVE <= 0 is not a good idea.")
@@ -67,10 +69,11 @@ def load_and_save(path, save):
   if save:
     for v in vgg7.get_parameters():
       nansbane(v)
-  kn = KinneDir(model, save)
-  kn.parameters(vgg7.get_parameters())
-  kn.close()
-  if not save:
+    st = get_state_dict(vgg7)
+    safe_save(st, path)
+  else:
+    st = safe_load(path)
+    load_state_dict(vgg7, st)
     for v in vgg7.get_parameters():
       nansbane(v)
 
@@ -80,8 +83,19 @@ if cmd == "import":
 
   vgg7.load_waifu2x_json(json.load(open(src, "rb")))
 
-  if not os.path.isdir(model):
-    os.mkdir(model)
+  load_and_save(model, True)
+elif cmd == "import_kinne":
+  # tinygrad wasn't doing safetensors when this example was written
+  # it's possible someone might have a model around using the resulting interim format
+  src = sys.argv[2]
+  model = sys.argv[3]
+
+  index = 0
+  for t in vgg7.get_parameters():
+    fn = src + "/snoop_bin_" + str(index) + ".bin"
+    t.assign(Tensor(numpy.fromfile(fn, "<f4")).reshape(shape=t.shape))
+    index += 1
+
   load_and_save(model, True)
 elif cmd == "execute":
   model = sys.argv[2]
@@ -102,8 +116,6 @@ elif cmd == "execute_full":
 elif cmd == "new":
   model = sys.argv[2]
 
-  if not os.path.isdir(model):
-    os.mkdir(model)
   load_and_save(model, True)
 elif cmd == "train":
   model = sys.argv[2]
@@ -118,7 +130,7 @@ elif cmd == "train":
   # This is used to try and get the network to focus on "interesting" samples,
   #  which works nicely with the microsample system.
   sample_probs = None
-  sample_probs_path = model + "/sample_probs.bin"
+  sample_probs_path = model + "_sample_probs.bin"
   try:
     # try to read...
     sample_probs = numpy.fromfile(sample_probs_path, "<f8")
@@ -179,7 +191,7 @@ elif cmd == "train":
     optim.step()
 
     # warning: used by sample probability adjuster
-    loss_indicator = loss.max().numpy()[0]
+    loss_indicator = loss.max().numpy()
     print("Round " + str(rnum) + " : " + str(loss_indicator))
 
     if (rnum % rounds_per_save) == 0:

examples/vgg7_helpers/kinne.py

@@ -1,71 +0,0 @@
-from tinygrad.tensor import Tensor
-import numpy
-import os
-
-# Format Details:
-#  A KINNE parameter set is stored as a set of files named "snoop_bin_*.bin",
-#   where the * is a number starting at 0.
-#  Each file is simply raw little-endian floats,
-#   as readable by: numpy.fromfile(path, "<f4")
-#   and as writable by: t.numpy().astype("<f4", "C").tofile(path)
-# This format is intended to be extremely simple to get into literally anything.
-# It is not intended to be structural or efficient - reloading a network when
-#  unnecessary is inefficient anyway.
-# Ultimately, the idea behind this is as a format that, while it will always
-#  require code to implement, requires as little code as possible, and therefore
-#  works as a suitable interchange for any situation.
-# To add to the usability of the format, some informal metadata is provided,
-#  in "meta.txt", which provides human-readable shape information.
-# This is intended to help with debugging other implementations of the network,
-#  by providing concrete human-readable information on tensor shapes.
-# It is NOT meant to be read by machines.
-
-class KinneDir:
-  """
-  A KinneDir is an intermediate object used to save or load a model.
-  """
-
-  def __init__(self, base: str, save: bool):
-    """
-    Opens a new KINNE directory with the given base path.
-    If save is true, the directory is created if possible.
-    (This does not create parents.)
-    Save being true or false determines if tensors are loaded or saved.
-    The base path is of the form "models/abc" - no trailing slash.
-    It is important that if you wish to save in the current directory,
-     you use ".", not the empty string.
-    """
-    if save and not os.path.isdir(base):
-      os.mkdir(base)
-    self.base = base + "/snoop_bin_"
-    self.next_part_index = 0
-    self.save = save
-    if save:
-      self.metadata = open(base + "/meta.txt", "w")
-
-  def parameter(self, t: Tensor):
-    """
-    parameter loads or saves a parameter, given as a tensor.
-    """
-    path = f"{self.base}{self.next_part_index}.bin"
-    if self.save:
-      t.numpy().astype("<f4", "C").tofile(path)
-      self.metadata.write(f"{self.next_part_index}: {t.shape}\n")
-    else:
-      t.assign(Tensor(numpy.fromfile(path, "<f4")).reshape(shape=t.shape))
-    self.next_part_index += 1
-
-  def parameters(self, params):
-    """
-    parameters loads or saves a sequence of parameters.
-    It's intended for easily attaching to an existing model,
-     assuming that your parameters list orders are consistent.
-    (In other words, usage with tinygrad.utils.get_parameters isn't advised -
-      it's too 'implicit'.)
-    """
-    for t in params:
-      self.parameter(t)
-
-  def close(self):
-    if self.save:
-      self.metadata.close()

examples/vgg7_helpers/waifu2x.py

@@ -4,6 +4,7 @@
 import numpy
 from tinygrad.tensor import Tensor
 from PIL import Image
+from tinygrad.helpers import fetch
 
 # File Formats
 
@@ -17,6 +18,9 @@ def image_load(path) -> numpy.ndarray:
   """
   # file
   na = numpy.array(Image.open(path))
+  if na.shape[2] == 4:
+    # RGBA -> RGB (covers opaque images with alpha channels)
+    na = na[:,:,0:3]
   # fix shape
   na = numpy.moveaxis(na, [2,0,1], [0,1,2])
   # shape is now (3,h,w), add 1
@@ -47,33 +51,37 @@ class Conv3x3Biased:
   A 3x3 convolution layer with some utility functions.
   """
   def __init__(self, inC, outC, last = False):
+    # The properties must be named as "W" and "b".
+    # This is in an attempt to try and be roughly compatible with https://github.com/FHPythonUtils/Waifu2x
+    #  though this cannot necessarily account for transposition and other such things.
+
     # Massively overstate the weights to get them to be focused on,
     #  since otherwise the biases overrule everything
-    self.weight = Tensor.uniform(outC, inC, 3, 3) * 16.0
+    self.W = Tensor.uniform(outC, inC, 3, 3) * 16.0
     # Layout-wise, blatant cheat, but serious_mnist does it. I'd guess channels either have to have a size of 1 or whatever the target is?
     # Values-wise, entirely different blatant cheat.
     # In most cases, use uniform bias, but tiny.
     # For the last layer, use just 0.5, constant.
     if last:
-      self.bias = Tensor.zeros(1, outC, 1, 1) + 0.5
+      self.b = Tensor.zeros(1, outC, 1, 1) + 0.5
     else:
-      self.bias = Tensor.uniform(1, outC, 1, 1)
+      self.b = Tensor.uniform(1, outC, 1, 1)
 
   def forward(self, x):
     # You might be thinking, "but what about padding?"
     # Answer: Tiling is used to stitch everything back together, though you could pad the image before providing it.
-    return x.conv2d(self.weight).add(self.bias)
+    return x.conv2d(self.W).add(self.b)
 
   def get_parameters(self) -> list:
-    return [self.weight, self.bias]
+    return [self.W, self.b]
 
   def load_waifu2x_json(self, layer: dict):
     # Weights in this file are outChannel,inChannel,X,Y.
     # Not outChannel,inChannel,Y,X.
     # Therefore, transpose it before assignment.
     # I have long since forgotten how I worked this out.
-    self.weight.assign(Tensor(layer["weight"]).reshape(shape=self.weight.shape).transpose(2, 3))
-    self.bias.assign(Tensor(layer["bias"]).reshape(shape=self.bias.shape))
+    self.W.assign(Tensor(layer["weight"]).reshape(shape=self.W.shape).transpose(2, 3))
+    self.b.assign(Tensor(layer["bias"]).reshape(shape=self.b.shape))
 
 class Vgg7:
   """
@@ -109,6 +117,14 @@ class Vgg7:
   def get_parameters(self) -> list:
     return self.conv1.get_parameters() + self.conv2.get_parameters() + self.conv3.get_parameters() + self.conv4.get_parameters() + self.conv5.get_parameters() + self.conv6.get_parameters() + self.conv7.get_parameters()
 
+  def load_from_pretrained(self, intent = "art", subtype = "scale2.0x"):
+    """
+    Downloads a nagadomi/waifu2x JSON weight file and loads it.
+    """
+    import json
+    data = json.loads(fetch("https://github.com/nagadomi/waifu2x/raw/master/models/vgg_7/" + intent + "/" + subtype + "_model.json").read_bytes())
+    self.load_waifu2x_json(data)
+
   def load_waifu2x_json(self, data: list):
     """
     Loads weights from one of the waifu2x JSON files, i.e. waifu2x/models/vgg_7/art/noise0_model.json
@@ -122,7 +138,6 @@ class Vgg7:
     self.conv6.load_waifu2x_json(data[5])
     self.conv7.load_waifu2x_json(data[6])
 
-
   def forward_tiled(self, image: numpy.ndarray, tile_size: int) -> numpy.ndarray:
     """
     Given an ndarray image as loaded by image_load (NOT a tensor), scales it, pads it, splits it up, forwards the pieces, and reconstitutes it.

examples/vit.py

@@ -1,11 +1,9 @@
 import ast
-import io
 import numpy as np
 from PIL import Image
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import getenv
-from models.vit import ViT
-from extra.utils import fetch
+from tinygrad.helpers import getenv, fetch
+from extra.models.vit import ViT
 """
 fn = "gs://vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npz"
 import tensorflow as tf
@@ -24,14 +22,13 @@ else:
 m.load_from_pretrained()
 
 # category labels
-lbls = fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt")
-lbls = ast.literal_eval(lbls.decode('utf-8'))
+lbls = ast.literal_eval(fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt").read_text())
 
 #url = "https://upload.wikimedia.org/wikipedia/commons/4/41/Chicken.jpg"
 url = "https://repository-images.githubusercontent.com/296744635/39ba6700-082d-11eb-98b8-cb29fb7369c0"
 
 # junk
-img = Image.open(io.BytesIO(fetch(url)))
+img = Image.open(fetch(url))
 aspect_ratio = img.size[0] / img.size[1]
 img = img.resize((int(224*max(aspect_ratio,1.0)), int(224*max(1.0/aspect_ratio,1.0))))
 img = np.array(img)

examples/vits.py

@@ -1,12 +1,15 @@
-import json, logging, math, os, re, sys, time, wave, argparse, numpy as np
+import json, logging, math, re, sys, time, wave, argparse, numpy as np
+from phonemizer.phonemize import default_separator, _phonemize
+from phonemizer.backend import EspeakBackend
+from phonemizer.punctuation import Punctuation
 from functools import reduce
 from pathlib import Path
 from typing import List
-from extra.utils import download_file
-from tinygrad import nn
-from tinygrad.helpers import dtypes
+from tinygrad import nn, dtypes
+from tinygrad.helpers import fetch
 from tinygrad.nn.state import torch_load
 from tinygrad.tensor import Tensor
+from tinygrad.jit import TinyJit
 from unidecode import unidecode
 
 LRELU_SLOPE = 0.1
@@ -20,14 +23,14 @@ class Synthesizer:
     self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)
     self.dp = StochasticDurationPredictor(hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels) if use_sdp else DurationPredictor(hidden_channels, 256, 3, 0.5, gin_channels=gin_channels)
     if n_speakers > 1: self.emb_g = nn.Embedding(n_speakers, gin_channels)
-  def infer(self, x, x_lengths, sid=None, noise_scale=1.0, length_scale=1, noise_scale_w=1., max_len=None, emotion_embedding=None, max_y_length_estimate_scale=None):
-    x, m_p, logs_p, x_mask = self.enc_p.forward(x, x_lengths, emotion_embedding)
+  def infer(self, x, x_lengths, sid=None, noise_scale=1.0, length_scale=1, noise_scale_w=1., max_len=None, emotion_embedding=None, max_y_length_estimate_scale=None, pad_length=-1):
+    x, m_p, logs_p, x_mask = self.enc_p.forward(x.realize(), x_lengths.realize(), emotion_embedding.realize() if emotion_embedding is not None else emotion_embedding)
     g = self.emb_g(sid.reshape(1, 1)).squeeze(1).unsqueeze(-1) if self.n_speakers > 0 else None
-    logw = self.dp.forward(x, x_mask, g=g, reverse=self.use_sdp, noise_scale=noise_scale_w if self.use_sdp else 1.0)
+    logw = self.dp.forward(x, x_mask.realize(), g=g.realize(), reverse=self.use_sdp, noise_scale=noise_scale_w if self.use_sdp else 1.0)
     w_ceil = Tensor.ceil(logw.exp() * x_mask * length_scale)
     y_lengths = Tensor.maximum(w_ceil.sum([1, 2]), 1).cast(dtypes.int64)
-    return self.generate(g, logs_p, m_p, max_len, max_y_length_estimate_scale, noise_scale, w_ceil, x, x_mask, y_lengths)
-  def generate(self, g, logs_p, m_p, max_len, max_y_length_estimate_scale, noise_scale, w_ceil, x, x_mask, y_lengths):
+    return self.generate(g, logs_p, m_p, max_len, max_y_length_estimate_scale, noise_scale, w_ceil, x, x_mask, y_lengths, pad_length)
+  def generate(self, g, logs_p, m_p, max_len, max_y_length_estimate_scale, noise_scale, w_ceil, x, x_mask, y_lengths, pad_length):
     max_y_length = y_lengths.max().numpy() if max_y_length_estimate_scale is None else max(15, x.shape[-1]) * max_y_length_estimate_scale
     y_mask = sequence_mask(y_lengths, max_y_length).unsqueeze(1).cast(x_mask.dtype)
     attn_mask = x_mask.unsqueeze(2) * y_mask.unsqueeze(-1)
@@ -35,9 +38,17 @@ class Synthesizer:
     m_p_2 = attn.squeeze(1).matmul(m_p.transpose(1, 2)).transpose(1, 2)        # [b, t', t], [b, t, d] -> [b, d, t']
     logs_p_2 = attn.squeeze(1).matmul(logs_p.transpose(1, 2)).transpose(1, 2)  # [b, t', t], [b, t, d] -> [b, d, t']
     z_p = m_p_2 + Tensor.randn(*m_p_2.shape, dtype=m_p_2.dtype) * logs_p_2.exp() * noise_scale
-    y_mask = y_mask.cast(z_p.dtype)
-    z = self.flow.forward(z_p, y_mask, g=g, reverse=True)
-    o = self.dec.forward((z * y_mask)[:, :, :max_len], g=g)
+    row_len = y_mask.shape[2]
+    if pad_length > -1:
+      # Pad flow forward inputs to enable JIT
+      assert pad_length > row_len, "pad length is too small"
+      y_mask = y_mask.pad(((0, 0), (0, 0), (0, pad_length - row_len)), 0).cast(z_p.dtype)
+      # New y_mask tensor to remove sts mask
+      y_mask = Tensor(y_mask.numpy(), device=y_mask.device, dtype=y_mask.dtype, requires_grad=y_mask.requires_grad)
+      z_p = z_p.squeeze(0).pad(((0, 0), (0, pad_length - z_p.shape[2])), 1).unsqueeze(0)
+    z = self.flow.forward(z_p.realize(), y_mask.realize(), g=g.realize(), reverse=True)
+    result_length = reduce(lambda x, y: x * y, self.dec.upsample_rates, row_len)
+    o = self.dec.forward((z * y_mask)[:, :, :max_len], g=g)[:, :, :result_length]
     if max_y_length_estimate_scale is not None:
       length_scaler = o.shape[-1] / max_y_length
       o.realize()
@@ -69,6 +80,7 @@ class StochasticDurationPredictor:
     self.pre, self.proj = nn.Conv1d(in_channels, filter_channels, 1), nn.Conv1d(filter_channels, filter_channels, 1)
     self.convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
     if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
+  @TinyJit
   def forward(self, x: Tensor, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
     x = self.pre(x.detach())
     if g is not None: x = x + self.cond(g.detach())
@@ -97,13 +109,13 @@ class StochasticDurationPredictor:
         z, log_det = flow.forward(z, x_mask, g=x, reverse=reverse)
         log_det_tot = log_det_tot + log_det
       nll = Tensor.sum(0.5 * (math.log(2*math.pi) + (z**2)) * x_mask, [1,2]) - log_det_tot
-      return nll + log_q # [b]
+      return (nll + log_q).realize() # [b]
     flows = list(reversed(self.flows))
     flows = flows[:-2] + [flows[-1]] # remove a useless vflow
     z = Tensor.randn(x.shape[0], 2, x.shape[2], dtype=x.dtype).to(device=x.device) * noise_scale
     for flow in flows: z = flow.forward(z, x_mask, g=x, reverse=reverse)
     z0, z1 = split(z, [1, 1], 1)
-    return z0
+    return z0.realize()
 
 class DurationPredictor:
   def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0):
@@ -128,6 +140,7 @@ class TextEncoder:
     if emotion_embedding: self.emo_proj = nn.Linear(1024, hidden_channels)
     self.encoder = Encoder(hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout)
     self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+  @TinyJit
   def forward(self, x: Tensor, x_lengths: Tensor, emotion_embedding=None):
     if self.n_vocab!=0: x = (self.emb(x) * math.sqrt(self.hidden_channels))
     if emotion_embedding: x = x + self.emo_proj(emotion_embedding).unsqueeze(1)
@@ -135,7 +148,7 @@ class TextEncoder:
     x_mask = sequence_mask(x_lengths, x.shape[2]).unsqueeze(1).cast(x.dtype)
     x = self.encoder.forward(x * x_mask, x_mask)
     m, logs = split(self.proj(x) * x_mask, self.out_channels, dim=1)
-    return x, m, logs, x_mask
+    return x.realize(), m.realize(), logs.realize(), x_mask.realize()
 
 class ResidualCouplingBlock:
   def __init__(self, channels, hidden_channels, kernel_size, dilation_rate, n_layers, n_flows=4, gin_channels=0):
@@ -144,9 +157,10 @@ class ResidualCouplingBlock:
     for _ in range(n_flows):
       self.flows.append(ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels, mean_only=True))
       self.flows.append(Flip())
+  @TinyJit
   def forward(self, x, x_mask, g=None, reverse=False):
     for flow in reversed(self.flows) if reverse else self.flows: x = flow.forward(x, x_mask, g=g, reverse=reverse)
-    return x
+    return x.realize()
 
 class PosteriorEncoder:
   def __init__(self, in_channels, out_channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0):
@@ -167,22 +181,23 @@ class Generator:
     resblock = ResBlock1 if resblock == '1' else ResBlock2
     self.ups = [nn.ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)), k, u, padding=(k-u)//2) for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes))]
     self.resblocks = []
+    self.upsample_rates = upsample_rates
     for i in range(len(self.ups)):
       ch = upsample_initial_channel // (2 ** (i + 1))
       for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
         self.resblocks.append(resblock(ch, k, d))
     self.conv_post = nn.Conv1d(ch, 1, 7, 1, padding=3, bias=False)
     if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+  @TinyJit
   def forward(self, x: Tensor, g=None):
     x = self.conv_pre(x)
     if g is not None:  x = x + self.cond(g)
     for i in range(self.num_upsamples):
-      x, xs = self.ups[i](x.leakyrelu(LRELU_SLOPE)), None
-      for j in range(self.num_kernels):
-        if xs is None: xs = self.resblocks[i * self.num_kernels + j].forward(x)
-        else: xs += self.resblocks[i * self.num_kernels + j].forward(x)
-      x = xs / self.num_kernels
-    return self.conv_post(x.leakyrelu()).tanh()
+      x = self.ups[i](x.leakyrelu(LRELU_SLOPE))
+      xs = sum(self.resblocks[i * self.num_kernels + j].forward(x) for j in range(self.num_kernels))
+      x = (xs / self.num_kernels).realize()
+    res = self.conv_post(x.leakyrelu()).tanh().realize()
+    return res
 
 class LayerNorm(nn.LayerNorm):
   def __init__(self, channels, eps=1e-5): super().__init__(channels, eps, elementwise_affine=True)
@@ -446,7 +461,7 @@ def rational_quadratic_spline(inputs: Tensor, unnormalized_widths: Tensor, unnor
   derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * theta.pow(2) + 2 * input_delta * theta_one_minus_theta + input_derivatives * (1 - theta).pow(2))
   return input_cum_heights + numerator / denominator, derivative_numerator.log() - 2 * denominator.log()
 
-def sequence_mask(length: Tensor, max_length): return Tensor.arange(max_length, dtype=length.dtype, device=length.device).unsqueeze(0).__lt__(length.unsqueeze(1))
+def sequence_mask(length: Tensor, max_length): return Tensor.arange(max_length, dtype=length.dtype, device=length.device).unsqueeze(0) < length.unsqueeze(1)
 def generate_path(duration: Tensor, mask: Tensor):  # duration: [b, 1, t_x], mask: [b, 1, t_y, t_x]
   b, _, t_y, t_x = mask.shape
   path = sequence_mask(duration.cumsum(axis=2).reshape(b * t_x), t_y).cast(mask.dtype).reshape(b, t_x, t_y)
@@ -516,13 +531,11 @@ class HParams:
 
 # MODEL LOADING
 def load_model(symbols, hps, model) -> Synthesizer:
-  weights_path = model[1]
-  download_if_not_present(weights_path, model[3])
   net_g = Synthesizer(len(symbols), hps.data.filter_length // 2 + 1, hps.train.segment_size // hps.data.hop_length, n_speakers = hps.data.n_speakers, **hps.model)
-  _ = load_checkpoint(weights_path, net_g, None)
+  _ = load_checkpoint(fetch(model[1]), net_g, None)
   return net_g
 def load_checkpoint(checkpoint_path, model: Synthesizer, optimizer=None, skip_list=[]):
-  assert os.path.isfile(checkpoint_path)
+  assert Path(checkpoint_path).is_file()
   start_time = time.time()
   checkpoint_dict = torch_load(checkpoint_path)
   iteration, learning_rate = checkpoint_dict['iteration'], checkpoint_dict['learning_rate']
@@ -555,18 +568,15 @@ def load_checkpoint(checkpoint_path, model: Synthesizer, optimizer=None, skip_li
   logging.info(f"Loaded checkpoint '{checkpoint_path}' (iteration {iteration}) in {time.time() - start_time:.4f}s")
   return model, optimizer, learning_rate, iteration
 
-def download_if_not_present(file_path: Path, url: str):
-  if not os.path.isfile(file_path):
-    logging.info(f"Did not find {file_path}, downloading...")
-    download_file(url, file_path)
-  return file_path
-
 # Used for cleaning input text and mapping to symbols
 class TextMapper: # Based on https://github.com/keithito/tacotron
   def __init__(self, symbols, apply_cleaners=True):
     self.apply_cleaners, self.symbols, self._inflect = apply_cleaners, symbols, None
     self._symbol_to_id, _id_to_symbol = {s: i for i, s in enumerate(symbols)}, {i: s for i, s in enumerate(symbols)}
     self._whitespace_re, self._abbreviations = re.compile(r'\s+'), [(re.compile('\\b%s\\.' % x[0], re.IGNORECASE), x[1]) for x in [('mrs', 'misess'), ('mr', 'mister'), ('dr', 'doctor'), ('st', 'saint'), ('co', 'company'), ('jr', 'junior'), ('maj', 'major'), ('gen', 'general'), ('drs', 'doctors'), ('rev', 'reverend'), ('lt', 'lieutenant'), ('hon', 'honorable'), ('sgt', 'sergeant'), ('capt', 'captain'), ('esq', 'esquire'), ('ltd', 'limited'), ('col', 'colonel'), ('ft', 'fort'), ]]
+    self.phonemizer = EspeakBackend(
+        language="en-us", punctuation_marks=Punctuation.default_marks(), preserve_punctuation=True, with_stress=True,
+    )
   def text_to_sequence(self, text, cleaner_names):
     if self.apply_cleaners:
       for name in cleaner_names:
@@ -575,18 +585,16 @@ class TextMapper: # Based on https://github.com/keithito/tacotron
         text = cleaner(text)
     else: text = text.strip()
     return [self._symbol_to_id[symbol] for symbol in text]
-  def get_text(self, text, add_blank=False, cleaners=('english_cleaners',)):
+  def get_text(self, text, add_blank=False, cleaners=('english_cleaners2',)):
     text_norm = self.text_to_sequence(text, cleaners)
     return Tensor(self.intersperse(text_norm, 0) if add_blank else text_norm, dtype=dtypes.int64)
   def intersperse(self, lst, item):
     (result := [item] * (len(lst) * 2 + 1))[1::2] = lst
     return result
+  def phonemize(self, text, strip=True): return _phonemize(self.phonemizer, text, default_separator, strip, 1, False, False)
   def filter_oov(self, text): return "".join(list(filter(lambda x: x in self._symbol_to_id, text)))
-  def base_english_cleaners(self, text, preserve_punctuation=False, with_stress=False):
-    from phonemizer import phonemize
-    return self.collapse_whitespace(phonemize(self.expand_abbreviations(unidecode(text.lower())), language='en-us', backend='espeak', strip=True, preserve_punctuation=preserve_punctuation, with_stress=with_stress))
-  def english_cleaners(self, text): return self.base_english_cleaners(text)
-  def english_cleaners2(self, text): return self.base_english_cleaners(text, preserve_punctuation=True, with_stress=True)
+  def base_english_cleaners(self, text): return self.collapse_whitespace(self.phonemize(self.expand_abbreviations(unidecode(text.lower()))))
+  def english_cleaners2(self, text): return self.base_english_cleaners(text)
   def transliteration_cleaners(self, text): return self.collapse_whitespace(unidecode(text.lower()))
   def cjke_cleaners(self, text): return re.sub(r'([^\.,!\?\-…~])$', r'\1.', re.sub(r'\s+$', '', self.english_to_ipa2(text).replace('ɑ', 'a').replace('ɔ', 'o').replace('ɛ', 'e').replace('ɪ', 'i').replace('ʊ', 'u')))
   def cjke_cleaners2(self, text): return re.sub(r'([^\.,!\?\-…~])$', r'\1.', re.sub(r'\s+$', '', self.english_to_ipa2(text)))
@@ -649,14 +657,14 @@ class TextMapper: # Based on https://github.com/keithito/tacotron
 # anime lady 1 | --model_to_use uma_trilingual --speaker_id 36
 # anime lady 2 | --model_to_use uma_trilingual --speaker_id 121
 #########################################################################################
-VITS_PATH = Path(__file__).parent.parent / "weights/VITS/"
-MODELS = { # config_path, weights_path, config_url, weights_url
-  "ljs": (VITS_PATH / "config_ljs.json", VITS_PATH / "pretrained_ljs.pth", "https://raw.githubusercontent.com/jaywalnut310/vits/main/configs/ljs_base.json", "https://drive.google.com/uc?export=download&id=1q86w74Ygw2hNzYP9cWkeClGT5X25PvBT&confirm=t"),
-  "vctk": (VITS_PATH / "config_vctk.json", VITS_PATH / "pretrained_vctk.pth", "https://raw.githubusercontent.com/jaywalnut310/vits/main/configs/vctk_base.json", "https://drive.google.com/uc?export=download&id=11aHOlhnxzjpdWDpsz1vFDCzbeEfoIxru&confirm=t"),
-  "mmts-tts": (VITS_PATH / "config_mmts-tts.json", VITS_PATH / "pretrained_mmts-tts.pth", "https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/config.json", "https://huggingface.co/facebook/mms-tts/resolve/main/full_models/eng/G_100000.pth"),
-  "uma_trilingual": (VITS_PATH / "config_uma_trilingual.json", VITS_PATH / "pretrained_uma_trilingual.pth", "https://huggingface.co/spaces/Plachta/VITS-Umamusume-voice-synthesizer/raw/main/configs/uma_trilingual.json", "https://huggingface.co/spaces/Plachta/VITS-Umamusume-voice-synthesizer/resolve/main/pretrained_models/G_trilingual.pth"),
-  "cjks": (VITS_PATH / "config_cjks.json", VITS_PATH / "pretrained_cjks.pth", "https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/14/config.json", "https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/14/model.pth"),
-  "voistock": (VITS_PATH / "config_voistock.json", VITS_PATH / "pretrained_voistock.pth", "https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/15/config.json", "https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/15/model.pth"),
+VITS_PATH = Path(__file__).parents[1] / "weights/VITS/"
+MODELS = { # config_url, weights_url
+  "ljs": ("https://raw.githubusercontent.com/jaywalnut310/vits/main/configs/ljs_base.json", "https://drive.google.com/uc?export=download&id=1q86w74Ygw2hNzYP9cWkeClGT5X25PvBT&confirm=t"),
+  "vctk": ("https://raw.githubusercontent.com/jaywalnut310/vits/main/configs/vctk_base.json", "https://drive.google.com/uc?export=download&id=11aHOlhnxzjpdWDpsz1vFDCzbeEfoIxru&confirm=t"),
+  "mmts-tts": ("https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/config.json", "https://huggingface.co/facebook/mms-tts/resolve/main/full_models/eng/G_100000.pth"),
+  "uma_trilingual": ("https://huggingface.co/spaces/Plachta/VITS-Umamusume-voice-synthesizer/raw/main/configs/uma_trilingual.json", "https://huggingface.co/spaces/Plachta/VITS-Umamusume-voice-synthesizer/resolve/main/pretrained_models/G_trilingual.pth"),
+  "cjks": ("https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/14/config.json", "https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/14/model.pth"),
+  "voistock": ("https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/15/config.json", "https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/15/model.pth"),
 }
 Y_LENGTH_ESTIMATE_SCALARS = {"ljs": 2.8, "vctk": 1.74, "mmts-tts": 1.9, "uma_trilingual": 2.3, "cjks": 3.3, "voistock": 3.1}
 if __name__ == '__main__':
@@ -665,7 +673,7 @@ if __name__ == '__main__':
   parser.add_argument("--model_to_use", default="vctk", help="Specify the model to use. Default is 'vctk'.")
   parser.add_argument("--speaker_id", type=int, default=6, help="Specify the speaker ID. Default is 6.")
   parser.add_argument("--out_path", default=None, help="Specify the full output path. Overrides the --out_dir and --name parameter.")
-  parser.add_argument("--out_dir", default=str(Path(__file__).parent.parent / "temp"), help="Specify the output path.")
+  parser.add_argument("--out_dir", default=str(Path(__file__).parents[1] / "temp"), help="Specify the output path.")
   parser.add_argument("--base_name", default="test", help="Specify the base of the output file name. Default is 'test'.")
   parser.add_argument("--text_to_synthesize", default="""Hello person. If the code you are contributing isn't some of the highest quality code you've written in your life, either put in the effort to make it great, or don't bother.""", help="Specify the text to synthesize. Default is a greeting message.")
   parser.add_argument("--noise_scale", type=float, default=0.667, help="Specify the noise scale. Default is 0.667.")
@@ -681,9 +689,7 @@ if __name__ == '__main__':
   model_config = MODELS[args.model_to_use]
 
   # Load the hyperparameters from the config file.
-  config_path = model_config[0]
-  download_if_not_present(config_path, model_config[2])
-  hps = get_hparams_from_file(config_path)
+  hps = get_hparams_from_file(fetch(model_config[0]))
 
   # If model has multiple speakers, validate speaker id and retrieve name if available.
   model_has_multiple_speakers = hps.data.n_speakers > 0
@@ -705,7 +711,7 @@ if __name__ == '__main__':
 
   # Load symbols, instantiate TextMapper and clean the text.
   if hps.__contains__("symbols"): symbols = hps.symbols
-  elif args.model_to_use == "mmts-tts": symbols = [x.replace("\n", "") for x in open(download_if_not_present(VITS_PATH / "vocab_mmts-tts.txt", "https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/vocab.txt"), encoding="utf-8").readlines()]
+  elif args.model_to_use == "mmts-tts": symbols = [x.replace("\n", "") for x in fetch("https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/vocab.txt").open(encoding="utf-8").readlines()]
   else: symbols = ['_'] + list(';:,.!?¡¿—…"«»“” ') + list('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz') + list("ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ")
   text_mapper = TextMapper(apply_cleaners=True, symbols=symbols)
 

examples/webgl/yolov8/compile.py

@@ -0,0 +1,23 @@
+from pathlib import Path
+from examples.yolov8 import YOLOv8
+from tinygrad.tensor import Tensor
+from tinygrad.nn.state import safe_save
+from extra.export_model import export_model
+from tinygrad.helpers import fetch
+from tinygrad.helpers import getenv
+from tinygrad.device import Device
+from tinygrad.nn.state import safe_load, load_state_dict
+
+if __name__ == "__main__":
+    Device.DEFAULT = "WEBGL"
+    yolo_variant = 'n'
+    yolo_infer = YOLOv8(w=0.25, r=2.0, d=0.33, num_classes=80)
+    weights_location = Path(__file__).parents[1] / "weights" / f'yolov8{yolo_variant}.safetensors'
+    fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{yolo_variant}.safetensors', weights_location)
+    state_dict = safe_load(weights_location)
+    load_state_dict(yolo_infer, state_dict)
+    prg, inp_sizes, out_sizes, state = export_model(yolo_infer, Device.DEFAULT.lower(), Tensor.randn(1,3,640,640))
+    dirname = Path(__file__).parent
+    safe_save(state, (dirname / "net.safetensors").as_posix())
+    with open(dirname / f"net.js", "w") as text_file:
+       text_file.write(prg)

examples/webgl/yolov8/index.html

@@ -0,0 +1,223 @@
+<!-- Pre and post-processing functions from: https://github.com/AndreyGermanov/yolov8_onnx_javascript -->
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <title>YOLOv8 tinygrad WebGL</title>
+    <script src="./net.js"></script>
+    <style>
+       body {
+            text-align: center;
+            font-family: Arial, sans-serif;
+            margin: 0;
+            padding: 0;
+            overflow: hidden;
+        }
+
+        .video-container {
+            position: relative;
+            width: 100%;
+            margin: 0 auto;
+        }
+
+        #video, #canvas {
+            position: absolute;
+            top: 0;
+            left: 0;
+            width: 100%;
+            height: auto;
+        }
+
+        #canvas {
+            background: transparent;
+        }
+
+        h1 {
+            margin-top: 20px;
+        }
+    </style>
+</head>
+<body>
+    <h1>YOLOv8 tinygrad WebGL</h1>
+    <div class="video-container">
+        <video id="video" muted autoplay playsinline></video>
+        <canvas id="canvas"></canvas>
+    </div>
+    
+    <script>
+        let net = null;
+
+        const video = document.getElementById('video');
+        const canvas = document.getElementById('canvas');
+        const context = canvas.getContext('2d');
+        const offscreenCanvas = document.createElement('canvas');
+        offscreenCanvas.width = 640;
+        offscreenCanvas.height = 640;
+        const offscreenContext = offscreenCanvas.getContext('2d');
+
+        if (navigator.mediaDevices && navigator.mediaDevices.getUserMedia) {
+            navigator.mediaDevices.getUserMedia({ audio: false, video: true }).then(function (stream) {
+                video.srcObject = stream;
+                video.onloadedmetadata = function() {
+                    canvas.width = video.clientWidth;
+                    canvas.height = video.clientHeight;
+                }
+            });
+        }
+
+        async function processFrame() {
+            offscreenContext.drawImage(video, 0, 0, 640, 640);
+            const boxes = await detectObjectsOnFrame(offscreenContext);
+            drawBoxes(offscreenCanvas, boxes);
+            requestAnimationFrame(processFrame);
+        }
+
+        requestAnimationFrame(processFrame);
+
+        function drawBoxes(offscreenCanvas, boxes) {
+            const canvas = document.querySelector("canvas");
+            const ctx = canvas.getContext("2d");
+            ctx.clearRect(0, 0, canvas.width, canvas.height);
+            ctx.lineWidth = 3;
+            ctx.font = "20px serif";
+            const scaleX = canvas.width / 640;
+            const scaleY = canvas.height / 640;
+
+            boxes.forEach(([x1, y1, x2, y2, label]) => {
+                const classIndex = yolo_classes.indexOf(label);
+                const color = classColors[classIndex];
+                const textWidth = ctx.measureText(label).width;
+                ctx.strokeStyle = color;
+                ctx.fillStyle = color;
+
+                let adjustedX1 = x1 * scaleX;
+                let adjustedY1 = y1 * scaleY;
+                let adjustedX2 = x2 * scaleX;
+                let adjustedY2 = y2 * scaleY;
+                let boxWidth = adjustedX2 - adjustedX1;
+                let boxHeight = adjustedY2 - adjustedY1;
+
+                ctx.strokeRect(adjustedX1, adjustedY1, boxWidth, boxHeight);
+                ctx.fillRect(adjustedX1, adjustedY1 - 25, textWidth + 10, 25);
+                ctx.fillStyle = "#000000";
+                ctx.fillText(label, adjustedX1, adjustedY1 - 7);
+            });
+        }
+
+        async function detectObjectsOnFrame(offscreenContext) {
+            if (!net) net = await loadNet();
+            let start = performance.now();
+            const [input,img_width,img_height] = await prepareInput(offscreenContext);
+            console.log("Preprocess took: " + (performance.now() - start) + " ms");
+            start = performance.now();
+            const output = net(new Float32Array(input));
+            console.log("Inference took: " + (performance.now() - start) + " ms");
+            start = performance.now();
+            let out = processOutput(output,img_width,img_height);
+            console.log("Postprocess took: " + (performance.now() - start) + " ms");
+            return out;
+        }
+
+        async function prepareInput(offscreenContext) {
+            return new Promise(resolve => {
+                const [img_width,img_height] = [640, 640]
+                const imgData = offscreenContext.getImageData(0,0,640,640);
+                const pixels = imgData.data;
+                const red = [], green = [], blue = [];
+
+                for (let index=0; index<pixels.length; index+=4) {
+                    red.push(pixels[index]/255.0);
+                    green.push(pixels[index+1]/255.0);
+                    blue.push(pixels[index+2]/255.0);
+                }
+                const input = [...red, ...green, ...blue];
+                resolve([input, img_width, img_height])
+            })
+        }
+        
+        const loadNet = async () => {
+            try {
+                const safetensor = await (new Uint8Array(await (await fetch("./net.safetensors")).arrayBuffer()));
+                const gl = document.createElement("canvas").getContext("webgl2");
+                return setupNet(gl, safetensor);
+            } catch (e) {
+                console.log(e);
+                return null;
+            }
+        }
+
+        function processOutput(output, img_width, img_height) {
+            let boxes = [];
+            for (let index=0;index<8400;index++) {
+                const [class_id,prob] = [...Array(80).keys()]
+                    .map(col => [col, output[8400*(col+4)+index]])
+                    .reduce((accum, item) => item[1]>accum[1] ? item : accum,[0,0]);
+                if (prob < 0.25) {
+                    continue;
+                }
+                const label = yolo_classes[class_id];
+                const xc = output[index];
+                const yc = output[8400+index];
+                const w = output[2*8400+index];
+                const h = output[3*8400+index];
+                const x1 = (xc-w/2)/640*img_width;
+                const y1 = (yc-h/2)/640*img_height;
+                const x2 = (xc+w/2)/640*img_width;
+                const y2 = (yc+h/2)/640*img_height;
+                boxes.push([x1,y1,x2,y2,label,prob]);
+            }
+
+            boxes = boxes.sort((box1,box2) => box2[5]-box1[5])
+            const result = [];
+            while (boxes.length>0) {
+                result.push(boxes[0]);
+                boxes = boxes.filter(box => iou(boxes[0],box)<0.7);
+            }
+            return result;
+        }
+
+        function iou(box1,box2) {
+            return intersection(box1,box2)/union(box1,box2);
+        }
+
+        function union(box1,box2) {
+            const [box1_x1,box1_y1,box1_x2,box1_y2] = box1;
+            const [box2_x1,box2_y1,box2_x2,box2_y2] = box2;
+            const box1_area = (box1_x2-box1_x1)*(box1_y2-box1_y1)
+            const box2_area = (box2_x2-box2_x1)*(box2_y2-box2_y1)
+            return box1_area + box2_area - intersection(box1,box2)
+        }
+
+        function intersection(box1,box2) {
+            const [box1_x1,box1_y1,box1_x2,box1_y2] = box1;
+            const [box2_x1,box2_y1,box2_x2,box2_y2] = box2;
+            const x1 = Math.max(box1_x1,box2_x1);
+            const y1 = Math.max(box1_y1,box2_y1);
+            const x2 = Math.min(box1_x2,box2_x2);
+            const y2 = Math.min(box1_y2,box2_y2);
+            return (x2-x1)*(y2-y1)
+        }
+
+        const yolo_classes = [
+            'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
+            'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
+            'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase',
+            'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
+            'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
+            'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant',
+            'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
+            'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
+        ];
+
+        function generateColors(numColors) {
+            const colors = [];
+            for (let i = 0; i < 360; i += 360 / numColors) {
+                colors.push(`hsl(${i}, 100%, 50%)`);
+            }
+            return colors;
+        }
+
+        const classColors = generateColors(yolo_classes.length);
+    </script>
+</body>
+</html>

examples/webgpu/stable_diffusion/compile.py

@@ -0,0 +1,231 @@
+import os
+from extra.export_model import compile_net, jit_model
+from examples.stable_diffusion import StableDiffusion
+from tinygrad.nn.state import get_state_dict, safe_save, safe_load_metadata, torch_load, load_state_dict
+from tinygrad.tensor import Tensor
+from tinygrad import Device
+from tinygrad.helpers import fetch
+from typing import NamedTuple, Any, List
+from pathlib import Path
+import argparse
+import numpy as np
+
+def convert_f32_to_f16(input_file, output_file):
+  with open(input_file, 'rb') as f:
+    metadata_length_bytes = f.read(8)
+    metadata_length = int.from_bytes(metadata_length_bytes, byteorder='little', signed=False)
+    metadata_json_bytes = f.read(metadata_length)
+    float32_values = np.fromfile(f, dtype=np.float32)
+
+  first_text_model_offset = 3772703308
+  num_elements = int((first_text_model_offset)/4)
+  front_float16_values = float32_values[:num_elements].astype(np.float16)
+  rest_float32_values = float32_values[num_elements:]
+
+  with open(output_file, 'wb') as f:
+    f.write(metadata_length_bytes)
+    f.write(metadata_json_bytes)
+    front_float16_values.tofile(f)
+    rest_float32_values.tofile(f)
+
+def split_safetensor(fn):
+  _, json_len, metadata = safe_load_metadata(fn)
+  text_model_offset = 3772703308
+  chunk_size = 536870912
+
+  for k in metadata:
+    # safetensor is in fp16, except for text moel
+    if (metadata[k]["data_offsets"][0] < text_model_offset):
+      metadata[k]["data_offsets"][0] = int(metadata[k]["data_offsets"][0]/2)
+      metadata[k]["data_offsets"][1] = int(metadata[k]["data_offsets"][1]/2)
+
+  last_offset = 0
+  part_end_offsets = []
+
+  for k in metadata:
+    offset = metadata[k]['data_offsets'][0]
+
+    if offset == text_model_offset:
+      break
+
+    part_offset = offset - last_offset
+
+    if (part_offset >= chunk_size):
+      part_end_offsets.append(8+json_len+offset)
+      last_offset = offset
+
+  text_model_start = int(text_model_offset/2)
+  net_bytes = bytes(open(fn, 'rb').read())
+  part_end_offsets.append(text_model_start+8+json_len)
+  cur_pos = 0
+
+  for i, end_pos in enumerate(part_end_offsets):
+    with open(f'./net_part{i}.safetensors', "wb+") as f:
+      f.write(net_bytes[cur_pos:end_pos])
+      cur_pos = end_pos
+
+  with open(f'./net_textmodel.safetensors', "wb+") as f:
+    f.write(net_bytes[text_model_start+8+json_len:])
+
+  return part_end_offsets
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(description='Run Stable Diffusion', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  parser.add_argument('--remoteweights', action='store_true', help="Use safetensors from Huggingface, or from local")
+  args = parser.parse_args()
+  Device.DEFAULT = "WEBGPU"
+
+  Tensor.no_grad = True
+  model = StableDiffusion()
+
+  # load in weights
+  load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], strict=False)
+
+  class Step(NamedTuple):
+    name: str = ""
+    input: List[Tensor] = []
+    forward: Any = None
+
+  sub_steps = [
+    Step(name = "textModel", input = [Tensor.randn(1, 77)], forward = model.cond_stage_model.transformer.text_model),
+    Step(name = "diffusor", input = [Tensor.randn(1, 77, 768), Tensor.randn(1, 77, 768), Tensor.randn(1,4,64,64), Tensor.rand(1), Tensor.randn(1), Tensor.randn(1), Tensor.randn(1)], forward = model),
+    Step(name = "decoder", input = [Tensor.randn(1,4,64,64)], forward = model.decode)
+  ]
+
+  prg = ""
+
+  def compile_step(model, step: Step):
+    run, special_names = jit_model(step, *step.input)
+    functions, statements, bufs, _ = compile_net(run, special_names)
+    state = get_state_dict(model)
+    weights = {id(x.lazydata.base.realized): name for name, x in state.items()}
+    kernel_code = '\n\n'.join([f"const {key} = `{code.replace(key, 'main')}`;" for key, code in functions.items()])
+    kernel_names = ', '.join([name for (name, _, _, _) in statements])
+    kernel_calls = '\n        '.join([f"addComputePass(device, commandEncoder, piplines[{i}], [{', '.join(args)}], {global_size});" for i, (_name, args, global_size, _local_size) in enumerate(statements) ])
+    bufs =  '\n    '.join([f"const {name} = " + (f"createEmptyBuf(device, {size});" if _key not in weights else f"createWeightBuf(device, {size}, getTensorBuffer(safetensor, metadata['{weights[_key]}'], '{weights[_key]}'))") + ";"  for name,(size,dtype,_key) in bufs.items()])
+    gpu_write_bufs =  '\n    '.join([f"const gpuWriteBuffer{i} = device.createBuffer({{size:input{i}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});" for i,(_,value) in enumerate(special_names.items()) if "output" not in value])
+    input_writer = '\n    '.join([f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n    new Float32Array(gpuWriteBuffer{i}.getMappedRange()).set(" + f'data{i});' + f"\n    gpuWriteBuffer{i}.unmap();\ncommandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, input{i}, 0, gpuWriteBuffer{i}.size);"  for i,(_,value) in enumerate(special_names.items()) if value != "output0"])
+    return f"""\n    var {step.name} = function() {{
+
+    {kernel_code}
+
+    return {{
+      "setup": async (device, safetensor) => {{
+        const metadata = getTensorMetadata(safetensor[0]);
+
+        {bufs}
+
+        {gpu_write_bufs}
+        const gpuReadBuffer = device.createBuffer({{ size: output0.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ }});
+
+        const kernels = [{kernel_names}];
+        const piplines = await Promise.all(kernels.map(name => device.createComputePipelineAsync({{layout: "auto", compute: {{ module: device.createShaderModule({{ code: name }}), entryPoint: "main" }}}})));
+
+        return async ({",".join([f'data{i}' for i,(k,v) in enumerate(special_names.items()) if v != "output0"])}) => {{
+            const commandEncoder = device.createCommandEncoder();
+
+            {input_writer}
+
+            {kernel_calls}
+            commandEncoder.copyBufferToBuffer(output0, 0, gpuReadBuffer, 0, output0.size);
+            const gpuCommands = commandEncoder.finish();
+            device.queue.submit([gpuCommands]);
+
+            await gpuReadBuffer.mapAsync(GPUMapMode.READ);
+            const resultBuffer = new Float32Array(gpuReadBuffer.size/4);
+            resultBuffer.set(new Float32Array(gpuReadBuffer.getMappedRange()));
+            gpuReadBuffer.unmap();
+            return resultBuffer;
+        }}
+      }}
+    }}
+  }}
+  """
+
+  for step in sub_steps:
+    print(f'Executing step={step.name}')
+    prg += compile_step(model, step)
+
+    if step.name == "diffusor":
+      if args.remoteweights:
+        base_url = "https://huggingface.co/wpmed/tinygrad-sd-f16/resolve/main"
+      else:
+        state = get_state_dict(model)
+        safe_save(state, os.path.join(os.path.dirname(__file__), "net.safetensors"))
+        convert_f32_to_f16("./net.safetensors", "./net_conv.safetensors")
+        split_safetensor("./net_conv.safetensors")
+        os.remove("net.safetensors")
+        os.remove("net_conv.safetensors")
+        base_url = "."
+
+  prekernel = f"""
+    window.MODEL_BASE_URL= "{base_url}";
+    const getTensorMetadata = (safetensorBuffer) => {{
+      const metadataLength = Number(new DataView(safetensorBuffer.buffer).getBigUint64(0, true));
+      const metadata = JSON.parse(new TextDecoder("utf8").decode(safetensorBuffer.subarray(8, 8 + metadataLength)));
+      return Object.fromEntries(Object.entries(metadata).filter(([k, v]) => k !== "__metadata__").map(([k, v]) => [k, {{...v, data_offsets: v.data_offsets.map(x => 8 + metadataLength + x)}}]));
+    }};
+
+  const getTensorBuffer = (safetensorParts, tensorMetadata, key) => {{
+    let selectedPart = 0;
+    let counter = 0;
+    let partStartOffsets = [1131408336, 2227518416, 3308987856, 4265298864];
+    let correctedOffsets = tensorMetadata.data_offsets;
+    let prev_offset = 0;
+
+    for (let start of partStartOffsets) {{
+      prev_offset = (counter == 0) ? 0 : partStartOffsets[counter-1];
+
+      if (tensorMetadata.data_offsets[0] < start) {{
+        selectedPart = counter;
+        correctedOffsets = [correctedOffsets[0]-prev_offset, correctedOffsets[1]-prev_offset];
+        break;
+      }}
+
+      counter++;
+    }}
+
+    let allZero = true;
+    let out = safetensorParts[selectedPart].subarray(...correctedOffsets);
+
+    for (let i = 0; i < out.length; i++) {{
+        if (out[i] !== 0) {{
+            allZero = false;
+            break;
+        }}
+    }}
+
+    if (allZero) {{
+        console.log("Error: weight '" + key + "' is all zero.");
+    }}
+
+    return safetensorParts[selectedPart].subarray(...correctedOffsets);
+  }}
+
+  const getWeight = (safetensors, key) => {{
+    let uint8Data = getTensorBuffer(safetensors, getTensorMetadata(safetensors[0])[key], key);
+    return new Float32Array(uint8Data.buffer, uint8Data.byteOffset, uint8Data.byteLength / Float32Array.BYTES_PER_ELEMENT);
+  }}
+
+  const createEmptyBuf = (device, size) => {{
+      return device.createBuffer({{size, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST }});
+  }};
+
+  const createWeightBuf = (device, size, data) => {{
+    const buf = device.createBuffer({{ mappedAtCreation: true, size, usage: GPUBufferUsage.STORAGE }});
+    new Uint8Array(buf.getMappedRange()).set(data);
+    buf.unmap();
+    return buf;
+  }};
+
+  const addComputePass = (device, commandEncoder, pipeline, bufs, workgroup) => {{
+    const bindGroup = device.createBindGroup({{layout: pipeline.getBindGroupLayout(0), entries: bufs.map((buffer, index) => ({{ binding: index, resource: {{ buffer }} }}))}});
+    const passEncoder = commandEncoder.beginComputePass();
+    passEncoder.setPipeline(pipeline);
+    passEncoder.setBindGroup(0, bindGroup);
+    passEncoder.dispatchWorkgroups(...workgroup);
+    passEncoder.end();
+  }};"""
+
+  with open(os.path.join(os.path.dirname(__file__), "net.js"), "w") as text_file:
+    text_file.write(prekernel + prg)

examples/webgpu/stable_diffusion/f16_to_f32.js

@@ -0,0 +1,64 @@
+const f16tof32 = `
+fn u16_to_f16(x: u32) -> f32 {
+    let sign = f32((x >> 15) & 0x1);
+    let exponent = f32((x >> 10) & 0x1F);
+    let fraction = f32(x & 0x3FF);
+
+    let sign_multiplier = select(1.0, -1.0, sign == 1.0);
+    if (exponent == 0.0) {
+        return sign_multiplier * 6.103515625e-5 * (fraction / 1024.0);
+    } else {
+        return sign_multiplier * exp2(exponent - 15.0) * (1.0 + fraction / 1024.0);
+    }
+}
+
+@group(0) @binding(0) var<storage,read_write> data0: array<u32>;
+@group(0) @binding(1) var<storage,read_write> data1: array<f32>;
+@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let gidx = gid.x;
+    let outgidx = gidx*2;
+
+    if (gidx >= arrayLength(&data0)) {
+        return;
+    }
+
+    let oo = data0[gidx];
+    let oo1 = (oo >> 16);
+    let oo2 = oo & 0xFFFFu;
+
+    let f1 = u16_to_f16(oo2);
+    let f2 = u16_to_f16(oo1);
+    
+    data1[outgidx] = f1;
+    data1[outgidx + 1] = f2;
+}`;
+
+window.f16tof32GPU = async(device, inf16) => {
+    const input = device.createBuffer({size: inf16.length, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });
+    const output = device.createBuffer({size: inf16.length*2, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });
+
+    const gpuWriteBuffer = device.createBuffer({size: input.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE });
+    const gpuReadBuffer = device.createBuffer({ size: output.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });
+    const commandEncoder = device.createCommandEncoder();
+    await gpuWriteBuffer.mapAsync(GPUMapMode.WRITE);
+
+    const alignedUint32View = new Uint32Array(inf16.buffer, inf16.byteOffset, inf16.length / 4);
+    new Uint32Array(gpuWriteBuffer.getMappedRange()).set(alignedUint32View);
+
+    gpuWriteBuffer.unmap();
+    commandEncoder.copyBufferToBuffer(gpuWriteBuffer, 0, input, 0, gpuWriteBuffer.size);
+    const pipeline = await device.createComputePipelineAsync({layout: "auto", compute: { module: device.createShaderModule({ code: f16tof32 }), entryPoint: "main" }});
+
+    addComputePass(device, commandEncoder, pipeline, [input, output], [Math.ceil(inf16.length/(4*256)), 1, 1]);
+
+    commandEncoder.copyBufferToBuffer(output, 0, gpuReadBuffer, 0, output.size);
+    const gpuCommands = commandEncoder.finish();
+    device.queue.submit([gpuCommands]);
+
+    await gpuReadBuffer.mapAsync(GPUMapMode.READ);
+    const resultBuffer = new Float32Array(gpuReadBuffer.size/4);
+    resultBuffer.set(new Float32Array(gpuReadBuffer.getMappedRange()));
+    gpuReadBuffer.unmap();
+
+    return resultBuffer;
+}

examples/webgpu/stable_diffusion/index.html

@@ -0,0 +1,511 @@
+<!DOCTYPE html>
+<html lang="en">
+
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>tinygrad has WebGPU</title>
+
+    <style>
+        body {
+            font-family: 'Arial', sans-serif;
+            text-align: center;
+            padding: 30px;
+        }
+
+        a {
+            text-decoration: none;
+            color: #4A90E2;
+        }
+
+        h1 {
+            font-size: 36px;
+            font-weight: normal;
+            margin-bottom: 20px;
+        }
+
+        #mybox {
+            display: flex;
+            flex-direction: column;
+            align-items: center;
+            gap: 20px;
+            width: 50%;
+            margin: 0 auto;
+        }
+
+        #promptText, #stepRange, #btnRunNet, #guidanceRange {
+            font-size: 18px;
+            width: 100%;
+        }
+
+        #result {
+            font-size: 48px;
+        }
+
+        #time {
+            font-size: 16px;
+            color: grey;
+        }
+
+        canvas {
+            margin-top: 20px;
+            border: 1px solid #000; 
+        }
+
+        label {
+            display: flex;
+            align-items: center;
+            gap: 10px;
+            width: 100%;
+        }
+        
+        #sliderValue {
+            margin-right: 10px;/
+        }
+    </style>
+
+    <script type="module">
+        import ClipTokenizer from 'https://softwired.nyc3.cdn.digitaloceanspaces.com/sd/clip_tokenizer.js';
+        window.clipTokenizer = new ClipTokenizer();
+    </script>
+    <script src="./f16_to_f32.js"></script>
+    <script src="./net.js"></script>
+</head>
+
+<body>
+<h1 id="wgpuError" style="display: none; color: red;">WebGPU is not supported in this browser</h1>
+<h1 id="sdTitle">StableDiffusion by <a href="https://github.com/tinygrad/tinygrad" target="_blank">tinygrad</a> WebGPU</h1>
+<div id="mybox">
+    <input id="promptText" type="text" placeholder="Enter your prompt here" value="a horse sized cat eating a bagel">
+    
+    <label>
+        Steps: <span id="stepValue">8</span>
+        <input id="stepRange" type="range" min="5" max="20" value="8" step="1">
+    </label>
+
+    <label>
+        Guidance: <span id="guidanceValue">7.5</span>
+        <input id="guidanceRange" type="range" min="3" max="15" value="7.5" step="0.1">
+    </label>
+    
+    <input id="btnRunNet" type="button" value="Run" disabled>
+
+    <div id="divModelDl" style="display: flex; align-items: center; width: 100%; gap: 10px;">
+        <span id="modelDlTitle">Downloading model</span>
+        <progress id="modelDlProgressBar" value="0" max="100" style="flex-grow: 1;"></progress>
+        <span id="modelDlProgressValue"></span>
+    </div>
+
+    <div id="divStepProgress" style="display: none; align-items: center; width: 100%; gap: 10px;">
+        <progress id="progressBar" value="0" max="100" style="flex-grow: 1;"></progress>
+        <span id="progressFraction"></span>
+    </div>
+</div>
+<canvas id="canvas" width="512" height="512"></canvas>
+
+<script>
+    function initDb() {
+        return new Promise((resolve, reject) => {
+            let db;
+            const request = indexedDB.open('tinydb', 1);
+            request.onerror = (event) => {
+                console.error('Database error:', event.target.error);
+                resolve(null);
+            };
+
+            request.onsuccess = (event) => {
+                db = event.target.result;
+                console.log("Db initialized.");
+                resolve(db);
+            };
+
+            request.onupgradeneeded = (event) => {
+                db = event.target.result;
+                if (!db.objectStoreNames.contains('tensors')) {
+                    db.createObjectStore('tensors', { keyPath: 'id' });
+                }
+            };
+        });
+    }
+
+    function saveTensorToDb(db, id, tensor) {
+        return new Promise((resolve, reject) => {
+            if (db == null) {
+                resolve(null);
+            }
+
+            const transaction = db.transaction(['tensors'], 'readwrite');
+            const store = transaction.objectStore('tensors');
+            const request = store.put({ id: id, content: tensor });
+
+            transaction.onabort = (event) => {
+                console.log("Transaction error while saving tensor: " + event.target.error);
+                resolve(null);
+            };
+
+            request.onsuccess = () => {
+                console.log('Tensor saved successfully.');
+                resolve();
+            };
+
+            request.onerror = (event) => {
+                console.error('Tensor save failed:', event.target.error);
+                resolve(null);
+            };
+        });
+    }
+
+    function readTensorFromDb(db, id) {
+        return new Promise((resolve, reject) => {
+            if (db == null) {
+                resolve(null);
+            }
+            
+            const transaction = db.transaction(['tensors'], 'readonly');
+            const store = transaction.objectStore('tensors');
+            const request = store.get(id);
+
+            transaction.onabort = (event) => {
+                console.log("Transaction error while reading tensor: " + event.target.error);
+                resolve(null);
+            };
+
+            request.onsuccess = (event) => {
+                const result = event.target.result;
+                if (result) {
+                    console.log("Cache hit: " + id);
+                    resolve(result);
+                } else {
+                    console.log("Cache miss: " + id);
+                    resolve(null);
+                }
+            };
+
+            request.onerror = (event) => {
+                console.error('Tensor retrieve failed: ', event.target.error);
+                resolve(null);
+            };
+        });
+    }
+
+    window.addEventListener('load', async function() {
+        if (!navigator.gpu) {
+            document.getElementById("wgpuError").style.display = "";
+            document.getElementById("sdTitle").style.display = "none";
+            return;
+        }
+
+        let db = await initDb();
+
+        const ctx = document.getElementById("canvas").getContext("2d", { willReadFrequently: true });
+        let labels, nets, safetensorParts;
+
+        const getDevice = async () => {
+            const adapter = await navigator.gpu.requestAdapter();
+            const requiredLimits = {};
+            const maxBufferSizeInSDModel = 1073741824;
+            requiredLimits.maxStorageBufferBindingSize = maxBufferSizeInSDModel;
+            requiredLimits.maxBufferSize = maxBufferSizeInSDModel;
+            
+            return await adapter.requestDevice({
+                requiredLimits
+            });
+        };
+
+        const timer = async (func, label = "") => {
+            const start = performance.now();
+            const out = await func();
+            const delta = (performance.now() - start).toFixed(1)
+            console.log(`${delta} ms ${label}`);
+            return out;
+        }
+
+        const getProgressDlForPart = async (part, progressCallback) => {
+            const response = await fetch(part);
+            const contentLength = response.headers.get('content-length');
+            const total = parseInt(contentLength, 10);
+
+            const res = new Response(new ReadableStream({
+                async start(controller) {
+                    const reader = response.body.getReader();
+                    for (;;) {
+                        const { done, value } = await reader.read();
+                        if (done) break;
+                        progressCallback(part, value.byteLength, total);
+                        controller.enqueue(value);
+                    }
+                    
+                    controller.close();
+                },
+            }));
+        
+            return res.arrayBuffer();
+        };
+  
+        const getAndDecompressF16Safetensors = async (device, progress) => {
+            let totalLoaded = 0;
+            let totalSize = 0;
+            let partSize = {};
+
+            const progressCallback = (part, loaded, total) => {
+                totalLoaded += loaded;
+
+                if (!partSize[part]) {
+                    totalSize += total;
+                    partSize[part] = true;
+                }
+                
+                progress(totalLoaded, totalSize);
+            };
+
+            let combinedBuffer = await readTensorFromDb(db, "net.f16");
+            let textModelU8 = await readTensorFromDb(db, "net.text");
+            let textModelFetched = false;
+
+            if (combinedBuffer == null) {
+                let dlParts = [
+                    getProgressDlForPart(window.MODEL_BASE_URL + '/net_part0.safetensors', progressCallback),
+                    getProgressDlForPart(window.MODEL_BASE_URL + '/net_part1.safetensors', progressCallback),
+                    getProgressDlForPart(window.MODEL_BASE_URL + '/net_part2.safetensors', progressCallback),
+                    getProgressDlForPart(window.MODEL_BASE_URL + '/net_part3.safetensors', progressCallback)
+                ];
+
+                if (textModelU8 == null) {
+                    dlParts.push(getProgressDlForPart(window.MODEL_BASE_URL + '/net_textmodel.safetensors', progressCallback));
+                }
+
+                let buffers = await Promise.all(dlParts);
+
+                // Combine everything except for text model, since that's alreafy f32
+                const totalLength = buffers.reduce((acc, buffer, index, array) => {
+                    if (index < 4) {
+                        return acc + buffer.byteLength;
+                    } else {
+                        return acc;
+                    }
+                    }, 0
+                );
+
+                combinedBuffer = new Uint8Array(totalLength);
+
+                let offset = 0;
+                buffers.forEach((buffer, index) => {
+                    if (index < 4) {
+                        combinedBuffer.set(new Uint8Array(buffer), offset);
+                        offset += buffer.byteLength;
+                        buffer = null;
+                    }
+                });
+
+                await saveTensorToDb(db, "net.f16", combinedBuffer);
+
+                if (textModelU8 == null) {
+                    textModelFetched = true;
+                    textModelU8 = new Uint8Array(buffers[4]);
+                    await saveTensorToDb(db, "net.text", textModelU8);
+                }
+            } else {
+                combinedBuffer = combinedBuffer.content;
+            }
+
+            if (textModelU8 == null) {
+                textModelU8 = new Uint8Array(await getProgressDlForPart(window.MODEL_BASE_URL + '/net_textmodel.safetensors', progressCallback));
+                await saveTensorToDb(db, "net.text", textModelU8);
+            } else if (!textModelFetched) {
+                textModelU8 = textModelU8.content;
+            }
+
+            document.getElementById("modelDlTitle").innerHTML = "Decompressing model";
+
+            const textModelOffset = 3772703308;
+            const metadataLength = Number(new DataView(combinedBuffer.buffer).getBigUint64(0, true));
+            const metadata = JSON.parse(new TextDecoder("utf8").decode(combinedBuffer.subarray(8, 8 + metadataLength)));
+
+            const allToDecomp = combinedBuffer.byteLength - (8 + metadataLength);
+            const decodeChunkSize = 67107840;
+            const numChunks = Math.ceil(allToDecomp/decodeChunkSize);
+
+            console.log(allToDecomp + " bytes to decompress");
+            console.log("Will be decompressed in " + numChunks+ " chunks");
+
+            let partOffsets = [{start: 0, end: 1131408336}, {start: 1131408336, end: 2227518416}, {start: 2227518416, end: 3308987856}, {start: 3308987856, end: 4265298864}];
+            let parts = [];
+
+            for (let offsets of partOffsets) {
+                parts.push(new Uint8Array(offsets.end-offsets.start));
+            }
+            parts[0].set(new Uint8Array(new BigUint64Array([BigInt(metadataLength)]).buffer), 0);
+            parts[0].set(combinedBuffer.subarray(8, 8 + metadataLength), 8);
+            parts[3].set(textModelU8, textModelOffset+8+metadataLength - partOffsets[3].start);
+
+            let start = Date.now();
+            let cursor = 0;
+
+            for (let i = 0; i < numChunks; i++) {
+                progress(i, numChunks);
+                let chunkStartF16 = 8 + metadataLength + (decodeChunkSize * i);
+                let chunkEndF16 = chunkStartF16 + decodeChunkSize;
+                let chunk = combinedBuffer.subarray(chunkStartF16, chunkEndF16);
+
+                if (chunk.byteLength %4 != 0) {
+                    const paddingBytes = 4 - (chunk.byteLength % 4);
+                    const alignedBuffer = new ArrayBuffer(chunk.byteLength + paddingBytes);
+                    const alignedView = new Uint8Array(alignedBuffer);
+                    alignedView.set(new Uint8Array(chunk));
+                    chunk = alignedView;
+                }
+
+                let result = await f16tof32GPU(device, chunk);
+                let resultUint8 = new Uint8Array(result.buffer);
+                let chunkStartF32 = 8 + metadataLength + (decodeChunkSize * i * 2);
+                let chunkEndF32 = chunkStartF32 + resultUint8.byteLength;
+                let offsetInPart = chunkStartF32 - partOffsets[cursor].start;
+
+                if (chunkEndF32 < partOffsets[cursor].end || cursor === parts.length - 1) {
+                    parts[cursor].set(resultUint8, offsetInPart);
+                } else {
+                    let spaceLeftInCurrentPart = partOffsets[cursor].end - chunkStartF32;
+                    parts[cursor].set(resultUint8.subarray(0, spaceLeftInCurrentPart), offsetInPart);
+
+                    cursor++;
+
+                    if (cursor < parts.length) {
+                        let nextPartOffset = spaceLeftInCurrentPart;
+                        let nextPartLength = resultUint8.length - nextPartOffset;
+                        parts[cursor].set(resultUint8.subarray(nextPartOffset, nextPartOffset + nextPartLength), 0);
+                    }
+                }
+
+                resultUint8 = null;
+                result = null;
+            }
+
+            combinedBuffer = null;
+
+            let end = Date.now();
+            console.log("Decoding took: " + ((end - start) / 1000) + " s");
+            console.log("Avarage " + ((end - start) / numChunks) + " ms per chunk");
+
+            return parts;
+        };
+
+        const loadNet = async () => {
+            const modelDlTitle = document.getElementById("modelDlTitle");
+
+            const progress = (loaded, total) => {
+                document.getElementById("modelDlProgressBar").value = (loaded/total) * 100
+                document.getElementById("modelDlProgressValue").innerHTML = Math.trunc((loaded/total) * 100) + "%"
+            }
+
+            const device = await getDevice();
+            safetensorParts = await getAndDecompressF16Safetensors(device, progress);
+
+            modelDlTitle.innerHTML = "Compiling model"
+
+            let models = ["textModel", "diffusor", "decoder"];
+
+            nets = await timer(() => Promise.all([
+                textModel().setup(device, safetensorParts),
+                diffusor().setup(device, safetensorParts),
+                decoder().setup(device, safetensorParts)
+            ]).then((loadedModels) => loadedModels.reduce((acc, model, index) => { acc[models[index]] = model; return acc; }, {})), "(compilation)")
+
+            progress(1, 1);
+
+            modelDlTitle.innerHTML = "Model ready"
+            setTimeout(() => {
+                document.getElementById("modelDlProgressBar").style.display = "none";
+                document.getElementById("modelDlProgressValue").style.display = "none";
+                document.getElementById("divStepProgress").style.display = "flex";
+            }, 1000);
+            document.getElementById("btnRunNet").disabled = false;
+        }
+
+        function runStableDiffusion(prompt, steps, guidance) {
+            return new Promise(async (resolve, reject) => {
+                let context = await timer(() => nets["textModel"](clipTokenizer.encodeForCLIP(prompt)));
+                let unconditional_context = await timer(() => nets["textModel"](clipTokenizer.encodeForCLIP("")));
+
+                let timesteps = [];
+
+                for (let i = 1; i < 1000; i += (1000/steps)) {
+                    timesteps.push(i);
+                }
+
+                console.log("Timesteps: " + timesteps);
+
+                let alphasCumprod = getWeight(safetensorParts,"alphas_cumprod");
+                let alphas = [];
+
+                for (t of timesteps) {
+                    alphas.push(alphasCumprod[Math.floor(t)]);
+                }
+
+                alphas_prev = [1.0];
+
+                for (let i = 0; i < alphas.length-1; i++) {
+                    alphas_prev.push(alphas[i]);
+                }
+
+                let inpSize = 4*64*64;
+                latent = new Float32Array(inpSize);
+
+                for (let i = 0; i < inpSize; i++) {
+                    latent[i] = Math.sqrt(-2.0 * Math.log(Math.random())) * Math.cos(2.0 * Math.PI * Math.random());
+                }
+                
+                for (let i = timesteps.length - 1; i >= 0; i--) {
+                    let timestep = new Float32Array([timesteps[i]]);
+                    let x_prev = await timer(() => nets["diffusor"](unconditional_context, context, latent, timestep, new Float32Array([alphas[i]]), new Float32Array([alphas_prev[i]]), new Float32Array([guidance])));
+                    latent = x_prev;
+                    document.getElementById("progressBar").value = ((steps - i) / steps) * 100
+                    document.getElementById("progressFraction").innerHTML = (steps - i) + "/" + steps
+                }
+                
+                resolve(await timer(() => nets["decoder"](latent)));
+            });
+        }
+
+        document.getElementById("btnRunNet").addEventListener("click", function(e) {
+            e.target.disabled = true;
+
+            runStableDiffusion(document.getElementById("promptText").value, document.getElementById("stepRange").value, document.getElementById("guidanceRange").value).then((image) => {
+                let pixels = []
+                let pixelCounter = 0
+
+                for (let j = 0; j < 512; j++) {
+                    for (let k = 0; k < 512; k++) {
+                        pixels.push(image[pixelCounter])
+                        pixels.push(image[pixelCounter+1])
+                        pixels.push(image[pixelCounter+2])
+                        pixels.push(255)
+                        pixelCounter += 3
+                    }
+                }
+                
+                ctx.putImageData(new ImageData(new Uint8ClampedArray(pixels), 512, 512), 0, 0);
+                console.log(image);
+                console.log("Success");
+                e.target.disabled = false;
+            });
+        }, false);
+
+        const stepSlider = document.getElementById('stepRange');
+        const stepValue = document.getElementById('stepValue');
+
+        stepSlider.addEventListener('input', function() {
+            stepValue.textContent = stepSlider.value;
+        });
+
+        const guidanceSlider = document.getElementById('guidanceRange');
+        const guidanceValue = document.getElementById('guidanceValue');
+
+        guidanceSlider.addEventListener('input', function() {
+            guidanceValue.textContent = guidanceSlider.value;
+        });
+
+        loadNet();
+    });
+</script>
+</body>
+</html>

examples/whisper.py

@@ -5,59 +5,79 @@ import pathlib
 import base64
 import multiprocessing
 import numpy as np
-from typing import Optional
-from extra.utils import download_file
+from typing import Optional, Union, Literal, List
+from tinygrad.jit import TinyJit
 from tinygrad.nn.state import torch_load, load_state_dict
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, DEBUG, CI, fetch
 import tinygrad.nn as nn
+from tinygrad.shape.symbolic import Variable
 from tinygrad.tensor import Tensor
 import itertools
 import librosa
 
-# TODO: you have written this fifteen times
 class MultiHeadAttention:
-  def __init__(self, n_state, n_head):
+  def __init__(self, n_state, n_head, kv_caching: Literal['cross', 'self']=None, max_self_attn_cache_len=None):
     self.n_head = n_head
     self.query = nn.Linear(n_state, n_state)
     self.key = nn.Linear(n_state, n_state, bias=False)
     self.value = nn.Linear(n_state, n_state)
     self.out = nn.Linear(n_state, n_state)
 
-  def __call__(self, x:Tensor, xa:Optional[Tensor]=None, mask:Optional[Tensor]=None):
+    self.kv_caching = kv_caching
+    self.max_self_attn_cache_len = max_self_attn_cache_len
+
+  def __call__(self, x:Tensor, xa:Optional[Tensor]=None, mask:Optional[Tensor]=None, len: Union[Variable,int]=None):
+    if self.kv_caching == 'cross':
+      if xa is not None:
+        k, v = self.key(xa), self.value(xa)
+        if not hasattr(self, 'cache_k'):
+          self.cache_k, self.cache_v = k, v
+        else:
+          # see test_jitted_read_assign in test_jit.py. more context https://github.com/tinygrad/tinygrad/pull/2360#issuecomment-1817989994
+          self.cache_k.assign(k+1-1).realize()
+          self.cache_v.assign(v+1-1).realize()
+      else:
+        k, v = self.cache_k, self.cache_v
+    else:
+      k, v = self.key(x), self.value(x)
+      if self.kv_caching == 'self':
+        if not hasattr(self, 'cache_k'):
+          self.cache_k = Tensor.zeros(x.shape[0], self.max_self_attn_cache_len, x.shape[2])
+          self.cache_v = Tensor.zeros(x.shape[0], self.max_self_attn_cache_len, x.shape[2])
+        k = self.cache_k.shrink((None, (0, len), None)).cat(k, dim=1)
+        v = self.cache_v.shrink((None, (0, len), None)).cat(v, dim=1)
+        padding = self.max_self_attn_cache_len-len-x.shape[1]
+        self.cache_k.assign(k.pad((None, (0, padding), None)).contiguous()).realize()
+        self.cache_v.assign(v.pad((None, (0, padding), None)).contiguous()).realize()
+
     q = self.query(x)
-    k = self.key(xa or x)
-    v = self.value(xa or x)
-    wv, qk = self.qkv_attention(q, k, v, mask)
-    # NOTE: we aren't returning qk
+    n_ctx = q.shape[1]
+    assert(q.shape[-1] == k.shape[-1] == v.shape[-1])
+    head_dim = q.shape[-1] // self.n_head
+    q = q.reshape(*q.shape[:2], self.n_head, head_dim).permute(0, 2, 1, 3)
+    k = k.reshape(*k.shape[:2], self.n_head, head_dim).permute(0, 2, 1, 3)
+    v = v.reshape(*v.shape[:2], self.n_head, head_dim).permute(0, 2, 1, 3)
+    attn = Tensor.scaled_dot_product_attention(q, k, v, mask[:n_ctx,:n_ctx] if mask is not None else None)
+    wv = attn.permute(0, 2, 1, 3).flatten(start_dim=2)
     return self.out(wv)
 
-  def qkv_attention(self, q, k, v, mask=None):
-    n_batch, n_ctx, n_state = q.shape
-    scale = (n_state // self.n_head) ** -0.25
-    q = q.reshape(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale
-    k = k.reshape(*k.shape[:2], self.n_head, -1).permute(0, 2, 3, 1) * scale
-    v = v.reshape(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
-    qk = q @ k
-    if mask is not None: qk = qk + mask[:n_ctx, :n_ctx]
-    w = qk.softmax(-1)
-    return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2), qk.detach()
 
 class ResidualAttentionBlock:
-  def __init__(self, n_state, n_head, cross_attention=False):
-    self.attn = MultiHeadAttention(n_state, n_head)
+  def __init__(self, n_state, n_head, is_decoder_block=False, max_self_attn_cache_len=None):
+    self.attn = MultiHeadAttention(n_state, n_head, kv_caching='self' if is_decoder_block else None, max_self_attn_cache_len=max_self_attn_cache_len)
     self.attn_ln = nn.LayerNorm(n_state)
 
-    self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None
-    self.cross_attn_ln = nn.LayerNorm(n_state) if cross_attention else None
+    self.cross_attn = MultiHeadAttention(n_state, n_head, kv_caching='cross') if is_decoder_block else None
+    self.cross_attn_ln = nn.LayerNorm(n_state) if is_decoder_block else None
 
     self.mlp = [nn.Linear(n_state, n_state*4), Tensor.gelu, nn.Linear(n_state*4, n_state)]
     self.mlp_ln = nn.LayerNorm(n_state)
 
-  def __call__(self, x, xa=None, mask=None):
-    x = x + self.attn(self.attn_ln(x), mask=mask)
+  def __call__(self, x, xa=None, mask=None, len: Union[Variable, int]=None):
+    x = x + self.attn(self.attn_ln(x), mask=mask, len=len)
     if self.cross_attn: x = x + self.cross_attn(self.cross_attn_ln(x), xa)
     x = x + self.mlp_ln(x).sequential(self.mlp)
-    return x
+    return x.realize()
 
 class AudioEncoder:
   def __init__(self, n_mels, n_audio_ctx, n_audio_state, n_audio_head, n_audio_layer, **_):
@@ -66,6 +86,7 @@ class AudioEncoder:
     self.blocks = [ResidualAttentionBlock(n_audio_state, n_audio_head) for _ in range(n_audio_layer)]
     self.ln_post = nn.LayerNorm(n_audio_state)
     self.positional_embedding = Tensor.empty(n_audio_ctx, n_audio_state)
+    self.encode = TinyJit(self.__call__)
 
   def __call__(self, x):
     x = self.conv1(x).gelu()
@@ -74,53 +95,89 @@ class AudioEncoder:
     x = x + self.positional_embedding[:x.shape[1]]
     x = x.sequential(self.blocks)
     x = self.ln_post(x)
-    return x
+    return x.realize()
 
 class TextDecoder:
   def __init__(self, n_vocab, n_text_ctx, n_text_state, n_text_head, n_text_layer, **_):
+    self.max_tokens_to_sample = n_text_ctx // 2
+    self.max_self_attn_cache_len = self.max_tokens_to_sample * 2 + 5  # roughly prompt + start toks + max_tokens_to_sample
+
     self.token_embedding = nn.Embedding(n_vocab, n_text_state)
     self.positional_embedding = Tensor.empty(n_text_ctx, n_text_state)
-    self.blocks = [ResidualAttentionBlock(n_text_state, n_text_head, cross_attention=True) for _ in range(n_text_layer)]
+    self.blocks = [ResidualAttentionBlock(n_text_state, n_text_head, is_decoder_block=True, max_self_attn_cache_len=self.max_self_attn_cache_len) for _ in range(n_text_layer)]
     self.ln = nn.LayerNorm(n_text_state)
-    #mask = torch.empty(n_ctx, n_ctx).fill_(-np.inf).triu_(1)
+    self.mask = Tensor.full((n_text_ctx, n_text_ctx), -np.inf).triu(1).realize()
+    self.blocks_start_tok = [TinyJit(block.__call__) for block in self.blocks]
+    self.blocks_after_start_tok = [TinyJit(block.__call__) for block in self.blocks]
+    self.start_output_tok = TinyJit(self.output_tok)
+    self.after_start_output_tok = TinyJit(self.output_tok)
 
-  def __call__(self, x, xa):
-    offset = 0
-    x = self.token_embedding(x) + self.positional_embedding[offset : offset + x.shape[-1]]
+  # if layernorm supported symbolic shapes, we wouldn't need this hacky 'streaming' param (which should be called something more descriptive like 'x_is_start_toks_only')
+  def __call__(self, x: Tensor, pos: int, encoded_audio: Tensor, streaming=False):
+    seqlen = x.shape[-1]
+    x = self.token_embedding(x) + self.positional_embedding[pos:pos+seqlen]
+    if pos == 0:
+      for block in (self.blocks if streaming else self.blocks_start_tok):
+        x = block(x, xa=encoded_audio, mask=self.mask, len=0)  # pass xa for cross attn kv caching
+      return self.output_tok(x) if streaming else self.start_output_tok(x)
+    else:
+      for block in self.blocks_after_start_tok:
+        len_v = Variable("self_attn_cache_len", 1, self.max_self_attn_cache_len).bind(pos)
+        x = block(x, mask=self.mask, len=len_v)
+      return self.after_start_output_tok(x)
 
-    seqlen, start_pos = x.shape[1], 0
-
-    mask = np.full((1, 1, seqlen, start_pos + seqlen), float("-inf"), dtype=np.float32)
-    mask = np.triu(mask, k=start_pos + 1)  # TODO: this is hard to do in tinygrad
-    mask = Tensor(mask)
-
-    for block in self.blocks: x = block(x, xa, mask)
-    x = self.ln(x)
-    return x @ self.token_embedding.weight.T
+  def output_tok(self, x):
+    return (self.ln(x) @ self.token_embedding.weight.T).realize()
 
 class Whisper:
-  def __init__(self, dims):
+  def __init__(self, dims, batch_size=1):
     self.encoder = AudioEncoder(**dims)
     self.decoder = TextDecoder(**dims)
+    self.is_multilingual = dims["n_vocab"] == 51865
+    self.batch_size = batch_size
 
-  def __call__(self, mel:Tensor, tokens:Tensor):
-    return self.decoder(tokens, self.encoder(mel))
 
 RATE = 16000
-CHUNK = 1600
-RECORD_SECONDS = 10
+SEGMENT_SECONDS=30
+SAMPLES_PER_SEGMENT = RATE * SEGMENT_SECONDS # 480000
+N_FFT = 400
+HOP_LENGTH = 160
+N_MELS = 80
+FRAMES_PER_SEGMENT = SAMPLES_PER_SEGMENT // HOP_LENGTH # 3000
 
-def prep_audio(waveform=None, sr=RATE) -> Tensor:
-  N_FFT = 400
-  HOP_LENGTH = 160
-  N_MELS = 80
-  if waveform is None: waveform = np.zeros(N_FFT, dtype=np.float32)
-  stft = librosa.stft(waveform, n_fft=N_FFT, hop_length=HOP_LENGTH, window='hann', dtype=np.float32)
-  magnitudes = stft[..., :-1] ** 2
-  mel_spec = librosa.filters.mel(sr=sr, n_fft=N_FFT, n_mels=N_MELS) @ magnitudes
-  log_spec = np.log10(np.clip(mel_spec, 1e-10, mel_spec.max() + 1e8))
+def prep_audio(waveforms: List[np.ndarray], batch_size: int, truncate=False) -> np.ndarray:
+  """
+  :param waveforms: A list of possibly variable length 16000Hz audio samples
+  :param batch_size: The batch_size associated with the Whisper model being used to transcribe the audio.
+                     Used to prevent JIT mismatch errors since the encoder does not accept symbolic shapes
+  :param truncate: If true, truncates (or pads) audio to exactly 30s for a single encoder pass
+  :return: mel spectrogram of the given waveforms
+  """
+  def pad_or_trim(arr, target_len):
+    curr_len = len(arr)
+    if curr_len == target_len:
+      return arr
+    elif curr_len < target_len:
+      return np.pad(arr, (0, target_len - curr_len), 'constant')
+    else:
+      return arr[:target_len]
+
+  max_len = SAMPLES_PER_SEGMENT if truncate else max(len(wav) for wav in waveforms)
+  if (r := max_len % SAMPLES_PER_SEGMENT) > 0: max_len += SAMPLES_PER_SEGMENT - r
+  waveforms = np.array(list(map(lambda w: pad_or_trim(w, max_len), waveforms)))
+  assert waveforms.shape[0] <= batch_size
+  if waveforms.shape[0] < batch_size:
+    # we could have a symbolic batch_size dim instead of manually padding here if conv/layernorm supported symbolic shapes
+    waveforms = np.pad(waveforms, pad_width=((0, batch_size - waveforms.shape[0]), (0, 0)))
+
+  stft = librosa.stft(waveforms, n_fft=N_FFT, hop_length=HOP_LENGTH, window='hann', dtype=np.csingle)
+  magnitudes = np.absolute(stft[..., :-1]) ** 2
+  mel_spec = librosa.filters.mel(sr=RATE, n_fft=N_FFT, n_mels=N_MELS) @ magnitudes
+
+  log_spec = np.log10(np.clip(mel_spec, 1e-10, None))
+  log_spec = np.maximum(log_spec, log_spec.max() - 8.0)
   log_spec = (log_spec + 4.0) / 4.0
-  #print(waveform.shape, log_spec.shape)
+
   return log_spec
 
 LANGUAGES = {
@@ -136,10 +193,9 @@ LANGUAGES = {
   "as": "assamese", "tt": "tatar", "haw": "hawaiian", "ln": "lingala", "ha": "hausa", "ba": "bashkir", "jw": "javanese", "su": "sundanese",
 }
 
-BASE = pathlib.Path(__file__).parent.parent / "weights"
-def get_encoding(n_vocab_in):
-  download_file("https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/gpt2.tiktoken", BASE / "gpt2.tiktoken")
-  ranks = {base64.b64decode(token): int(rank) for token, rank in (line.split() for line in open(BASE / "gpt2.tiktoken") if line)}
+def get_encoding(encoding_name):
+  with fetch(f"https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/{encoding_name}.tiktoken").open() as f:
+    ranks = {base64.b64decode(token): int(rank) for token, rank in (line.split() for line in f if line)}
   n_vocab = len(ranks)
   specials = [
     "<|endoftext|>",
@@ -155,82 +211,142 @@ def get_encoding(n_vocab_in):
   ]
   special_tokens = dict(zip(specials, itertools.count(n_vocab)))
   n_vocab += len(specials)
-  assert n_vocab == n_vocab_in
   import tiktoken
   return tiktoken.Encoding(
-    name="bob",
+    name=encoding_name,
     explicit_n_vocab=n_vocab,
     pat_str=r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""",
     mergeable_ranks=ranks,
     special_tokens=special_tokens)
 
-def img(x):
-  import matplotlib.pyplot as plt
-  plt.imshow(x.numpy())
-  plt.show()
+MODEL_URLS = {
+  "tiny.en": "https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt",
+  "tiny": "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt",
+  "base.en": "https://openaipublic.azureedge.net/main/whisper/models/25a8566e1d0c1e2231d1c762132cd20e0f96a85d16145c3a00adf5d1ac670ead/base.en.pt",
+  "base": "https://openaipublic.azureedge.net/main/whisper/models/ed3a0b6b1c0edf879ad9b11b1af5a0e6ab5db9205f891f668f8b0e6c6326e34e/base.pt",
+  "small.en": "https://openaipublic.azureedge.net/main/whisper/models/f953ad0fd29cacd07d5a9eda5624af0f6bcf2258be67c92b79389873d91e0872/small.en.pt",
+  "small": "https://openaipublic.azureedge.net/main/whisper/models/9ecf779972d90ba49c06d968637d720dd632c55bbf19d441fb42bf17a411e794/small.pt",
+  "medium.en": "https://openaipublic.azureedge.net/main/whisper/models/d7440d1dc186f76616474e0ff0b3b6b879abc9d1a4926b7adfa41db2d497ab4f/medium.en.pt",
+  "medium": "https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt",
+  "large-v1": "https://openaipublic.azureedge.net/main/whisper/models/e4b87e7e0bf463eb8e6956e646f1e277e901512310def2c24bf0e11bd3c28e9a/large-v1.pt",
+  "large-v2": "https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt",
+  "large": "https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt",
+}
+def init_whisper(model_name="tiny.en", batch_size=1):
+  assert MODEL_URLS[model_name] is not None
+
+  filename = fetch(MODEL_URLS[model_name])
+  state = torch_load(filename)
+  model = Whisper(state['dims'], batch_size)
+  load_state_dict(model, state['model_state_dict'], strict=False)
+  enc = get_encoding("multilingual" if model.is_multilingual else "gpt2")
+  return model, enc
+
+def load_file_waveform(filename):
+  waveform, _ = librosa.load(filename, sr=RATE)
+  return waveform
+
+def transcribe_file(model, enc, filename):
+  return transcribe_waveform(model, enc, [load_file_waveform(filename)])
+
+def transcribe_waveform(model, enc, waveforms, truncate=False):
+  """
+  Expects an array of shape (N,S) where N is the number waveforms to transcribe in parallel and S is number of 16000Hz samples
+  Returns the transcribed text if a single waveform is provided, or an array of transcriptions if multiple are provided
+  """
+  N_audio = len(waveforms)
+  log_spec = prep_audio(waveforms, model.batch_size, truncate)
+
+  if log_spec.shape[-1] > FRAMES_PER_SEGMENT and N_audio > 1:
+    # we don't support multi-segment batching because the size of the prompt tokens would be different for each item in the batch
+    # if we really want this feature, we can consider padding or trimming prompt tokens of varying lengths to make them consistent
+    raise Exception("Multi-segment transcription not supported with batch audio input")
+
+  start_tokens = [enc._special_tokens["<|startoftranscript|>"]]
+  if model.is_multilingual:
+    # TODO detect language
+    language_token = enc._special_tokens["<|startoftranscript|>"] + 1 + tuple(LANGUAGES.keys()).index("en")
+    start_tokens.append(language_token)
+    start_tokens.append(enc._special_tokens["<|transcribe|>"])
+  start_tokens.append(enc._special_tokens["<|notimestamps|>"])
+  transcription_start_index = len(start_tokens)
+  eot = enc._special_tokens["<|endoftext|>"]
+  transcription_tokens = [np.array([], dtype=np.int32)] * log_spec.shape[0]
+
+  for curr_frame in range(0, log_spec.shape[-1], FRAMES_PER_SEGMENT):
+    encoded_audio = model.encoder.encode(Tensor(log_spec[:, :, curr_frame:curr_frame + FRAMES_PER_SEGMENT]))
+    pos = 0
+    curr_segment_tokens = np.tile(start_tokens, (log_spec.shape[0], 1))
+    if curr_frame > 0:
+      # pass the previously inferred tokens as 'prompt' - https://github.com/openai/whisper/discussions/117#discussioncomment-3727051
+      prompt = np.concatenate((
+        [enc._special_tokens["<|startofprev|>"]],
+        transcription_tokens[0][-model.decoder.max_tokens_to_sample+1:],
+        start_tokens))
+      curr_segment_tokens = np.tile(prompt, (log_spec.shape[0], 1))
+      transcription_start_index = len(curr_segment_tokens[0])
+
+    for i in range(model.decoder.max_tokens_to_sample):
+      out = model.decoder(Tensor(curr_segment_tokens if i == 0 else curr_segment_tokens[:, -1:]), pos, encoded_audio, streaming=curr_frame > 0)
+      next_tokens = out[:, -1].argmax(axis=-1).numpy().astype(np.int32)
+      next_tokens[curr_segment_tokens[:, -1] == eot] = eot
+      curr_segment_tokens = np.concatenate((curr_segment_tokens, next_tokens.reshape(-1, 1)), axis=1)
+      pos = curr_segment_tokens.shape[-1] - 1
+      if DEBUG >= 1: print(i, list(map(lambda tokens: enc.decode(tokens), curr_segment_tokens)))
+      if (curr_segment_tokens[:, -1] == eot).all():
+        break
+
+    for i, t in enumerate(curr_segment_tokens):
+      eot_index = np.where(t == eot)[0]
+      eot_index = None if len(eot_index) == 0 else eot_index[0]
+      transcription_tokens[i] = np.concatenate((transcription_tokens[i], t[transcription_start_index:eot_index]))
+
+  transcriptions = list(map(lambda tokens: enc.decode(tokens).strip(), transcription_tokens))
+  return transcriptions[:N_audio] if N_audio > 1 else transcriptions[0]
+
+CHUNK = 1600
+RECORD_SECONDS = 10
 
 def listener(q):
-  prep_audio()
   import pyaudio
   p = pyaudio.PyAudio()
   stream = p.open(format=pyaudio.paInt16, channels=1, rate=RATE, input=True, frames_per_buffer=CHUNK)
   print("listening")
   for _ in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
     data = stream.read(CHUNK)
-    waveform = ((np.frombuffer(data, np.int16)/32768).astype(np.float32)*3).reshape(1, -1)
+    waveform = ((np.frombuffer(data, np.int16)/32768).astype(np.float32)*3)
     q.put(waveform)
   print("done listening")
 
 if __name__ == "__main__":
-  if getenv("SMALL"):
-    fn = BASE / "whisper-small.en.pt"
-    download_file("https://openaipublic.azureedge.net/main/whisper/models/f953ad0fd29cacd07d5a9eda5624af0f6bcf2258be67c92b79389873d91e0872/small.en.pt", fn)
-  else:
-    fn = BASE / "whisper-tiny.en.pt"
-    download_file("https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt", fn)
-  state = torch_load(fn)
-  model = Whisper(state['dims'])
-  load_state_dict(model, state['model_state_dict'])
-  enc = get_encoding(state['dims']['n_vocab'])
+  model, enc = init_whisper("small.en" if getenv("SMALL") else "tiny.en", batch_size=1)
 
   if len(sys.argv) > 1:
-    # offline
-    waveform, sample_rate = librosa.load(sys.argv[1], normalize=True)
-    log_spec = prep_audio(waveform, sample_rate)
-    lst = [enc._special_tokens["<|startoftranscript|>"]]
-    dat = model.encoder(Tensor(log_spec)).realize()
-    for i in range(50):
-      out = model.decoder(Tensor([lst]), dat)
-      out.realize()
-      idx = out[0,-1].argmax().numpy()
-      lst.append(idx)
-      print(enc.decode(lst))
+    print(transcribe_file(model, enc, sys.argv[1]))
   else:
     # online
-
     q = multiprocessing.Queue()
     p = multiprocessing.Process(target=listener, args=(q,))
     p.daemon = True
     p.start()
 
-    lst = [enc._special_tokens["<|startoftranscript|>"]]
+    lst = [enc._special_tokens["<|startoftranscript|>"], enc._special_tokens["<|notimestamps|>"]]
     total = None
     did_read = False
     for i in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
       while not q.empty() or total is None:
         waveform = q.get()
         if total is None: total = waveform
-        else: total = np.concatenate([total, waveform], axis=1)
+        else: total = np.concatenate([total, waveform])
         did_read = True
       if did_read:
-        last_total = total.shape[1]
-        log_spec = prep_audio(waveform=Tensor(total).numpy(), sr=RATE)
-        encoded_audio = model.encoder(Tensor(log_spec)).realize()
-      out = model.decoder(Tensor([lst]), encoded_audio).realize()
-      idx = out[0,-1].argmax().numpy()
+        log_spec = prep_audio(total.reshape(1, -1), model.batch_size, truncate=True)
+        encoded_audio = model.encoder.encode(Tensor(log_spec))
+      # pass the previously inferred tokens as 'prefix' - https://github.com/openai/whisper/discussions/117#discussioncomment-3727051
+      out = model.decoder(Tensor([lst]), 0, encoded_audio, streaming=True).realize()
+      idx = int(out[0,-1].argmax().numpy().item())
       lst.append(idx)
       dec = enc.decode(lst)
       print(dec) # DO NOT REMOVE PRINT. IT'S VERY IMPORTANT
       if dec.endswith("<|endoftext|>"):
-        #total = total[:, 320*(len(lst)-1):]
-        lst = [enc._special_tokens["<|startoftranscript|>"]]
+        lst.pop()

examples/yolov3.py

@@ -8,10 +8,10 @@ import numpy as np
 from PIL import Image
 from tinygrad.tensor import Tensor
 from tinygrad.nn import BatchNorm2d, Conv2d
-from extra.utils import fetch
+from tinygrad.helpers import fetch
 
 def show_labels(prediction, confidence=0.5, num_classes=80):
-  coco_labels = fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names')
+  coco_labels = fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names').read_bytes()
   coco_labels = coco_labels.decode('utf-8').split('\n')
   prediction = prediction.detach().numpy()
   conf_mask = (prediction[:,:,4] > confidence)
@@ -38,7 +38,7 @@ def show_labels(prediction, confidence=0.5, num_classes=80):
 def add_boxes(img, prediction):
   if isinstance(prediction, int): # no predictions
     return img
-  coco_labels = fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names')
+  coco_labels = fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names').read_bytes()
   coco_labels = coco_labels.decode('utf-8').split('\n')
   height, width = img.shape[0:2]
   scale_factor = 608 / width
@@ -281,7 +281,7 @@ class Darknet:
           print("None biases for layer", i)
 
   def load_weights(self, url):
-    weights = np.frombuffer(fetch(url), dtype=np.float32)[5:]
+    weights = np.frombuffer(fetch(url).read_bytes(), dtype=np.float32)[5:]
     ptr = 0
     for i in range(len(self.module_list)):
       module_type = self.blocks[i + 1]["type"]
@@ -369,7 +369,7 @@ class Darknet:
     return detections
 
 if __name__ == "__main__":
-  model = Darknet(fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg'))
+  model = Darknet(fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg').read_bytes())
   print("Loading weights file (237MB). This might take a while…")
   model.load_weights('https://pjreddie.com/media/files/yolov3.weights')
   if len(sys.argv) > 1:
@@ -392,7 +392,7 @@ if __name__ == "__main__":
     cap.release()
     cv2.destroyAllWindows()
   elif url.startswith('http'):
-    img_stream = io.BytesIO(fetch(url))
+    img_stream = io.BytesIO(fetch(url).read_bytes())
     img = cv2.imdecode(np.frombuffer(img_stream.read(), np.uint8), 1)
   else:
     img = cv2.imread(url)

examples/yolov8-onnx.py

@@ -2,11 +2,12 @@
 import os
 from ultralytics import YOLO
 import onnx
+from pathlib import Path
 from extra.onnx import get_run_onnx
 from tinygrad.tensor import Tensor
 
 os.chdir("/tmp")
-if not os.path.isfile("yolov8n-seg.onnx"):
+if not Path("yolov8n-seg.onnx").is_file():
   model = YOLO("yolov8n-seg.pt")
   model.export(format="onnx", imgsz=[480,640])
 onnx_model = onnx.load(open("yolov8n-seg.onnx", "rb"))

examples/yolov8.py

@@ -2,15 +2,13 @@ from tinygrad.nn import Conv2d, BatchNorm2d
 from tinygrad.tensor import Tensor
 import numpy as np
 from itertools import chain
-from extra.utils import get_child, fetch, download_file
 from pathlib import Path
 import cv2
 from collections import defaultdict
-import os
-import time, io, sys
+import time, sys
+from tinygrad.helpers import fetch
 from tinygrad.nn.state import safe_load, load_state_dict
 
-
 #Model architecture from https://github.com/ultralytics/ultralytics/issues/189
 #The upsampling class has been taken from this pull request https://github.com/tinygrad/tinygrad/pull/784 by dc-dc-dc. Now 2(?) models use upsampling. (retinet and this)
 
@@ -245,7 +243,7 @@ class Upsample:
     tmp = x.reshape([b, c, -1] + [1] * _lens) * Tensor.ones(*[1, 1, 1] + [self.scale_factor] * _lens)
     return tmp.reshape(list(x.shape) + [self.scale_factor] * _lens).permute([0, 1] + list(chain.from_iterable([[y+2, y+2+_lens] for y in range(_lens)]))).reshape([b, c] + [x * self.scale_factor for x in x.shape[2:]])
 
-class Conv_Block():
+class Conv_Block:
   def __init__(self, c1, c2, kernel_size=1, stride=1, groups=1, dilation=1, padding=None):
     self.conv = Conv2d(c1,c2, kernel_size, stride, padding=autopad(kernel_size, padding, dilation), bias=False, groups=groups, dilation=dilation)
     self.bn = BatchNorm2d(c2, eps=0.001)
@@ -398,13 +396,12 @@ if __name__ == '__main__':
   yolo_variant = sys.argv[2] if len(sys.argv) >= 3 else (print("No variant given, so choosing 'n' as the default. Yolov8 has different variants, you can choose from ['n', 's', 'm', 'l', 'x']") or 'n')
   print(f'running inference for YOLO version {yolo_variant}')
 
-  output_folder_path = './outputs_yolov8'
-  if not os.path.exists(output_folder_path):
-    os.makedirs(output_folder_path)
+  output_folder_path = Path('./outputs_yolov8')
+  output_folder_path.mkdir(parents=True, exist_ok=True)
   #absolute image path or URL
-  image_location = [np.frombuffer(io.BytesIO(fetch(img_path)).read(), np.uint8)]
+  image_location = [np.frombuffer(fetch(img_path).read_bytes(), np.uint8)]
   image = [cv2.imdecode(image_location[0], 1)]
-  out_paths = [os.path.join(output_folder_path, img_path.split("/")[-1].split('.')[0] + "_output" + '.' + img_path.split("/")[-1].split('.')[1])]
+  out_paths = [(output_folder_path / f"{Path(img_path).stem}_output{Path(img_path).suffix}").as_posix()]
   if not isinstance(image[0], np.ndarray):
     print('Error in image loading. Check your image file.')
     sys.exit(1)
@@ -414,10 +411,7 @@ if __name__ == '__main__':
   depth, width, ratio = get_variant_multiples(yolo_variant)
   yolo_infer = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)
 
-  weights_location = Path(__file__).parent.parent / "weights" / f'yolov8{yolo_variant}.safetensors'
-  download_file(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{yolo_variant}.safetensors', weights_location)
-
-  state_dict = safe_load(weights_location)
+  state_dict = safe_load(fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{yolo_variant}.safetensors'))
   load_state_dict(yolo_infer, state_dict)
 
   st = time.time()
@@ -427,8 +421,7 @@ if __name__ == '__main__':
   post_predictions = postprocess(preds=predictions, img=pre_processed_image, orig_imgs=image)
 
   #v8 and v3 have same 80 class names for Object Detection
-  class_labels = fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names')
-  class_labels = class_labels.decode('utf-8').split('\n')
+  class_labels = fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names').read_text().split("\n")
 
   draw_bounding_boxes_and_save(orig_img_paths=image_location, output_img_paths=out_paths, all_predictions=post_predictions, class_labels=class_labels)
 

extra/accel/ane/lib/ane.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-import os
+from pathlib import Path
 from ctypes import *
 import json
 import collections
@@ -8,13 +8,13 @@ import faulthandler
 import struct
 faulthandler.enable()
 
-basedir = os.path.dirname(os.path.abspath(os.path.realpath(__file__)))
+basedir = Path(__file__).resolve().parent
 
 libane = None
 aneregs = None
 def init_libane():
   global libane, aneregs
-  libane = cdll.LoadLibrary(os.path.join(basedir, "libane.dylib"))
+  libane = cdll.LoadLibrary((basedir / "libane.dylib").as_posix())
 
   libane.ANE_Compile.argtypes = [c_char_p, c_int]
   libane.ANE_Compile.restype = c_void_p
@@ -29,7 +29,7 @@ def init_libane():
 
   #libane.ANE_RegDebug.restype = c_char_p
 
-  with open(os.path.join(basedir, "aneregs.json")) as f:
+  with open(basedir / "aneregs.json") as f:
     aneregs = json.load(f)
 
 ANE_Struct = [
@@ -214,7 +214,7 @@ if __name__ == "__main__":
   mdf = ane.pack(dd, md)
   assert(md == mdf)
 
-  comp = ane.compile(dat) 
+  comp = ane.compile(dat)
   ret = ane.run(comp, tin, tout)
   print("** after **")
   print(tind)

extra/accel/triton/ops_triton.py

@@ -1,141 +0,0 @@
-from __future__ import annotations
-import hashlib
-from weakref import WeakValueDictionary
-from torch import float32
-import numpy as np
-import pycuda.autoprimaryctx # type: ignore # noqa: F401
-import pycuda.driver as cuda # type: ignore
-
-import triton # type: ignore # noqa: F401
-import triton.language as tl  # type: ignore # noqa: F401
-
-from typing import Union, Tuple, Optional, Dict
-from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, LazyOp, Op, ExplicitExecAST, GlobalCounters
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.helpers import prod, DEBUG
-from tinygrad.runtime.cuda import CLBuffer
-from tinygrad.compiler.ast import ASTKernel
-
-stream = cuda.Stream()
-
-class TritonASTKernel(ASTKernel):
-  code_for_op : Dict[Op, str] = {
-    UnaryOps.NOOP: "(A)", UnaryOps.NEG: "(-(A))", UnaryOps.RELU: "tl.maximum(A, 0.0)", UnaryOps.GT0: "tl.where(A>0,1,0)",
-    UnaryOps.EXP: "tl.exp(A)", UnaryOps.LOG: "tl.log(A)", UnaryOps.RECIPROCAL: "(1.0/A)", UnaryOps.SQRT: "tl.sqrt(A)",
-    BinaryOps.ADD: "(A+B)", BinaryOps.SUB: "(A-B)", BinaryOps.MUL: "(A*B)",
-    BinaryOps.DIV: "(A/B)", BinaryOps.CMPEQ: "(A==B)",
-    ReduceOps.SUM: "A += B", ReduceOps.MAX: "A = tl.maximum(A,B)"
-  }
-  start_for_op = {ReduceOps.SUM: "0.0", ReduceOps.MAX: "float('-inf')"}
-
-  def ast_parse(self, x:Union[TritonBuffer, LazyOp], acc:str, do_reduce=False) -> str:
-    if not isinstance(x, LazyOp):
-      # this is a load
-      buf_index = self.bufs.index(x)
-      if buf_index not in self.loaded:
-        idx, valid = self.sts[buf_index].expr_idxs()
-        valid_expr = valid.render().replace("&&", "*1*")
-        self.kernel.append(self.kernel_prefix + f"  val{buf_index} = tl.where({valid_expr}, tl.load(data{buf_index} + {idx.render()}, mask={valid_expr}), 0.0)")
-        self.loaded.add(buf_index)
-      return f"val{buf_index}"
-    if isinstance(x.op, ReduceOps) and not do_reduce: return acc
-
-    values = ([acc] if isinstance(x.op, ReduceOps) else []) + [self.ast_parse(v, acc, do_reduce) for v in x.src]
-
-    code = TritonASTKernel.code_for_op[x.op]  # TODO: replace this with a function
-    code = code.replace("A", values[0])
-    if len(values) == 2: code = code.replace("B", values[1])
-    return code
-
-  func_cache: WeakValueDictionary = WeakValueDictionary()
-  def codegen(self):
-    if self.key in self.func_cache: return self.func_cache[self.key]
-
-    self.process()
-    self.kernel_prefix = ""
-    self.loaded = set()
-    self.kernel = ["@triton.jit"]
-    self.kernel.append("def fxn("+','.join(f"data{i}" for i in range(len(self.bufs)))+"):")
-
-    self.output_shape = list(self.sts[0].shape[:self.first_reduce])
-
-    # copied from ops_gpu
-    # TODO CUDA only supports a grid of (2^31-1, 65535, 65535), that results in invalid kernel launches for some shapes, so flattern the grid for now.
-    MAX_OUTPUT_SHAPE = 1
-    self.kernel += [f"  idx{len(self.output_shape)-1-i} = tl.program_id({i})" for i in range(min(MAX_OUTPUT_SHAPE, len(self.output_shape)))]
-    if len(self.output_shape) > MAX_OUTPUT_SHAPE:
-      final_dimension = len(self.output_shape)-MAX_OUTPUT_SHAPE
-      for i in range(final_dimension-1, -1, -1):
-        self.kernel += [f"  idx{i} = idx{final_dimension} % {self.output_shape[i]}", f"  idx{final_dimension} = idx{final_dimension} // {self.output_shape[i]}"]
-      self.output_shape = [prod(self.output_shape[0:final_dimension+1])] + list(self.output_shape[final_dimension+1:])
-      if DEBUG >= 3: print(f"replaced output shape with {self.output_shape}")
-    elif len(self.output_shape) == 0: self.output_shape = [1]
-
-    if self.reduceop:
-      full_shape = [st.shape for st in self.sts if st.shape != self.sts[0].shape]
-      full_shape = self.sts[0].shape if len(full_shape) == 0 else full_shape[0]
-      self.kernel += [f"  acc = {TritonASTKernel.start_for_op[self.reduceop.op]}"]
-      self.kernel += [("  "*(i-self.first_reduce)+f"  for idx{i} in range(0, {full_shape[i]}):") for i in range(self.first_reduce, self.shape_len)]
-      self.kernel_prefix =  "  "*(self.shape_len - self.first_reduce)
-      self.kernel.append("  "+self.kernel_prefix+self.ast_parse(self.reduceop, "acc", True))
-      self.kernel_prefix =  ""
-
-    code = self.ast_parse(self.ast, "acc")
-
-    # store
-    idx, valid = self.sts[0].expr_idxs()
-    self.kernel.append(f"  tl.store(data0 + {idx.render()}, {code})")
-
-    # Torch inductor seems to write out files too!
-    hash = hashlib.md5(self.key.encode('utf-8')).hexdigest()
-    fn = f"/tmp/{hash}.py"
-    kernel = '\n'.join(self.kernel)
-    if DEBUG >= 4: print(kernel)
-    with open(fn, "w") as f: f.write(kernel)
-    codeObject = compile(kernel, fn, "exec")
-    exec(codeObject, globals())
-    program = globals()['fxn']
-    mem_estimate = sum(prod(x._base_shape) for x in self.bufs)
-    def runner(*bufs):
-      GlobalCounters.log_kernel(self.info.flops, mem_estimate)
-      return program[tuple(self.output_shape[::-1])](*[x.cuda for x in bufs], stream=stream.handle)
-    self.func_cache[self.key] = runner
-    return runner
-
-class TritonBuffer(ExplicitExecAST):
-  def __init__(self, shape:Union[ShapeTracker, Tuple[int, ...]], hostbuf:Optional[TritonBuffer]=None, backing:Optional[np.ndarray]=None, force_create=False):
-    super().__init__(shape, hostbuf)
-    self._buf : Optional[TritonDeviceAllocation] = hostbuf._buf if hostbuf is not None else None
-    self._base_shape : Tuple[int, ...] = hostbuf._base_shape if hostbuf is not None else self.shape
-    self._backing : Optional[np.ndarray] = hostbuf._backing if hostbuf is not None else backing
-    if force_create: self.cuda
-
-  @property
-  def cuda(self):
-    if self._buf is None:
-      self._buf = TritonDeviceAllocation(4*prod(self._base_shape))
-      if self._backing is not None: self._buf.copyin(self._backing, stream)
-    return self._buf
-
-  @staticmethod
-  def fromCPU(x): return TritonBuffer(x.shape, backing=x.view(np.ndarray).astype(np.float32).ravel())
-
-  def toCPU(self):
-    data = np.empty(self.shape, dtype=np.float32)
-    buf = self.contiguous()
-    buf.cuda
-    buf._buf.copyout(data)
-    return data
-
-  @classmethod
-  def exec_ast(cls, ast:LazyOp, output_buffer:Optional[TritonBuffer]=None):
-    k = TritonASTKernel(ast, output_buffer)
-    k.codegen()(*k.bufs)
-    return k.ret
-
-class TritonDeviceAllocation(CLBuffer):
-  def __init__(self, size):
-    super().__init__(size)
-    self.dtype = float32
-
-  def data_ptr(self): return int(self._cl)

extra/archprobe.py

@@ -2,7 +2,7 @@
 import numpy as np
 import pickle
 from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer
-from tinygrad.helpers import dtypes
+from tinygrad import dtypes
 from tqdm import trange, tqdm
 from matplotlib import pyplot as plt
 

extra/assembly/assembly.py

@@ -1,13 +1,14 @@
 from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
-from tinygrad.codegen.linearizer import UOps, Token, ConstOp, MemOp, UOp
+from tinygrad.codegen.linearizer import UOps, MemOp, UOp
 from tinygrad.ops import BinaryOps, UnaryOps
-from tinygrad.helpers import DType, dtypes, DEBUG
+from tinygrad.dtype import DType, dtypes
+from tinygrad.helpers import DEBUG
 from tinygrad.shape.symbolic import Variable, NumNode, MulNode, DivNode, ModNode, LtNode, SumNode, AndNode
 import functools
 import math
 from collections import defaultdict
 
-_type_to_letter = {dtypes.float32: 'f', dtypes.bool: 'p', dtypes.int32: 'i', dtypes.int64: 'a', dtypes.uint32: 'u', dtypes.uint64: 'b', dtypes._float4: 'x', dtypes.uint8: 'uc', dtypes.float16: 'h',
+_type_to_letter = {dtypes.float32: 'f', dtypes.bool: 'p', dtypes.int32: 'i', dtypes.int64: 'a', dtypes.uint32: 'u', dtypes.uint64: 'b', dtypes.float.vec(4): 'x', dtypes.uint8: 'uc', dtypes.float16: 'h',
                    dtypes.int8: 'c', dtypes.uint16: 'us', dtypes.float64: 'd'}
 
 class Register(NamedTuple):
@@ -17,7 +18,7 @@ class Register(NamedTuple):
   off:Optional[int] = None
   def __repr__(self): return self.nm if self.off is None else f"{self.nm}:{self.off}"
   def subregs(self):
-    if self.dtype == dtypes._float4:
+    if self.dtype == dtypes.float.vec(4):
       return [Register(self.nm, dtypes.float, False, off=off) for off in range(4)]
     return []
 
@@ -38,16 +39,15 @@ class AssemblyLanguage:
   ins: List[AssemblyInstruction] = []
 
   def type_to_letter(self,x): return _type_to_letter[x[0]].upper() if x[1] else _type_to_letter[x[0]]
-  def newreg(self, tok, dtype=dtypes.float32, scalar=False):
-    if isinstance(tok, Token): dtype = tok.dtype  # this
+  def newreg(self, tok, dtype=dtypes.float32, scalar=False) -> Register:
     self.tor[tok] = ret = Register(f"%{self.type_to_letter((dtype, scalar))}{self.cnts[(dtype, scalar)]}", dtype, scalar)
-    if dtype == dtypes._float4:
+    if dtype == dtypes.float.vec(4):
       for off in range(4):
-        self.tor[Token(tok.name, tok.dtype, off)] = Register(ret.nm, dtypes.float, ret.scalar, off)
+        self.tor[tok] = Register(ret.nm, dtypes.float, ret.scalar, off)
     self.cnts[(dtype, scalar)] += 1
     return ret
 
-  def render_numnode(self, b):
+  def render_numnode(self, b) -> Register:
     key = ("num", b)
     if key not in self.tor: self.ins.append(AssemblyInstruction(UOps.LOAD, self.newreg(key, scalar=True, dtype=dtypes.int32), [], b))
     return self.tor[key]
@@ -98,11 +98,12 @@ def uops_to_asmstyle(lang, function_name:str, uops:List[UOp]):
   lang.ins.clear()
   lang.tor.clear()
   lang.cnts.clear()
-  buf_to_dtype = {args[0]:args[1] for uop,_,_,args in uops if uop == UOps.DEFINE_GLOBAL}
+  buf_to_dtype = {args:dtype for uop,dtype,_,args,_ in uops if uop == UOps.DEFINE_GLOBAL}
   global_size, local_size = [], []
   skipload_branch = 0
   lang.ins += [AssemblyInstruction(UOps.SPECIAL, lang.newreg(buf, dtype=dtypes.uint64, scalar=True), [], buf) for buf in buf_to_dtype]
-  for uop,newvar,vin,args in uops:
+  for u in uops:
+    uop,dtype,vin,args,_ = u
     if uop == UOps.DEFINE_LOCAL:
       lang.ins.append(AssemblyInstruction(UOps.DEFINE_LOCAL, None, [], args))
       lang.ins.append(AssemblyInstruction(UOps.ALU, lang.newreg(args[0], dtype=dtypes.uint64), [args[0]], UnaryOps.NOOP))
@@ -133,57 +134,55 @@ def uops_to_asmstyle(lang, function_name:str, uops:List[UOp]):
       elif args[1] == 'local':
         for i, var in enumerate(reversed(args[0])):
           lang.ins.append(AssemblyInstruction(UOps.ENDLOOP, None, [lang.tor[var]], (var.max+1, f"lid{i}")))
-    elif uop == UOps.CAST and newvar is not None:
+    elif uop == UOps.CAST:
       # TODO: we should reconsider outputting CAST in the linearizer. these are needless copies
-      out = lang.newreg(newvar)
+      out = lang.newreg(u, dtype)
       for i,sr in enumerate(out.subregs()):
         lang.ins.append(AssemblyInstruction(UOps.ALU, sr, [lang.tor[vin[i]]], UnaryOps.NOOP))
-    elif uop == UOps.ALU and newvar is not None:
-      out = lang.newreg(newvar) if newvar not in lang.tor else lang.tor[newvar]
+    elif uop == UOps.ALU:
+      out = lang.newreg(u, dtype) if u not in lang.tor else lang.tor[u]
       # this is the only thing that can violate SSA
       if args in [BinaryOps.CMPLT]:
-        pred_reg = lang.newreg((newvar, 'pred'), dtype=dtypes.bool)
+        pred_reg = lang.newreg((u, 'pred'), dtype=dtypes.bool)
         lang.ins.append(AssemblyInstruction(UOps.ALU, pred_reg, [lang.tor[x] for x in vin], args))
         lang.ins.append(AssemblyInstruction(UOps.CAST, out, [pred_reg], args))
       elif args == BinaryOps.DIV and lang.no_div:
-        tmp = lang.newreg((newvar, "rcp"))
+        tmp = lang.newreg((u, "rcp"))
         lang.ins.append(AssemblyInstruction(UOps.ALU, tmp, [lang.tor[vin[1]]], UnaryOps.RECIP))
         lang.ins.append(AssemblyInstruction(UOps.ALU, out, [lang.tor[vin[0]], tmp], BinaryOps.MUL))
       elif args == UnaryOps.SIN and lang.sin_is_sin2pi:
-        tmp = lang.newreg((newvar, "2pi"))
+        tmp = lang.newreg((u, "2pi"))
         lang.ins.append(AssemblyInstruction(UOps.ALU, tmp, [lang.tor[vin[0]], 1/(math.pi*2)], BinaryOps.MUL))
         lang.ins.append(AssemblyInstruction(UOps.ALU, out, [tmp], args))
       else:
         lang.ins.append(AssemblyInstruction(UOps.ALU, out, [lang.tor[x] for x in vin], args))
-    elif uop == UOps.LOAD and newvar is not None:
-      if isinstance(args, ConstOp):
-        if args.valid.min == 0 and args.valid.max == 1:
-          reg = lang.newreg(newvar, dtype=newvar.dtype)
-          lang.ins.append(AssemblyInstruction(UOps.LOAD, reg, [], args.invalid_value))
+    elif uop == UOps.DEFINE_ACC:
+      reg = lang.newreg(u, dtype=dtype)
+      lang.ins.append(AssemblyInstruction(UOps.LOAD, reg, [], args))
+    elif uop == UOps.SPECIAL:
+      lang.tor[u] = lang.tor[args]
+    elif uop == UOps.CONST:
+      lang.ins.append(AssemblyInstruction(UOps.LOAD, lang.newreg(u, dtype=dtype), [], args))
+    elif uop == UOps.LOAD:
+      idx, treg, off = lang.addr_w_offset(args)
+      reg = lang.newreg(u, dtype=dtype, scalar=(idx.scalar and (not isinstance(treg, Register) or treg.scalar)))
+      if args.valid.min == 0:
+        lang.ins.append(AssemblyInstruction(UOps.LOAD, reg, [], 0))
+        if args.valid.max == 1:
           pred = args.valid.render(lang.render_ops, lang)
           lang.ins.append(AssemblyInstruction(UOps.COND_BRANCH, None, [pred], (f"$skipload_{skipload_branch}", False)))
-          lang.ins.append(AssemblyInstruction(UOps.LOAD, reg, [], args.value))
-          lang.ins.append(AssemblyInstruction(UOps.LABEL, None, [], f"$skipload_{skipload_branch}"))
-          skipload_branch += 1
-        else:
-          lang.ins.append(AssemblyInstruction(UOps.LOAD, lang.newreg(newvar, dtype=newvar.dtype), [], args.value if args.valid.min == 1 else args.invalid_value))
+      if args.valid.max == 1:
+          # NOTE: you can't compute the index in here, because it assumes it's all available later
+        lang.ins.append(AssemblyInstruction(UOps.LOAD, reg, [idx] + ([treg] if treg is not None else []), (off, 'global' if not args.local else 'shared', args.memory_dtype if args.memory_dtype != dtypes.float else None)))
+      if args.valid.min == 0 and args.valid.max == 1:
+        lang.ins.append(AssemblyInstruction(UOps.LABEL, None, [], f"$skipload_{skipload_branch}"))
+        skipload_branch += 1
+    elif uop == UOps.STORE:
+      if args is None:
+        lang.ins.append(AssemblyInstruction(UOps.ALU, lang.tor[vin[0]], [lang.tor[vin[1]]], UnaryOps.NOOP))
       else:
         idx, treg, off = lang.addr_w_offset(args)
-        reg = lang.newreg(newvar, dtype=newvar.dtype, scalar=(idx.scalar and (not isinstance(treg, Register) or treg.scalar))) # and not dtypes.is_float(newvar.dtype)))
-        if args.valid.min == 0:
-          lang.ins.append(AssemblyInstruction(UOps.LOAD, reg, [], 0))
-          if args.valid.max == 1:
-            pred = args.valid.render(lang.render_ops, lang)
-            lang.ins.append(AssemblyInstruction(UOps.COND_BRANCH, None, [pred], (f"$skipload_{skipload_branch}", False)))
-        if args.valid.max == 1:
-            # NOTE: you can't compute the index in here, because it assumes it's all available later
-          lang.ins.append(AssemblyInstruction(UOps.LOAD, reg, [idx] + ([treg] if treg is not None else []), (off, 'global' if not args.local else 'shared', args.memory_dtype if args.memory_dtype != dtypes.float else None)))
-        if args.valid.min == 0 and args.valid.max == 1:
-          lang.ins.append(AssemblyInstruction(UOps.LABEL, None, [], f"$skipload_{skipload_branch}"))
-          skipload_branch += 1
-    elif uop == UOps.STORE:
-      idx, treg, off = lang.addr_w_offset(args)
-      lang.ins.append(AssemblyInstruction(UOps.STORE, None, [idx, lang.tor[vin[0]]] + ([treg] if treg is not None else []), (off, 'global' if not args.local else 'shared', args.memory_dtype if args.memory_dtype != dtypes.float else None)))
+        lang.ins.append(AssemblyInstruction(UOps.STORE, None, [idx, lang.tor[vin[0]]] + ([treg] if treg is not None else []), (off, 'global' if not args.local else 'shared', args.memory_dtype if args.memory_dtype != dtypes.float else None)))
 
   if DEBUG >= 4:
     for tins in lang.ins: print(tins)

extra/assembly/assembly_arm64.py

@@ -1,9 +1,10 @@
 import struct
 from platform import system
 from typing import Tuple, Dict, List, Optional
+from tinygrad import dtypes
 from tinygrad.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.codegen.linearizer import UOps, UOp
-from tinygrad.helpers import dtypes, CI
+from tinygrad.helpers import CI
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
 
 def float_to_hex(x): return "%02X%02X%02X%02X" % tuple(struct.pack("f",x)[::-1])
@@ -25,6 +26,7 @@ def specialize_to_arm64(fn_nm, asm):
   type_to_reg = {dtypes.double: "d", dtypes.half: 'h', dtypes.float32: 's', dtypes.bool: 'w', dtypes.int8:'w', dtypes.int32: 'w', dtypes.int64: 'x', dtypes.uint8:'w', dtypes.uint32: 'w', dtypes.uint64: 'x'}
   alu = {BinaryOps.ADD: "add", BinaryOps.SUB: "sub", BinaryOps.MUL: "mul", BinaryOps.DIV: "div", BinaryOps.MAX: "max",
           BinaryOps.MOD: "", BinaryOps.CMPLT: "subs",
+          UnaryOps.NOOP: "mov", UnaryOps.NEG: "neg",
           UnaryOps.SIN:'bl ' + get_name('sinf'), UnaryOps.LOG2: 'bl ' + get_name("log2f"), UnaryOps.EXP2: 'bl ' + get_name("exp2f"), UnaryOps.SQRT: 'bl ' + get_name("sqrtf"),
           TernaryOps.MULACC: "madd", TernaryOps.WHERE: "fcsel"}
 
@@ -90,9 +92,14 @@ def specialize_to_arm64(fn_nm, asm):
         ins.append(f"loop_{arg}:")
     elif uop == UOps.CAST:
       if arg == BinaryOps.CMPLT:
-        mov_imm(0.0, 's0')
-        mov_imm(1.0, 's1')
-        ins.append(f"fcsel {rtor[out.nm]}, s1, s0, lt")
+        if rtor[out.nm][0] == 's':
+          mov_imm(0.0, 's0')
+          mov_imm(1.0, 's1')
+          ins.append(f"fcsel {rtor[out.nm]}, s1, s0, lt")
+        if rtor[out.nm][0] == 'x':
+          mov_imm(0, 'x14')
+          mov_imm(1, 'x15')
+          ins.append(f"csel {rtor[out.nm]}, x15, x14, lt")
       else:
         ins.append(f"sxtw {rtor[out.nm]}, w{rtor[vin[0].nm][1:]}")
     elif uop == UOps.ALU:
@@ -100,7 +107,7 @@ def specialize_to_arm64(fn_nm, asm):
       if arg == BinaryOps.MUL and out.dtype == dtypes.bool:
         ins.append(f"ands {','.join('x15' if v.__class__ is int else rtor[v.nm] for v in [out] + vin)}")
       elif arg == TernaryOps.WHERE:
-        ins.append(f"fcmp {rtor[vin[0].nm]}, #0.0")
+        ins.append(f"fcmp {rtor[vin[0].nm]}, #0.0" if rtor[vin[0].nm][0] == 's' else f"cmp {rtor[vin[0].nm]}, #0")
         ins.append(f"{alu[arg]} {rtor[out.nm]}, {rtor[vin[1].nm]}, {rtor[vin[2].nm]}, ne")
       elif arg in [UnaryOps.LOG2, UnaryOps.SIN, UnaryOps.EXP2, UnaryOps.SQRT]:
         #NOTE: Not a real instruction, use to emulate a ext call in unicorn
@@ -124,8 +131,9 @@ def specialize_to_arm64(fn_nm, asm):
       elif arg == BinaryOps.CMPLT:
         ins.append(f"{alu[arg]} {','.join('x15' if v.__class__ is int else rtor[v.nm] for v in [out] + vin)}" if not dtypes.is_float(vin[0][1]) else f"fcmp {rtor[vin[0].nm]}, {rtor[vin[1].nm]}")
       elif arg == BinaryOps.MOD:
-        ins.append(f"udiv x14, {rtor[vin[0].nm]}, x15")
-        ins.append(f"msub {rtor[out.nm]}, x14, x15, {rtor[vin[0].nm]}")
+        rhs = 'x15' if vin[1].__class__ is int else rtor[vin[1].nm]
+        ins.append(f"udiv x14, {rtor[vin[0].nm]}, {rhs}")
+        ins.append(f"msub {rtor[out.nm]}, x14, {rhs}, {rtor[vin[0].nm]}")
       else:
         ins.append(f"{'f' if dtypes.is_float(vin[0][1]) else 's' if arg == BinaryOps.DIV else ''}{alu[arg]} {', '.join('x15' if v.__class__ is int else rtor[v.nm] for v in [out] + vin)}")
     elif uop == UOps.LOAD:

extra/assembly/assembly_ptx.py

@@ -2,7 +2,7 @@ from typing import List
 import struct
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
 from tinygrad.codegen.linearizer import UOps, UOp
-from tinygrad.helpers import dtypes
+from tinygrad import dtypes
 from tinygrad.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.runtime.ops_cuda import arch
 
@@ -15,7 +15,10 @@ def render_cast(ins, inp, out):
   if inp.dtype == dtypes.bool and (dtypes.is_float(out.dtype) or dtypes.is_int(out.dtype)):
     ins.append(f"selp.{dtype_to_nvtype[out.dtype]} {out}, {'0f3F800000, 0f00000000' if dtypes.is_float(out.dtype) else '1, 0'}, {inp};")
   elif out.dtype == dtypes.bool:
-    ins.append(f"setp.ne.{dtype_to_nvtype[inp.dtype]} {out}, {'0f00000000' if dtypes.is_float(inp.dtype) else '0'}, {inp};")
+    if inp.dtype == dtypes.bool:
+      ins.append(f"mov.pred {out}, {inp};")
+    else:
+      ins.append(f"setp.ne.{dtype_to_nvtype[inp.dtype]} {out}, {'0f00000000' if dtypes.is_float(inp.dtype) else '0'}, {inp};")
   else:
     round_mod = ".rzi" if dtypes.is_int(out.dtype) and dtypes.is_float(inp.dtype) else '.rz' if dtypes.is_float(out.dtype) and (dtypes.is_int(inp.dtype) or dtypes.is_float(inp.dtype) and inp.dtype.itemsize > out.dtype.itemsize) else ''
     ins.append(f"cvt{round_mod}.{dtype_to_nvtype[out.dtype]}.{dtype_to_nvtype[inp.dtype]} {out}, {inp};")
@@ -30,7 +33,8 @@ def specialize_to_ptx(lang, function_name):
   ins = []
   alu = {BinaryOps.ADD: "add", BinaryOps.SUB: "sub", BinaryOps.MUL: "mul", BinaryOps.DIV: "div", BinaryOps.MAX: "max",
          BinaryOps.MOD: "rem", BinaryOps.CMPLT: "setp.lt", UnaryOps.SQRT: "sqrt.approx",
-         UnaryOps.NOOP: "mov", UnaryOps.SIN: "sin.approx", UnaryOps.LOG2: "lg2.approx", UnaryOps.EXP2: "ex2.approx.ftz",
+         UnaryOps.NOOP: "mov", UnaryOps.NEG: "neg",
+         UnaryOps.SIN: "sin.approx", UnaryOps.LOG2: "lg2.approx", UnaryOps.EXP2: "ex2.approx.ftz",
          TernaryOps.MULACC: "fma.rn", TernaryOps.WHERE: "selp"}
   for uop, out, vin, arg in lang.ins:
     if uop == UOps.ENDLOOP:
@@ -53,8 +57,11 @@ def specialize_to_ptx(lang, function_name):
       else:
         otype = vin[0].dtype if arg in [BinaryOps.CMPLT] else out.dtype
         if arg == TernaryOps.WHERE:
-          reg = lang.newreg((vin[0], 'bool'), dtypes.bool)
-          ins.append(f"setp.ne.{dtype_to_nvtype[vin[0].dtype]} {reg}, {'0f00000000' if dtypes.is_float(vin[0].dtype) else '0'}, {vin[0]};")
+          if vin[0].dtype == dtypes.bool:
+            reg = vin[0]
+          else:
+            reg = lang.newreg((vin[0], 'bool'), dtypes.bool)
+            ins.append(f"setp.ne.{dtype_to_nvtype[vin[0].dtype]} {reg}, {'0f00000000' if dtypes.is_float(vin[0].dtype) else '0'}, {vin[0]};")
           vin = vin[1:] + [reg]
         ins.append(f"{alu[arg]}{'.lo' if arg == BinaryOps.MUL and out.dtype != dtypes.float32 else ''}{'.rn' if arg == BinaryOps.DIV and out.dtype == dtypes.float32 else ''}.{dtype_to_nvtype[otype]} {out}, {', '.join(str(x) for x in vin)};")
     elif uop == UOps.LOAD:

extra/assembly/assembly_rdna.py

@@ -1,6 +1,6 @@
 import yaml
 from typing import Tuple, Set, Dict
-from tinygrad.helpers import dtypes
+from tinygrad import dtypes
 from tinygrad.codegen.assembly import AssemblyCodegen, Register
 from tinygrad.codegen.linearizer import UOps
 from tinygrad.ops import BinaryOps, UnaryOps, TernaryOps
@@ -62,7 +62,7 @@ class RDNACodegen(AssemblyCodegen):
       return rtor[x]
     for uop, out, vin, arg in asm:
       if uop == UOps.DEFINE_REGISTER:
-        if arg[0][0] in [dtypes.uint32, dtypes.uint64, dtypes.int64, dtypes.int32, dtypes.float32, dtypes._float4]:
+        if arg[0][0] in [dtypes.uint32, dtypes.uint64, dtypes.int64, dtypes.int32, dtypes.float32, dtypes.float.vec(4)]:
           for i in range(arg[2]):
             # TODO: Re-use gaps created by this to avoid wasting registers
             align = int(arg[0][0].itemsize / 4)
@@ -76,7 +76,7 @@ class RDNACodegen(AssemblyCodegen):
               v_cnt += align
             rtor[Register(f"%{arg[1]}{i}", *arg[0])] = reg_name
 
-            if arg[0][0] == dtypes._float4:
+            if arg[0][0] == dtypes.float.vec(4):
               for off in range(4):
                 reg_name = f"s{s_cnt-align+off}" if arg[0][1] else f"v{v_cnt-align+off}"
                 rtor[Register(f"%{arg[1]}{i}", dtypes.float, False, off=off)] = reg_name
@@ -109,7 +109,7 @@ class RDNACodegen(AssemblyCodegen):
       elif uop == UOps.CONST:
         if arg == float('inf'): arg = "0x7f800000"
         elif arg == float('-inf'): arg = "0xff800000"
-        if out.dtype == dtypes._float4:
+        if out.dtype == dtypes.float.vec(4):
           for off in range(4):
             ins.append(f"{'s_' if out.scalar else 'v_'}mov_b32 {reg_out(Register(out.nm, dtypes.float, False, off=off))}, {arg}")
         else:
@@ -122,8 +122,8 @@ class RDNACodegen(AssemblyCodegen):
           if arg == TernaryOps.MULACC and out == vin[2]:
             alu_arg = "fmac"
             vin = vin[0:2]
-          if out.dtype == dtypes._float4:
-            for rr in zip(*[x.subregs() if x.dtype == dtypes._float4 else [x,x,x,x] for x in [out]+vin]):
+          if out.dtype == dtypes.float.vec(4):
+            for rr in zip(*[x.subregs() if x.dtype == dtypes.float.vec(4) else [x,x,x,x] for x in [out]+vin]):
               ins.append(f"{'s_' if rr[0].scalar else 'v_'}{alu_arg}_{dtype_to_rdnatype[rr[0].dtype]} {reg_out(rr[0])}, {', '.join(reg_in(x) if x.__class__ is Register else str(x) for x in rr[1:])}")
           else:
             ins.append(f"{'s_' if out.scalar else 'v_'}{alu_arg}_{dtype_to_rdnatype[out.dtype] if arg != UnaryOps.NOOP else 'b32'}{'_i24' if arg == BinaryOps.MUL and out.dtype != dtypes.float32 and not out.scalar else ''} {reg_out(out)}, {', '.join(reg_in(x) if x.__class__ is Register else str(x) for x in vin)}")
@@ -132,11 +132,11 @@ class RDNACodegen(AssemblyCodegen):
           # swap arg order
           ins.append(f's_load_b32 {reg_out(out)}, {reg_in(vin[0])}, {reg_in(vin[1])} offset:{arg[0]}')
         else:
-          ins.append(f'global_load_{"b128" if out.dtype == dtypes._float4 else "b32"} {reg_out(out)}, {reg_in(vin[1])}, {reg_in(vin[0])} offset:{arg[0]}')
+          ins.append(f'global_load_{"b128" if out.dtype == dtypes.float.vec(4) else "b32"} {reg_out(out)}, {reg_in(vin[1])}, {reg_in(vin[0])} offset:{arg[0]}')
         pend_regs.add(out)
         for r in out.subregs(): pend_regs.add(r)
       elif uop == UOps.STORE:
-        ins.append(f'global_store_{"b128" if vin[1].dtype == dtypes._float4 else "b32"} {reg_in(vin[2])}, {reg_in(vin[1])}, {reg_in(vin[0])} offset:{arg[0]}')
+        ins.append(f'global_store_{"b128" if vin[1].dtype == dtypes.float.vec(4) else "b32"} {reg_in(vin[2])}, {reg_in(vin[1])}, {reg_in(vin[0])} offset:{arg[0]}')
       elif uop == UOps.LABEL:
         ins.append(f"{arg}:")
       elif uop == UOps.COND_BRANCH:

extra/assembly/rocm/rdna3/asm.py

@@ -67,13 +67,13 @@ with open("/tmp/cc2.elf", "wb") as f:
   f.write(asm)
 
 print(colored("creating CLProgram", "green"))
-prg = CLProgram("code", asm, binary=True)
+prg = CLProgram("code", asm)
 
 print(colored("running program", "green"))
 G = 512
 FLOPS *= 100000*G*G  # loop * global_size
 for i in range(3):
-  tm = prg([G//256, G], [256, 1], buf, wait=True)
+  tm = prg(buf, global_size=[G//256, G, 1], local_size=[256, 1, 1], wait=True)
   print(f"ran in {tm*1e3:.2f} ms, {FLOPS/(tm*1e9):.2f} GFLOPS")
 
 print(colored("transferring buffer", "green"))

extra/augment.py

@@ -1,9 +1,10 @@
 import numpy as np
 from PIL import Image
-import os
+from pathlib import Path
 import sys
-sys.path.append(os.getcwd())
-sys.path.append(os.path.join(os.getcwd(), 'test'))
+cwd = Path.cwd()
+sys.path.append(cwd.as_posix())
+sys.path.append((cwd / 'test').as_posix())
 from extra.datasets import fetch_mnist
 from tqdm import trange
 

extra/autopad.py

@@ -0,0 +1,36 @@
+from tinygrad.tensor import Tensor
+from tinygrad.ops import LoadOps
+from tinygrad.codegen.linearizer import Linearizer
+from test.external.fuzz_linearizer import run_linearizer
+from tinygrad.codegen.kernel import Opt, OptOps
+
+N = 17**3
+
+a = Tensor.rand(N, N)
+b = Tensor.rand(N, N)
+c = a @ b
+sched = [si for si in c.lazydata.schedule() if si.ast.op not in LoadOps]
+assert len(sched) == 1
+lin = Linearizer(sched[0].ast)
+
+lin.apply_opt(Opt(op=OptOps.PADTO, axis=0, amt=32))
+lin.apply_opt(Opt(op=OptOps.PADTO, axis=1, amt=32))
+lin.hand_coded_optimizations()
+lin.linearize()
+print(f"{lin.applied_opts=}")
+
+run_linearizer(lin)
+
+###
+
+a = Tensor.rand(61, 61).sum(axis=0)
+sched = [si for si in a.lazydata.schedule() if si.ast.op not in LoadOps]
+assert len(sched) == 1
+lin = Linearizer(sched[0].ast)
+
+lin.apply_opt(Opt(op=OptOps.PADTO, axis=0, amt=32))
+lin.hand_coded_optimizations()
+lin.linearize()
+print(f"{lin.applied_opts=}")
+
+run_linearizer(lin)

extra/backends/cstyle.py

@@ -0,0 +1,70 @@
+# TODO: how much of this can be merged with above?
+class WGSLLanguage(CStyleLanguage):
+  code_for_workitem = {"g": lambda x: f"i32(gindex.{'xyz'[x]})", "l": lambda x: f"i32(lindex.{'xyz'[x]})"}
+  size_prefix = "let"
+  barrier="workgroupBarrier();"
+  generic_var_prefix = "var "
+  external_local_bufs = True
+  code_for_op = { **CStyleLanguage().code_for_op,
+                 BinaryOps.CMPLT: lambda x,y,dtype: f"f32({x}<{y})", BinaryOps.CMPEQ: lambda x,y,dtype: f"f32({x}=={y})",
+                 TernaryOps.MULACC: lambda x,y,z,dtype: f"fma({x},{y},{z})", TernaryOps.WHERE: lambda a,b,c,dtype: f"select({c},{b},bool({a}))" }
+  # HACK: write bool as f32
+  type_map = {dtypes.float: "f32", dtypes.half: "f16", dtypes.int32: "i32", dtypes.uint32: "u32", dtypes.bool: "f32"}
+
+  def render_local(self, name: str, dtype:DType, size: int): return f"var<workgroup> {name}: array<{self.type_map[dtype]},{size}>;"
+
+  def render_const(self, x:Union[float,int], var_dtype) -> str:
+    if math.isnan(x): return "nan()"
+    elif math.isinf(x): return ("-" if x < 0 else "") + "inf(1.0)"
+    return f"({super().render_const(x, var_dtype)})"
+
+  def render_if(self, cond: str): return f"if (bool({cond})) {{"
+
+  def render_kernel(self, function_name:str, kernel:List[str], bufs:List[Tuple[str,DType]], local_size:List[int], prekernel:List[str]) -> str:
+    local_size = local_size[::-1] if local_size else [1]
+    bind_it = iter(range(len(bufs)))
+    prg = "fn nan() -> f32 { let bits = 0xffffffffu; return bitcast<f32>(bits); }\nfn inf(a: f32) -> f32 { return a/0.0; }\n"
+    prg += "\n".join(prekernel+[f"@group(0) @binding({next(bind_it)}) {'var<storage,read_write>' if isinstance(dtype, PtrDType) else 'var<uniform>'} {name}: {f'array<{self.type_map[dtype]}>' if isinstance(dtype, PtrDType) else 'i32'};" for name,dtype in bufs])  # noqa: E501
+    prg += f"\n@compute @workgroup_size({','.join([str(x) for x in local_size])}) fn {function_name}(@builtin(workgroup_id) gindex: vec3<u32>, @builtin(local_invocation_id) lindex: vec3<u32>) {{\n" + "\n".join(kernel) + "\n}"  # noqa: E501
+    return prg
+
+  def render_cast(self, x:List[str], var_dtype:DType, bitcast=False) -> str:
+    if self.type_map[var_dtype]: return f"bitcast<{self.type_map[var_dtype]}>({x[0]})" if bitcast else f"{self.type_map[var_dtype]}({x[0]})"
+    raise NotImplementedError(f"no cast for {var_dtype}")
+WGSLRenderer = functools.partial(uops_to_cstyle, WGSLLanguage())
+
+
+class GLSLLanguage(CStyleLanguage):
+  type_map = {dtypes.float: "float", dtypes.half: "float", dtypes.int32: "int", dtypes.uint32: "uint", dtypes.bool: "bool"}
+  sampler_prefix = {dtypes.float64: "d", dtypes.float: "", dtypes.half: "", dtypes.int32: "i", dtypes.uint32: "u", dtypes.bool: "i"}
+  fragment_center_offset = 0.5
+  code_for_workitem = {"i": lambda x, offset=fragment_center_offset:f"int(gl_FragCoord.y-{offset}) * width + int(gl_FragCoord.x-{offset})"}
+  code_for_op = {**CStyleLanguage().code_for_op, **{op: lambda a,b,dtype,charforop=charforop: f"bool(int({a}){charforop}int({b}))" \
+    if dtype == dtypes.bool else f"({a}{charforop}{b})" for op,charforop in [(BinaryOps.MUL,"*"),(BinaryOps.ADD,"+"),(BinaryOps.DIV,"/")]},
+    BinaryOps.CMPLT: lambda a,b,dtype: f"(float({a})<float({b}))" if dtype == dtypes.bool else f"({a}<{b})",
+    BinaryOps.MOD: lambda a,b,dtype: f"(int({a})%int({b}))", TernaryOps.WHERE: lambda a,b,c,dtype: f"(float({a})!=0.0?{b}:{c})"}
+
+  def render_const(self, x:Union[float,int], var_dtype) -> str:
+    if math.isnan(x): return "(0.0 / 0.0)"
+    elif math.isinf(x): return ("-" if x < 0 else "") + "(1./0.)"
+    return self.render_cast(["({:.1f})".format(x) if x == int(x) and dtypes.is_float(var_dtype) else f"({x})"]*var_dtype.sz, var_dtype)
+
+  def render_kernel(self, function_name:str, kernel:List[str], bufs:List[Tuple[str,DType]], local_size:List[int], prekernel:List[str]) -> str:
+    prg = "#version 330\nprecision highp float;\nprecision highp int;\nin vec2 uv;\nuniform int width;\n"
+    prg += "\n".join([f"uniform {self.sampler_prefix[dtype]}sampler2D {name};" for name,dtype in bufs if name != "data0"])
+    prg += f"\nout {'int' if bufs[0][1] == dtypes.bool else self.type_map[bufs[0][1]]} out_data;\n"
+    return prg + "\nvoid main() {\n" + "\n".join(kernel) + "\n}"
+
+  def render_cast(self, x:List[str], var_dtype:DType, bitcast=False) -> str:
+    if self.type_map[var_dtype]: return f"{self.type_map[var_dtype]}({x[0]})"
+    raise NotImplementedError(f"no cast for {var_dtype}")
+
+  def render_load(self, output_dtype, buf_name, buf_dtype, idx, local=False) -> str:
+    x_calc = f"float(int({idx})%textureSize({buf_name}, 0).x)"
+    y_calc = f"float(int({idx})/textureSize({buf_name}, 0).x)"
+    out_val = f"texture({buf_name}, vec2(float({x_calc} + {self.fragment_center_offset}f)/float(textureSize({buf_name}, 0).x),\
+    float({y_calc} + {self.fragment_center_offset}f)/float(textureSize({buf_name}, 0).y))).r"
+    return f"{self.render_cast([out_val], output_dtype)}"
+
+  def render_store(self, buf_name:str, buf_dtype:DType, var_name:str, var_dtype:DType, idx, local=False) -> str:
+    return f"out_data = {'int' if buf_dtype == dtypes.bool else self.type_map[buf_dtype]}({var_name});"

extra/backends/ops_webgl.py

@@ -0,0 +1,52 @@
+import numpy as np
+import functools
+from tinygrad.dtype import dtypes, ImageDType
+from tinygrad.device import Compiled, Allocator
+from tinygrad.codegen.kernel import LinearizerOptions, OptOps
+from tinygrad.renderer.cstyle import uops_to_cstyle
+from tinygrad.renderer.cstyle import GLSLLanguage
+import moderngl
+
+ctx = moderngl.create_standalone_context()
+max_dims = 4096
+dtype_map = { dtypes.float64: "f8", dtypes.float: "f4", dtypes.half: "f2", dtypes.int32: "i4", dtypes.uint32: "u4", dtypes.bool: "i1"}
+vertex_shader="#version 330\nprecision highp float;\nin vec2 in_position;in vec2 in_uv;out vec2 uv;void main(){\
+gl_Position=vec4(in_position,0.0,1.0);uv=in_uv;}"
+class WebGLProgram:
+  def __init__(self, name: str, prg: str, bufs:int=0, vars:int=0):
+    self.name, self.prg = name, ctx.program(vertex_shader=vertex_shader, fragment_shader=prg)
+  def __call__(self, *bufs, global_size, local_size=None, vals=(), wait=False):
+    vert = ctx.buffer(np.asarray([-1, 1, -1, -1, 1, 1, 1, -1], dtype='f4').tobytes())
+    uv = ctx.buffer(np.asarray([0, 1, 0, 0, 1, 1, 1, 0], dtype='f4').tobytes())
+    self.vao = ctx.vertex_array(self.prg, [])
+    self.vao.bind(self.prg["in_position"].location if "in_position" in self.prg else 0, buffer=vert, cls='f', fmt='2f4')
+    self.vao.bind(self.prg["in_uv"].location if "in_uv" in self.prg else 1, buffer=uv, cls='f', fmt='2f4')
+    self.vao.vertices = vert.size//4//2
+    self.fbo = ctx.framebuffer(color_attachments=[bufs[0]])
+
+    for i, x in enumerate(bufs[1:], start=1):
+      if f"data{i}" in self.prg:
+        self.prg[f"data{i}"] = i
+        x.use(i)
+
+    if ("width" in self.prg): self.prg["width"].value = self.fbo.size[0]
+    ctx.viewport = (0, 0, self.fbo.size[0], self.fbo.size[1])
+    self.fbo.use()
+    self.vao.render(mode=moderngl.TRIANGLE_STRIP)
+
+class RawWebGLAllocator(Allocator):
+  def _alloc_image(self, dtype:ImageDType):
+    tex = ctx.texture(dtype.shape, 1, dtype=dtype_map[dtype.base])
+    tex.filter = (moderngl.NEAREST, moderngl.NEAREST)
+    return tex
+  def copyin(self, dest:moderngl.Texture, src: memoryview): dest.write(src)
+  def copyout(self, dest:memoryview, src: moderngl.Texture):
+    src.read_into(dest)
+    return dest
+
+class WebGlDevice(Compiled):
+  def __init__(self, device:str):
+    super().__init__(RawWebGLAllocator(),
+      LinearizerOptions(device="WEBGL", global_max=[4096*4096,1,1], unsupported_opts=[OptOps.UPCAST, OptOps.UPCASTMID],
+      supports_float4=False, supports_float4_alu=False, has_local=False, has_shared=False, dont_use_locals=True),
+      functools.partial(uops_to_cstyle, GLSLLanguage()), lambda x: x, WebGLProgram)

extra/backends/ops_webgpu.py

@@ -0,0 +1,41 @@
+from wgpu.utils.device import get_default_device
+from tinygrad.device import Compiled, Allocator
+from tinygrad.codegen.kernel import LinearizerOptions
+from tinygrad.renderer.cstyle import WGSLRenderer
+import wgpu
+
+wgpu_device = get_default_device()
+def create_uniform(val: int) -> wgpu.GPUBuffer:
+  buf = wgpu_device.create_buffer(size=4, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)
+  wgpu_device.queue.write_buffer(buf, 0, val.to_bytes(4, "little"))
+  return buf
+
+class WebGPUProgram:
+  def __init__(self, name:str, lib:bytes):
+    self.name, self.lib, self.prg = name, lib, wgpu_device.create_shader_module(code=lib)   # NOTE: this is the compiler
+  def __call__(self, *bufs, global_size, local_size, vals=(), wait=False):
+    assert len(bufs) <= 8, "WEBGPU only supports 8 buffers"
+    binding_layouts = [{"binding": i, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.uniform if i >= len(bufs) else wgpu.BufferBindingType.storage }} for i in range(len(bufs)+len(vals))]  # noqa: E501
+    bindings = [{"binding": i, "resource": {"buffer": create_uniform(x) if i >= len(bufs) else x, "offset": 0, "size": 4 if i >= len(bufs) else x.size}} for i,x in enumerate(bufs+vals)]  # noqa: E501
+    bind_group_layout = wgpu_device.create_bind_group_layout(entries=binding_layouts)
+    pipeline_layout = wgpu_device.create_pipeline_layout(bind_group_layouts=[bind_group_layout])
+    bind_group = wgpu_device.create_bind_group(layout=bind_group_layout, entries=bindings)
+    compute_pipeline = wgpu_device.create_compute_pipeline(layout=pipeline_layout,compute={"module": self.prg, "entry_point": self.name},)
+    command_encoder = wgpu_device.create_command_encoder()
+    compute_pass = command_encoder.begin_compute_pass()
+    compute_pass.set_pipeline(compute_pipeline)
+    compute_pass.set_bind_group(0, bind_group, [], 0, 999999) # last 2 not used
+    compute_pass.dispatch_workgroups(*global_size)  # x y z
+    compute_pass.end()
+    wgpu_device.queue.submit([command_encoder.finish()])
+
+class WebGpuAllocator(Allocator):
+  def _alloc(self, size: int):
+    return wgpu_device.create_buffer(size=size, usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_DST | wgpu.BufferUsage.COPY_SRC)
+  def copyin(self, dest, src: memoryview): wgpu_device.queue.write_buffer(dest, 0, src)
+  def copyout(self, dest, src: memoryview): dest[:] = wgpu_device.queue.read_buffer(src, 0)    # TODO: remove this copy
+
+class WebGpuDevice(Compiled):
+  def __init__(self, device:str):
+    super().__init__(WebGpuAllocator(), LinearizerOptions(device="WEBGPU", supports_float4=False, local_max=[256, 256, 64],
+                                                          global_max=[65535, 65535, 65535]), WGSLRenderer, lambda x: x, WebGPUProgram)

extra/backends/triton.py

@@ -0,0 +1,131 @@
+from typing import Dict, List, Final, Callable, DefaultDict
+from collections import defaultdict
+from tinygrad.ops import UnaryOps, BinaryOps, TernaryOps, Op
+from tinygrad.helpers import DType, PtrDType, dtypes, ImageDType, DEBUG, getenv
+from tinygrad.codegen.linearizer import  UOp, UOps
+from triton.compiler import compile as triton_compile
+import linecache
+import math
+import re
+
+triton_dtypes = {dtypes.double: "tl.float64", dtypes.float32: "tl.float32", dtypes.float16: "tl.float16", dtypes.bool: "tl.int1", dtypes.int8: "tl.int8", dtypes.uint8: "tl.uint8", dtypes.int32: "tl.int32", dtypes.int64: "tl.int64", dtypes.uint32: "tl.uint32", dtypes.uint64: "tl.uint64", dtypes.int16: "tl.int16", dtypes.uint16: "tl.uint16"}
+signature_dtypes = {dtypes.double: "fp64",dtypes.float32: "fp32", dtypes.float16: "fp16", dtypes.bool: "i8", dtypes.int8: "i1", dtypes.uint8: "u8", dtypes.int32: "i32", dtypes.int64: "i64", dtypes.uint32: "u32", dtypes.uint64: "u64", dtypes.int16: "i16", dtypes.uint16: "u16"}
+
+def next_power_of_2(x):
+  return 1 << (x - 1).bit_length()
+
+def render_valid(valid):
+  return '(' * (len(valid) -1) + ') and '.join(valid) if len(valid) else 'True'
+
+#NOTE Triton requires matching dimensions for load/store, disable this and see TestOps::test_output_padded_conv_transpose2d fail to compile
+def fill_dims_for_idx(idx, dims):
+  return "(" + idx + "+ (" + (f"0*({'+'.join(d for d in dims)})))") if len(dims) else idx
+
+def get_max(var):
+  if isinstance(var, int): return var
+  return re.sub(r'\[(.*?)\]', '', str(var))[1:-1]
+
+#NOTE can be removed after https://github.com/gpuocelot/gpuocelot/issues/8 gets resolved
+def remove_single_scalar_curly_braces(ptx_code):
+  return '\n'.join([re.sub(r'\{\s*(%\w+)\s*\}', r'\1', line) for line in ptx_code.split('\n')])
+
+def render_const(args,dtype:DType):
+  return (('-' if args<0 else '') + 'tl.where(1,float("inf"),0)') if math.isinf(args) else ('tl.where(1,float("nan"),0)' if math.isnan(args) else f"{int(args)}" if dtypes.is_int(dtype) else str(args))
+
+def render_cast(x:str, dtype:DType, bitcast=False):
+  return f"{x}.to({triton_dtypes[dtype]}, bitcast={bitcast})"
+
+def define_scalar(local_size, dtype, args):
+  if len(local_size) > 0: return f"tl.full(({','.join([str(next_power_of_2(x)) for x in local_size])},),{render_const(args,dtype)}, dtype={triton_dtypes[dtype]})"
+  return render_const(args,dtype)
+
+def uops_to_triton(function_name:str, uops:List[UOp]):
+  local_size: List[int] = []
+  depth = 1
+  signatures, dims, bufs, kernel, valid = [], [], [], [], [] #type: ignore
+
+  c: DefaultDict[str, int] = defaultdict(int)
+  r: Dict[UOp, str] = {}
+  def ssa(u, prefix="t"):
+    nonlocal c, r
+    c[prefix] += 1
+    r[u]=f"{prefix}{c[prefix]-1}"
+    return r[u]
+
+  child_count: DefaultDict[UOp, int] = defaultdict(int)
+  for ru in uops:
+    for v in ru.vin:
+      child_count[v] += 1
+
+  def kk(s): kernel.append("  "*depth+s)
+  code_for_op: Final[Dict[Op, Callable]] = {
+    UnaryOps.EXP2: lambda x,dtype,: f"tl.math.exp2({x})",
+    UnaryOps.LOG2: lambda x,dtype,: f"tl.math.log2({x})",
+    UnaryOps.SIN: lambda x,dtype: f"tl.sin({x})",
+    UnaryOps.SQRT: lambda x,dtype: f"tl.sqrt({x})",
+    UnaryOps.NEG: lambda x,dtype: f"-{x}" if dtype != dtypes.bool else f"tl.where({x}, 0, 1)",
+    BinaryOps.ADD: lambda x,y,dtype: f"({x}+{y})", BinaryOps.SUB: lambda x,y,: f"({x}-{y})",
+    BinaryOps.MUL: lambda x,y,dtype: f"({x}*{y})", BinaryOps.DIV: lambda x,y,: f"({x}/{y})" if y != '0.0' else f"{x}*tl.where({x}==0.0, float('nan'), float('inf'))",
+    BinaryOps.MAX: lambda x,y,dtype: f"tl.maximum({x},{y})",
+    BinaryOps.CMPLT: lambda x,y,dtype: f"({x}<{y})",
+    BinaryOps.MOD: lambda x,y,dtype: f"tl.abs({x})%tl.abs({y})*tl.where({x}<0,-1,1)",
+    TernaryOps.MULACC: lambda x,y,z,dtype: f"(({x}*{y})+{z})",
+    TernaryOps.WHERE: lambda x,y,z,dtype: f"tl.where({x},{y},{z})",
+  }
+  def int_div(x,y): return f"({x}//{y})" if y != '0' else f"{x}*tl.where({x}==0, float('nan'), float('inf'))"
+  for u in uops:
+    uop,dtype,vin,args = u.uop,u.dtype,u.vin,u.arg
+    if uop == UOps.LOOP:
+      kk(f"for {ssa(u, 'ridx')} in range({vin[0].arg}, {r[vin[1]]}):")
+      depth += 1
+    elif uop == UOps.END: depth -= 1
+    elif uop == UOps.ALU:
+      assert dtype is not None
+      val = code_for_op[args](*[r[x] for x in vin])
+      if child_count[u] <=1 or dtypes.is_int(dtype): r[u] = int_div(*[r[x] for x in vin]) if args == BinaryOps.DIV and dtypes.is_int(dtype) else val
+      else: kk(f"{ssa(u, 'alu')} = ({val})")
+    elif uop == UOps.LOAD:
+      assert dtype is not None
+      if len(vin) == 2: kk(f"{ssa(u, 'val')} = {render_cast(f'tl.load({r[vin[0]]} + { fill_dims_for_idx(r[vin[1]], dims)}, mask = {render_valid(valid)})', dtype)}")
+      else: kk(f"{ssa(u, 'val')} = {render_cast(f'tl.where({r[vin[2]]}, tl.load({r[vin[0]]}+{fill_dims_for_idx(r[vin[1]],dims)} , mask={render_valid(valid+[r[vin[2]]])}), 0.0)', dtype)}")
+    elif uop == UOps.DEFINE_ACC: kk(f"{ssa(u, 'acc')} = {define_scalar(local_size, dtype, args).replace('//', '/')}")
+    elif uop == UOps.CONST: r[u] = define_scalar([], dtype, args)
+    elif uop == UOps.PHI:
+      kk(f"{r[vin[0]]} = {r[vin[1]].replace('//', '/')}")
+      r[u] = r[vin[0]]
+    elif uop == UOps.STORE:
+      assert not isinstance(dtype, ImageDType), "unimplemented: image store"
+      kk(f"{'if '+r[vin[3]]+': ' if len(vin)>3 else ''}tl.store({r[vin[0]]} + {r[vin[1]]}, {r[vin[2]].replace('//', '/')}, mask = {render_valid(valid)}) ")
+    elif uop == UOps.DEFINE_GLOBAL:
+      bufs.append(args)
+      signatures.append("*" if isinstance(dtype, PtrDType) else "" +  signature_dtypes[dtype])
+      r[u] = args
+    elif uop == UOps.SPECIAL:
+      dims.append(args[1])
+      valid.append(f"{args[1]}<{get_max(args[2])}")
+      if args[1].startswith("g"): kk(f"{args[1]} = tl.program_id({args[0]}) # {args[2]}")
+      elif args[1].startswith("l"):
+        kk(f"{args[1]} = tl.arange({0}, {next_power_of_2(args[2])})")
+        local_size.append(args[2])
+      r[u] = args[1]
+    elif uop == UOps.CAST and dtype is not None: r[u] = render_cast(r[vin[0]], dtype, isinstance(args, tuple) and args[1])
+    else: raise NotImplementedError(f"unimplemented: {uop}")
+
+  prg = f"import triton\nimport triton.language as tl\ntl.core.TRITON_MAX_TENSOR_NUMEL = float('inf')\n@triton.jit\ndef {function_name}("+','.join(bufs)+"):\n"
+  for i, line in enumerate(list(filter(lambda line: "tl.arange" in line, kernel))): kernel[kernel.index(line)] +=  f"[{', '.join([':' if i == j else 'None' for j in range(len(local_size))])}]"
+  prg += "\n".join(kernel)
+
+  acc_local_size = 1
+  for x in local_size: acc_local_size *= next_power_of_2(x)
+  local_size = [acc_local_size] + [1] * (len(local_size) - 1)
+
+  if DEBUG >= 4: print(prg)
+  getlines = linecache.getlines
+  linecache.getlines = lambda filename, module_globals=None: prg.splitlines(keepends=True) if "<triton>" == filename else getlines(filename, module_globals)
+  exec(compile(prg, "<triton>", "exec"), globals()) # pylint: disable=W0122\
+  compiled = triton_compile(globals()[function_name], signature=",".join(signatures), device_type="cuda", debug=False, cc=(35 if getenv("CUDACPU", 0) else None))
+  prg = remove_single_scalar_curly_braces(compiled.asm["ptx"].split(".file")[0].split(".visible .func")[0])
+  max_local_size =  [int(x) for x in prg.split(".maxntid ")[1].split("\n")[0].split(", ")]
+  for i in range(len(local_size)): local_size[i] = min(local_size[i], max_local_size[i])
+
+  return prg, {"shared":compiled.metadata["shared"], "local_size":local_size + [1]*(3-len(local_size))}