add TestKernelOpts.test_tensor_core_opts run in benchmarks

* moveout pci logic

* fixes

* oops

* types

* more type

* one style

* thi is imp

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

@@ -29,6 +29,10 @@ inputs:
     description: "Install CUDA?"
     required: false
     default: 'false'
+  ocelot:
+    description: "Install gpuocelot?"
+    required: false
+    default: 'false'
   webgpu:
     description: "Install webgpu?"
     required: false
@@ -193,14 +197,14 @@ runs:
           sudo xargs curl -L -o /usr/local/lib/libamd_comgr.dylib
         cargo build --release --manifest-path ./extra/remu/Cargo.toml
 
-    # **** CUDA ****
+    # **** gpuocelot ****
 
     - name: Install gpuocelot dependencies (MacOS)
-      if: inputs.cuda == 'true' && runner.os == 'macOS'
+      if: inputs.ocelot == 'true' && runner.os == 'macOS'
       shell: bash
       run: brew install --quiet cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
     - name: Cache gpuocelot
-      if: inputs.cuda == 'true'
+      if: inputs.ocelot == 'true'
       id: cache-build
       uses: actions/cache@v4
       env:
@@ -209,7 +213,7 @@ runs:
         path: ${{ github.workspace }}/gpuocelot/ocelot
         key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-0
     - name: Clone/compile gpuocelot
-      if: inputs.cuda == 'true' && steps.cache-build.outputs.cache-hit != 'true'
+      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
       shell: bash
       run: |
         git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
@@ -220,7 +224,7 @@ runs:
         cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5
         ninja
     - name: Install gpuocelot
-      if: inputs.cuda == 'true'
+      if: inputs.ocelot == 'true'
       shell: bash
       run: |
         cd ${{ github.workspace }}/gpuocelot/ocelot/build

.github/workflows/benchmark.yml

@@ -67,7 +67,7 @@ jobs:
     - name: Test speed vs torch
       run: BIG=2 MPS=1 python3.11 test/test_speed_v_torch.py | tee torch_speed.txt
     - name: Test tensor cores
-      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
     - name: Test AMX tensor cores
       run: |
         DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
@@ -196,8 +196,8 @@ jobs:
       run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
     - name: Test tensor cores
       run: |
-        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
+        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
     - name: Run Tensor Core GEMM (CUDA)
       run: |
         CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
@@ -396,8 +396,8 @@ jobs:
       run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Test tensor cores
       run: |
-        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
-        AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
+        AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
         AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
     - name: Run Tensor Core GEMM (AMD)
       run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
@@ -607,12 +607,8 @@ jobs:
       run: test/external/process_replay/reset.py
     - name: validate openpilot 0.9.7
       run: PYTHONPATH=. FLOAT16=0 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
-    - name: benchmark openpilot 0.9.4
-      run: BENCHMARK_LOG=openpilot_0_9_4 PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_4.txt
     - name: benchmark openpilot 0.9.7
       run: BENCHMARK_LOG=openpilot_0_9_7 PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_7.txt
-    - name: benchmark openpilot w IMAGE=2 0.9.4
-      run: BENCHMARK_LOG=openpilot_0_9_4_image PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_4.txt
     - name: benchmark openpilot w IMAGE=2 0.9.7
       run: BENCHMARK_LOG=openpilot_0_9_7_image PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
     - name: openpilot compile3 0.9.7

.github/workflows/test.yml

@@ -450,6 +450,8 @@ jobs:
         run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
       - name: Test openpilot LLVM compile
         run: PYTHONPATH="." LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      - name: Test openpilot compile4
+        run: PYTHONPATH="." NOLOCALS=1 GPU=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -677,6 +679,7 @@ jobs:
           key: ${{ matrix.backend }}-minimal
           deps: testing_minimal
           cuda: 'true'
+          ocelot: 'true'
       - name: Set env
         run: printf "${{ matrix.backend == 'PTX' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nPTX=1\nMOCKGPU=1' || matrix.backend == 'nv' && 'NV=1\nMOCKGPU=1\nFORWARD_ONLY=1' }}" >> $GITHUB_ENV
       - name: Check Device.DEFAULT and print some source
@@ -744,6 +747,7 @@ jobs:
         python-version: '3.11'
         amd: 'true'
         cuda: 'true'
+        ocelot: 'true'
         llvm: 'true'
     - name: Run real world test
       run: METAL=1 python -m pytest -n=auto test/models/test_real_world.py --durations=20
@@ -812,15 +816,16 @@ jobs:
       run: npm install puppeteer
     - name: Run WEBGPU Efficientnet
       run: node test/web/test_webgpu.js
-    - name: Run VIZ tests as external package
-      run: |
-        mkdir $GITHUB_WORKSPACE/test_dir
-        cd $GITHUB_WORKSPACE/test_dir
-        python -m venv venv
-        source venv/bin/activate
-        pip install $GITHUB_WORKSPACE
-        cp $GITHUB_WORKSPACE/test/web/test_viz.js .
-        node test_viz.js
+    # this is flaky
+    #- name: Run VIZ tests as external package
+    #  run: |
+    #    mkdir $GITHUB_WORKSPACE/test_dir
+    #    cd $GITHUB_WORKSPACE/test_dir
+    #    python -m venv venv
+    #    source venv/bin/activate
+    #    pip install $GITHUB_WORKSPACE
+    #    cp $GITHUB_WORKSPACE/test/web/test_viz.js .
+    #    node test_viz.js
 
   osxremote:
     name: MacOS (remote metal)
@@ -852,9 +857,8 @@ jobs:
     timeout-minutes: 20
     env:
       REMOTE: 1
-      REMOTEDEV: 'AMD'
+      HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
       PYTHONPATH: ${{ github.workspace }}
-      MOCKGPU: 1
     steps:
       - name: Checkout Code
         uses: actions/checkout@v4
@@ -865,6 +869,21 @@ jobs:
           deps: testing_minimal
           amd: 'true'
           llvm: 'true'
+      - name: Start remote server
+        run: |
+          start_server() {
+            systemd-run --user \
+              --unit="$1" \
+              --setenv=REMOTEDEV=AMD \
+              --setenv=MOCKGPU=1 \
+              --setenv=PYTHONPATH=. \
+              --setenv=PORT="$2" \
+              --working-directory="$(pwd)" \
+              python tinygrad/runtime/ops_remote.py
+          }
+
+          start_server "remote-server-1" 6667
+          start_server "remote-server-2" 6668
       - name: Check Device.DEFAULT and print some source
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
@@ -872,7 +891,12 @@ jobs:
           DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run REMOTE=1 Test
         run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py
+      - name: Show remote server logs
+        if: always()
+        run: |
+          journalctl --user -u remote-server-1 --no-pager
+          journalctl --user -u remote-server-2 --no-pager
 
   osxtests:
     strategy:

autogen_stubs.sh

@@ -390,8 +390,8 @@ generate_am() {
     $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu14_driver_if_v14_0.h \
     extra/amdpci/headers/amdgpu_smu.h \
     --clang-args="-include stdint.h" \
-    -o $BASE/am/smu_v14_0_3.py
-  fixup $BASE/am/smu_v14_0_3.py
+    -o $BASE/am/smu_v14_0_2.py
+  fixup $BASE/am/smu_v14_0_2.py
 }
 
 generate_sqtt() {

docs/abstractions2.py

@@ -59,11 +59,11 @@ st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.view(ShapeTracker.from_shape((1,)
 s = UOp(Ops.SINK, dtypes.void, (st_0,))
 
 # convert the computation to a "linearized" format (print the format)
-from tinygrad.engine.realize import get_kernel, CompiledRunner
-kernel = get_kernel(Device[DEVICE].renderer, s).linearize()
+from tinygrad.engine.realize import get_program, CompiledRunner
+program = get_program(Device[DEVICE].renderer, s)
 
 # compile a program (and print the source)
-fxn = CompiledRunner(kernel.to_program())
+fxn = CompiledRunner(program)
 print(fxn.p.src)
 # NOTE: fxn.clprg is the CPUProgram
 
@@ -78,7 +78,7 @@ print("******** third, the UOp ***********")
 
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.engine.grouper import get_kernelize_map
+from tinygrad.engine.kernelize import get_kernelize_map
 
 # allocate some values + load in values
 a = UOp.new_buffer(DEVICE, 1, dtypes.int32)

docs/ramp.py

@@ -239,7 +239,7 @@ print("******* PART 3 *******")
 # it's much simpler than what's in LLVM or MLIR
 
 from tinygrad import dtypes
-from tinygrad.uop import UOp, Ops
+from tinygrad.uop.ops import UOp, Ops
 
 # first, we'll construct some const UOps
 a = UOp(Ops.CONST, dtypes.int, arg=2)

examples/handcode_opt.py

@@ -2,11 +2,11 @@ from extra.models.resnet import ResNet50
 from extra.mcts_search import mcts_search
 from examples.mlperf.helpers import get_mlperf_bert_model
 from tinygrad import Tensor, Device, dtypes, nn
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.uop.ops import Ops, sym_infer
 from tinygrad.device import Compiled
-from tinygrad.engine.search import beam_search, bufs_from_lin
+from tinygrad.opt.search import beam_search, bufs_from_lin
 from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
 from extra.optimization.helpers import time_linearizer
 

examples/llama3.py

@@ -157,7 +157,11 @@ MODEL_PARAMS = {
   "70B": {
     "args": {"dim": 8192, "n_heads": 64, "n_kv_heads": 8, "n_layers": 80, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256,  "hidden_dim": 28672},
     "files": 8
-  }
+  },
+  "405B": {
+    "args": {"dim": 16384, "n_heads": 128, "n_kv_heads": 8, "n_layers": 126, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256,  "hidden_dim": 53248},
+    "files": 191
+  },
 }
 def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dtype=dtypes.float16, device=None, max_context=8192, load_weights=True):
   # build model
@@ -236,7 +240,7 @@ if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("--download_model", action="store_true", help="Download a model")
   parser.add_argument("--model", type=Path, help="Model path")
-  parser.add_argument("--size", choices=["1B", "8B", "70B"], default="1B", help="Model size")
+  parser.add_argument("--size", choices=["1B", "8B", "70B", "405B"], default="1B", help="Model size")
   parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
   parser.add_argument("--quantize", choices=["int8", "nf4", "float16"], help="Quantization method")
   parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")

examples/mlperf/model_train.py

@@ -914,16 +914,24 @@ def train_rnnt():
   pass
 
 @TinyJit
-def train_step_bert(model, optimizer, scheduler, loss_scaler:float, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
-                    masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
-  for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
-    if len(GPUS) > 1: t.shard_(GPUS, axis=0)
-    else: t.to_(GPUS[0])
+def train_step_bert(model, optimizer, scheduler, loss_scaler:float, GPUS, grad_acc:int, **kwargs):
   optimizer.zero_grad()
 
-  lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
-  loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
-  (loss * loss_scaler).backward()
+  for i in range(grad_acc):
+    input_ids, segment_ids = kwargs[f"input_ids{i}"], kwargs[f"segment_ids{i}"]
+    # NOTE: these two have different names
+    attention_mask, masked_positions = kwargs[f"input_mask{i}"], kwargs[f"masked_lm_positions{i}"]
+    masked_lm_ids, masked_lm_weights, next_sentence_labels = kwargs[f"masked_lm_ids{i}"], kwargs[f"masked_lm_weights{i}"], kwargs[f"next_sentence_labels{i}"]
+
+    for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
+      if len(GPUS) > 1: t.shard_(GPUS, axis=0)
+      else: t.to_(GPUS[0])
+
+    lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
+    loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
+    (loss * loss_scaler).backward()
+    # TODO: OOM without this realize with large grad_acc
+    Tensor.realize(*[p.grad for p in optimizer.params])
 
   global_norm = Tensor([0.0], dtype=dtypes.float32, device=optimizer[0].device)
   for p in optimizer.params:
@@ -999,16 +1007,19 @@ def train_bert():
     MLLOGGER = None
 
   # ** hyperparameters **
-  BS                 = config["GLOBAL_BATCH_SIZE"]      = getenv("BS", 11 * len(GPUS) if dtypes.default_float in (dtypes.float16, dtypes.bfloat16) else 8 * len(GPUS))
+  BS                 = config["BS"]                     = getenv("BS", 11 * len(GPUS) if dtypes.default_float in (dtypes.float16, dtypes.bfloat16) else 8 * len(GPUS))
+  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
+  # TODO: mlperf logging
+  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
   EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 1 * len(GPUS))
-  max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.000175 * math.sqrt(BS/96))
+  max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.000175 * math.sqrt(GBS/96))
   opt_lamb_beta_1    = config["OPT_LAMB_BETA_1"]        = getenv("OPT_LAMB_BETA_1", 0.9)
   opt_lamb_beta_2    = config["OPT_LAMB_BETA_2"]        = getenv("OPT_LAMB_BETA_2", 0.999)
 
-  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3600000 // BS)
+  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3600000 // GBS)
   warmup_steps       = config["NUM_WARMUP_STEPS"]       = getenv("NUM_WARMUP_STEPS", 1)
   max_eval_steps     = config["MAX_EVAL_STEPS"]         = getenv("MAX_EVAL_STEPS", (10000 + EVAL_BS - 1) // EVAL_BS) # EVAL_BS * MAX_EVAL_STEPS >= 10000
-  eval_step_freq     = config["EVAL_STEP_FREQ"]         = getenv("EVAL_STEP_FREQ", int((math.floor(0.05 * (230.23 * BS + 3000000) / 25000) * 25000) / BS)) # Round down
+  eval_step_freq     = config["EVAL_STEP_FREQ"]         = getenv("EVAL_STEP_FREQ", int((math.floor(0.05 * (230.23 * GBS + 3000000) / 25000) * 25000) / GBS)) # Round down
   save_ckpt_freq     = config["SAVE_CKPT_FREQ"]         = getenv("SAVE_CKPT_FREQ", 1000)
   keep_ckpt_amount   = config["KEEP_CKPT_AMOUNT"]       = getenv("KEEP_CKPT_AMOUNT", 5)
   save_ckpt_dir      = config["SAVE_CKPT_DIR"]          = getenv("SAVE_CKPT_DIR", "./ckpts")
@@ -1066,7 +1077,7 @@ def train_bert():
   scheduler_wd = PolynomialDecayWithWarmup(optimizer_wd, max_lr, 0, train_steps, warmup_steps, power=poly_power)
   scheduler_no_wd = PolynomialDecayWithWarmup(optimizer_no_wd, max_lr, 0, train_steps, warmup_steps, power=poly_power)
   scheduler_group = LRSchedulerGroup(scheduler_wd, scheduler_no_wd)
-  print(f"training with batch size {BS} for one epoch with {train_steps} steps")
+  print(f"training with global batch size {GBS} for one epoch with {train_steps} steps")
 
   # log mlperf hparams
   if MLLOGGER:
@@ -1115,11 +1126,11 @@ def train_bert():
   # ** train loop **
   wc_start = time.perf_counter()
 
-  i, train_data = start_step, next(train_it)
+  i, train_data = start_step, [next(train_it) for _ in range(grad_acc)]
 
   if RUNMLPERF:
     if MLLOGGER:
-      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*BS, metadata={"epoch_num": i*BS})
+      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*GBS, metadata={"epoch_num": i*GBS})
 
   while train_data is not None and i < train_steps and not achieved:
     if getenv("TRAIN", 1):
@@ -1128,14 +1139,13 @@ def train_bert():
       st = time.perf_counter()
       GlobalCounters.reset()
       with WallTimeEvent(BenchEvent.STEP):
-        loss, global_norm, lr = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler,
-          train_data["input_ids"], train_data["segment_ids"], train_data["input_mask"], train_data["masked_lm_positions"], \
-          train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"], GPUS)
+        data = {f"{k}{i}":v for i,d in enumerate(train_data) for k,v in d.items()}
+        loss, global_norm, lr = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler, GPUS, grad_acc, **data)
 
         pt = time.perf_counter()
 
         try:
-          next_data = next(train_it)
+          next_data = [next(train_it) for _ in range(grad_acc)]
         except StopIteration:
           next_data = None
 
@@ -1156,7 +1166,7 @@ def train_bert():
       if WANDB:
         wandb.log({"lr": lr, "train/loss": loss, "train/global_norm": global_norm.item(), "train/step_time": cl - st,
                     "train/python_time": pt - st, "train/data_time": dt - pt, "train/cl_time": cl - dt,
-                    "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*BS})
+                    "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})
 
       train_data, next_data = next_data, None
       i += 1
@@ -1171,7 +1181,7 @@ def train_bert():
     # ** eval loop **
     if i % eval_step_freq == 0 or (BENCHMARK and i == BENCHMARK) or i == train_steps:
       if MLLOGGER and RUNMLPERF:
-        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*BS, "step_num": i})
+        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i})
       if getenv("RESET_STEP"): train_step_bert.reset()
       elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
       eval_lm_losses = []
@@ -1221,11 +1231,11 @@ def train_bert():
 
       if WANDB:
         wandb.log({"eval/lm_loss": avg_lm_loss, "eval/clsf_loss": avg_clsf_loss, "eval/lm_accuracy": avg_lm_acc, \
-                    "eval/clsf_accuracy": avg_clsf_acc, "eval/forward_time": avg_fw_time, "epoch": (i+1)*BS})
+                    "eval/clsf_accuracy": avg_clsf_acc, "eval/forward_time": avg_fw_time, "epoch": (i+1)*GBS})
 
       if MLLOGGER and RUNMLPERF:
-        MLLOGGER.end(key=mllog_constants.EVAL_STOP, value=i*BS, metadata={"epoch_count": i*BS, "step_num": i, "samples_count": config["EVAL_BS"] * config["MAX_EVAL_STEPS"]})
-        MLLOGGER.event(key=mllog_constants.EVAL_ACCURACY, value=avg_lm_acc, metadata={"epoch_num": i*BS, "masked_lm_accuracy": avg_lm_acc})
+        MLLOGGER.end(key=mllog_constants.EVAL_STOP, value=i*GBS, metadata={"epoch_count": i*GBS, "step_num": i, "samples_count": config["EVAL_BS"] * config["MAX_EVAL_STEPS"]})
+        MLLOGGER.event(key=mllog_constants.EVAL_ACCURACY, value=avg_lm_acc, metadata={"epoch_num": i*GBS, "masked_lm_accuracy": avg_lm_acc})
 
       # save model if achieved target
       if not achieved and avg_lm_acc >= target:
@@ -1240,10 +1250,10 @@ def train_bert():
         hours = int(total_seconds // 3600)
         minutes = int((total_seconds % 3600) // 60)
         seconds = total_seconds % 60
-        print(f"Reference Convergence point reached after {i * BS} datasamples and {hours}h{minutes}m{seconds:.2f}s.")
+        print(f"Reference Convergence point reached after {i * GBS} datasamples and {hours}h{minutes}m{seconds:.2f}s.")
         achieved = True
         if MLLOGGER and RUNMLPERF:
-          MLLOGGER.event(key=mllog_constants.EPOCH_STOP, value=i*BS, metadata={"epoch_num": i*BS})
+          MLLOGGER.event(key=mllog_constants.EPOCH_STOP, value=i*GBS, metadata={"epoch_num": i*GBS})
           MLLOGGER.end(key=mllog_constants.RUN_STOP, metadata=dict(status=mllog_constants.SUCCESS))
         # stop once hitting the target
         break
@@ -1271,13 +1281,9 @@ def train_bert():
         os.remove(os.path.join(ckpt_dir, last))
       if MLLOGGER and RUNMLPERF:
         MLLOGGER.end(key="checkpoint_stop", value=None, metadata={"step_num": i})
-        MLLOGGER.start(key=mllog_constants.BLOCK_START, value=None, metadata={"first_epoch_num": 1, "epoch_num": 1, "epoch_count": 1, "samples_count": i * BS, "step_num": i, "first_step_num": i+1})
+        MLLOGGER.start(key=mllog_constants.BLOCK_START, value=None, metadata={"first_epoch_num": 1, "epoch_num": 1, "epoch_count": 1, "samples_count": i * GBS, "step_num": i, "first_step_num": i+1})
         previous_step = i
 
-def train_maskrcnn():
-  # TODO: Mask RCNN
-  pass
-
 if __name__ == "__main__":
   multiprocessing.set_start_method('spawn')
 

examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_beam.sh

@@ -9,6 +9,6 @@ export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MI
 export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
 export BASEDIR="/raid/datasets/wiki"
 
-export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+export BENCHMARK=10 BERT_LAYERS=2
 
 python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh

@@ -22,8 +22,7 @@ export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
 LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log"
 
-# init  # TODO: without DEBUG=2 it hangs
-BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 DEBUG=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
 
 # run
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE

examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh

@@ -27,6 +27,4 @@ sleep 5 && sudo rmmod amdgpu || true
 BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
 
 # run
-# TODO: AM driver resulted in nan
-sudo modprobe amdgpu
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE

examples/openpilot/compile4.py

@@ -1,8 +1,8 @@
 import sys, onnx
 from tinygrad import Tensor, fetch, GlobalCounters
-from tinygrad.uop import UOp
+from tinygrad.uop.ops import UOp
 from tinygrad.frontend.onnx import OnnxRunner
-from tinygrad.engine.grouper import get_kernelize_map
+from tinygrad.engine.kernelize import get_kernelize_map
 from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import run_schedule
 
@@ -37,12 +37,12 @@ if __name__ == "__main__":
   independent = UOp.sink(*independent_set.keys())
   kernelized = get_kernelize_map(independent)
   independent = independent.substitute(kernelized)
-  schedule, var_vals, becomes_map = create_schedule_with_vars(independent)
+  schedule, var_vals = create_schedule_with_vars(independent)
   run_schedule(schedule)
 
   print("**** real ****")
   GlobalCounters.reset()
-  out.uop = root.substitute(kernelized).substitute(becomes_map)
+  out.uop = root.substitute(kernelized)
   out.kernelize()
 
   # realize

examples/so_vits_svc.py

@@ -5,7 +5,7 @@ from functools import partial, reduce
 from pathlib import Path
 from typing import Tuple, Optional, Type
 from tinygrad import nn, dtypes, Tensor
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, fetch
 from tinygrad.nn.state import torch_load
 from examples.vits import ResidualCouplingBlock, PosteriorEncoder, Encoder, ResBlock1, ResBlock2, LRELU_SLOPE, sequence_mask, split, get_hparams_from_file, load_checkpoint, weight_norm, HParams
 from examples.sovits_helpers import preprocess
@@ -19,10 +19,6 @@ F0_MIN = 50.0
 F0_MEL_MIN = 1127 * np.log(1 + F0_MIN / 700)
 F0_MEL_MAX = 1127 * np.log(1 + F0_MAX / 700)
 
-def download_if_not_present(file_path: Path, url: str):
-  if not os.path.isfile(file_path): download_file(url, file_path)
-  return file_path
-
 class SpeechEncoder:
   def __init__(self, hidden_dim, model:ContentVec): self.hidden_dim, self.model = hidden_dim, model
   def encode(self, ): raise NotImplementedError("implement me")
@@ -97,7 +93,7 @@ class ContentVec:
     return res, padding_mask
   @classmethod
   def load_from_pretrained(cls, checkpoint_path:str, checkpoint_url:str) -> ContentVec:
-    download_if_not_present(checkpoint_path, checkpoint_url)
+    fetch(checkpoint_url, checkpoint_path)
     cfg = load_fairseq_cfg(checkpoint_path)
     enc = cls(cfg.model)
     _ = load_checkpoint_enc(checkpoint_path, enc, None)
@@ -324,9 +320,9 @@ class Synthesizer:
     return f0_coarse
   @classmethod
   def load_from_pretrained(cls, config_path:str, config_url:str, weights_path:str, weights_url:str) -> Synthesizer:
-    download_if_not_present(config_path, config_url)
+    fetch(config_url, config_path)
     hps = get_hparams_from_file(config_path)
-    download_if_not_present(weights_path, weights_url)
+    fetch(weights_url, weights_path)
     net_g = cls(hps.data.filter_length // 2 + 1, hps.train.segment_size // hps.data.hop_length, **hps.model)
     _ = load_checkpoint(weights_path, net_g, None, skip_list=["f0_decoder"])
     logging.debug(f"{cls.__name__}:Loaded model with hps: {hps}")
@@ -602,7 +598,7 @@ if __name__=="__main__":
   speaker = args.speaker if args.speaker is not None else list(hps.spk.__dict__.keys())[0]
 
   ### Loading audio and slicing ###
-  if audio_path == DEMO_PATH: download_if_not_present(DEMO_PATH, DEMO_URL)
+  if audio_path == DEMO_PATH: fetch(DEMO_URL, DEMO_PATH)
   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)

extra/assembly/assembly.py

@@ -1,5 +1,5 @@
 from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
-from tinygrad.codegen.kernel import Ops, MemOp, UOp
+from tinygrad.opt.kernel import Ops, MemOp, UOp
 from tinygrad.uop.ops import BinaryOps, UnaryOps
 from tinygrad.dtype import DType, dtypes
 from tinygrad.helpers import DEBUG

extra/assembly/assembly_arm64.py

@@ -3,7 +3,7 @@ from platform import system
 from typing import Tuple, Dict, List, Optional
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.codegen.kernel import Ops, UOp
+from tinygrad.opt.kernel import Ops, UOp
 from tinygrad.helpers import CI
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
 

extra/assembly/assembly_ptx.py

@@ -1,7 +1,7 @@
 from typing import List
 import struct
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
-from tinygrad.codegen.kernel import Ops, UOp
+from tinygrad.opt.kernel import Ops, UOp
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.runtime.ops_cuda import arch

extra/assembly/assembly_rdna.py

@@ -2,7 +2,7 @@ import yaml
 from typing import Tuple, Set, Dict
 from tinygrad import dtypes
 from tinygrad.codegen.assembly import AssemblyCodegen, Register
-from tinygrad.codegen.kernel import Ops
+from tinygrad.opt.kernel import Ops
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH
 

extra/backends/triton.py

@@ -2,7 +2,7 @@ from typing import Dict, List, Final, Callable, DefaultDict
 from collections import defaultdict
 from tinygrad.uop.ops import UnaryOps, BinaryOps, TernaryOps, Op
 from tinygrad.helpers import DType, PtrDType, dtypes, ImageDType, DEBUG, getenv
-from tinygrad.codegen.kernel import  UOp, Ops
+from tinygrad.opt.kernel import  UOp, Ops
 from triton.compiler import compile as triton_compile
 import linecache
 import math

extra/gemm/amd_matmul.py

@@ -4,7 +4,7 @@ import numpy as np
 from dataclasses import replace
 from tinygrad import Tensor, Device, Context
 from tinygrad.helpers import getenv
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
 

extra/gemm/max_matmul.py

@@ -4,9 +4,9 @@ from tinygrad import dtypes
 from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
 
 # for copied uops
-from tinygrad.codegen.kernel import Kernel, KernelOptError
+from tinygrad.opt.kernel import Kernel, KernelOptError
 from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.dtype import PtrDType, DType, DTYPES_DICT
 from tinygrad.shape.shapetracker import ShapeTracker

extra/gemm/simple_matmul.py

@@ -2,7 +2,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
-from tinygrad.codegen.kernel import OptOps
+from tinygrad.opt.kernel import OptOps
 from tinygrad.engine.realize import lower_schedule
 
 dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float

extra/gemm/tinygrad_nv_matmul.py

@@ -1,6 +1,6 @@
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.helpers import getenv, DEBUG
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem
 from dataclasses import replace
 

extra/gemm/tvm_gemm.py

@@ -37,7 +37,7 @@ B = Tensor.rand(K, N, device="CPU")
 C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2)
 
 sched = C.schedule()
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.device import CompilerOptions
 lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
 lin.linearize()

extra/mcts_search.py

@@ -4,9 +4,9 @@ import numpy as np
 np.set_printoptions(suppress=True)
 import math, functools, time, random, statistics
 from tinygrad.helpers import DEBUG, getenv, CACHELEVEL, diskcache_get, diskcache_put, colored, Profiling
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.device import Buffer, Device, CompileError
-from tinygrad.engine.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
+from tinygrad.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
 
 class MCTSNode:
   def __init__(self, kernel:Kernel, parent=None):

extra/onnx.py

@@ -1,6 +1,6 @@
 from types import SimpleNamespace
 from typing import Any, Sequence, cast, Literal, Callable
-import dataclasses, functools, io, math, types
+import dataclasses, functools, io, math, types, warnings
 from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr
 from tinygrad.helpers import getenv, DEBUG, all_same, prod, flatten, make_tuple, argsort
 from tinygrad.dtype import DType, ConstType, dtypes, ImageDType
@@ -14,7 +14,7 @@ def has_field(onnx_type: TypeProto|SimpleNamespace, field):
   if isinstance(onnx_type, TypeProto): return onnx_type.HasField(field)
   return hasattr(onnx_type, field)
 
-def dtype_parse(onnx_dtype: int) -> DType:
+def dtype_parse(onnx_dtype: int, fallback_context: str | None = None) -> DType:
   supported: dict[int, DType] = {
     TensorProto.FLOAT:dtypes.float32, TensorProto.UINT8:dtypes.uint8, TensorProto.INT8:dtypes.int8,
     TensorProto.UINT16:dtypes.uint16, TensorProto.INT16:dtypes.int16, TensorProto.INT32:dtypes.int32, TensorProto.INT64:dtypes.int64,
@@ -26,7 +26,13 @@ def dtype_parse(onnx_dtype: int) -> DType:
     TensorProto.FLOAT8E5M2, TensorProto.FLOAT8E5M2FNUZ, TensorProto.UINT4, TensorProto.INT4
   }
   if onnx_dtype in unsupported: raise NotImplementedError(f"onnx dtype {TensorProto.DataType.Name(onnx_dtype)} is not supported")
-  return supported[onnx_dtype] if is_dtype_supported(supported[onnx_dtype]) else dtypes.float
+  if is_dtype_supported(dtype := supported[onnx_dtype]): return dtype
+  # if fallback_context is provided, we can fall back to a default dtype
+  if fallback_context is not None:
+    default_dtype = dtypes.float
+    warnings.warn(f"dtype {dtype} on {Device.DEFAULT} from {fallback_context} is not supported, falling back to {default_dtype}")
+    return default_dtype
+  raise RuntimeError(f"dtype {dtype} on device {Device.DEFAULT} is not supported")
 
 def attribute_parse(onnx_attribute: AttributeProto):
   supported: dict[AttributeProto.AttributeType, Callable[[AttributeProto], Any]] = {
@@ -46,7 +52,7 @@ def attribute_parse(onnx_attribute: AttributeProto):
 
 def buffer_parse(onnx_tensor: TensorProto) -> Tensor:
   if onnx_tensor.string_data: raise NotImplementedError("Parsing for buffer with string data is not implemented.")
-  dtype, shape = dtype_parse(onnx_tensor.data_type), tuple(onnx_tensor.dims)
+  dtype, shape = dtype_parse(onnx_tensor.data_type, "buffer parse"), tuple(onnx_tensor.dims)
   data = None
   if len(onnx_tensor.float_data): data = onnx_tensor.float_data
   elif len(onnx_tensor.int32_data): data = onnx_tensor.int32_data
@@ -57,12 +63,14 @@ def buffer_parse(onnx_tensor: TensorProto) -> Tensor:
     if len(data) == 1: return Tensor(data.tolist()[0], dtype=dtype).reshape(shape)
     return data.cast(dtype).reshape(shape).to(Device.DEFAULT)
   if has_field(onnx_tensor, "raw_data"):
+    raw_data = onnx_tensor.raw_data
+    if not isinstance(raw_data, Tensor): raw_data = Tensor(raw_data)
     if onnx_tensor.data_type == TensorProto.FLOAT16:
-      np_buffer = np.frombuffer(onnx_tensor.raw_data.data().tobytes(),
+      np_buffer = np.frombuffer(raw_data.data().tobytes(),
                                 dtype=helper.tensor_dtype_to_np_dtype(onnx_tensor.data_type)).copy().reshape(shape)
       if np_buffer.size == 1: return Tensor(np_buffer.item(), dtype=dtype).reshape(shape)
       return Tensor(np_buffer, dtype=dtype)
-    ret = onnx_tensor.raw_data.bitcast(dtype).reshape(shape).to(Device.DEFAULT)
+    ret = raw_data.bitcast(dtype).reshape(shape).to(Device.DEFAULT)
     if shape == (): ret = Tensor(ret.item(), dtype=dtype).reshape(shape)
     return ret
   return Tensor(None)
@@ -76,7 +84,7 @@ def type_parse(onnx_type: TypeProto):
   if has_field(elem_type, "tensor_type"):
     shape = tuple(getattr(d, "dim_param", None) or getattr(d, "dim_value") for d in elem_type.tensor_type.shape.dim) \
       if has_field(elem_type.tensor_type, "shape") else None # test_identity_sequence_cpu
-    dtype = dtype_parse(elem_type.tensor_type.elem_type)
+    dtype = dtype_parse(elem_type.tensor_type.elem_type, "input type spec parse")
     return OnnxValue(shape, dtype, is_optional, is_sequence)
   raise RuntimeError(f"TypeProto was not parsed properly: {onnx_type=}")
 
@@ -145,15 +153,15 @@ class OnnxRunner:
     if spec.is_optional and value is None: return None
     # TODO: need true float16 for dtype checking
     if spec.is_sequence:
-      if not isinstance(value, Sequence): raise RuntimeError(f"{name} received {value}, expected a sequence type")
+      if not isinstance(value, Sequence): raise RuntimeError(f"input {name} received {value}, expected a sequence type")
       sequence = [Tensor(v, dtype=spec.dtype, requires_grad=self.is_training) if not isinstance(v, Tensor) else v for v in value]
-      if not all_same(tuple(t.shape for t in sequence)): raise RuntimeError(f"Shapes for {name} sequence must be homogeneous")
+      if not all_same(tuple(t.shape for t in sequence)): raise RuntimeError(f"Shapes for input {name} sequence must be homogeneous")
       return sequence
     tensor = Tensor(value, dtype=spec.dtype, requires_grad=self.is_training) if not isinstance(value, Tensor) else value
     for dim, (onnx_dim, user_dim_input) in enumerate(zip(spec.shape, tensor.shape, strict=True)):
       if isinstance(onnx_dim, str):
         onnx_dim = self.variable_dims[onnx_dim] if onnx_dim in self.variable_dims else self.variable_dims.setdefault(onnx_dim, int(user_dim_input))
-      if user_dim_input != onnx_dim: raise RuntimeError(f"{name} has mismatch on {dim=}. Expected {onnx_dim}, received {user_dim_input}.")
+      if user_dim_input != onnx_dim: raise RuntimeError(f"input {name} has mismatch on {dim=}. Expected {onnx_dim}, received {user_dim_input}.")
     return tensor
 
   def _dispatch_op(self, op, inps, opts):
@@ -284,7 +292,7 @@ def get_onnx_ops():
     raise ValueError(f"pixel_format={pixel_format!r} is not supported.")
 
   def EyeLike(x:Tensor, dtype:int|None=None, k:int=0):
-    ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_parse(dtype) if dtype is not None else x.dtype)
+    ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_parse(dtype, "EyeLike op") if dtype is not None else x.dtype)
     return ret if x.size(0) == x.size(1) else ret.pad(tuple(None if d == ret.size(0) else (k, d-ret.shape[0]-k) for d in x.shape))
 
   def OptionalHasElement(x:Tensor|None=None): return Tensor(x is not None and x.numel() > 0)
@@ -338,7 +346,7 @@ def get_onnx_ops():
 
   # ***** Casting Ops *****
   # TODO: saturate
-  def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(dtype_parse(to))
+  def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(dtype_parse(to, "Cast op"))
   def CastLike(x:Tensor, target_type:Tensor, saturate:int=1): return x.cast(target_type.dtype)
 
   # ***** Reduce Ops *****
@@ -605,9 +613,13 @@ def get_onnx_ops():
 
   # Reimplemented here because you need legacy RNG for passing ONNX tests.
   def Dropout_7(data:Tensor, ratio:float=0.5, training_mode:bool=False, seed:int|None=None):
-    if not training_mode: return data, Tensor.ones(data.shape, dtype=dtypes.bool)  # if mask is requested as output it will contain all True's.
-    mask = Tensor(np.random.RandomState(seed).random(cast(tuple[int,...], data.shape)) >= ratio, requires_grad=False, device=data.device)
-    return data * mask * (1/(1.0 - ratio)), mask
+    if not training_mode: return data, data.full_like(True, dtype=dtypes.bool)
+    if seed is not None:
+      rand = Tensor(np.random.RandomState(seed).random(cast(tuple[int,...], data.shape)), requires_grad=False, dtype=data.dtype, device=data.device)
+    else:
+      rand = data.rand_like(requires_grad=False)
+    mask = rand >= ratio
+    return data * mask / (1.0 - ratio), mask
   # 6 with 'is_test' needed for https://github.com/MTlab/onnx2caffe/raw/refs/heads/master/model/MobileNetV2.onnx
   def Dropout_6(data:Tensor, ratio:float=0.5, is_test=0): return Dropout_7(data, ratio, training_mode=not is_test)
   Dropout = {6:Dropout_6, 7:Dropout_7}
@@ -731,7 +743,9 @@ def get_onnx_ops():
 
   # ***** Quantization Ops *****
   def QuantizeLinear(x:Tensor, y_scale:Tensor, y_zero_point:Tensor|int=0, axis:int=1, block_size:int=0, output_dtype:int=0, saturate=1):
-    out_dtype = y_zero_point.dtype if isinstance(y_zero_point, Tensor) else dtype_parse(output_dtype) if output_dtype else dtypes.uint8
+    if isinstance(y_zero_point, Tensor): out_dtype = y_zero_point.dtype
+    elif output_dtype != 0: out_dtype = dtype_parse(output_dtype, "QuantizeLinear op")
+    else: out_dtype = dtypes.uint8
     y_scale, y_zero_point = _prepare_quantize(x, y_scale, y_zero_point, axis, block_size)
     if out_dtype == dtypes.uchar:
       # this appears to work in practice, at least for uchar out_dtype. it folds with the quantize stuff

extra/optimization/extract_policynet.py

@@ -5,9 +5,9 @@ from tinygrad.nn import Linear
 from tinygrad.tensor import Tensor
 from tinygrad.nn.optim import Adam
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.engine.search import actions
+from tinygrad.opt.search import actions
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.helpers import getenv
 
 # stuff needed to unpack a kernel
@@ -17,7 +17,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.uop.ops import Variable
 inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 
 INNER = 256
 class PolicyNet:

extra/optimization/extract_sa_pairs.py

@@ -10,11 +10,11 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.uop.ops import Variable
 inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 
 # more stuff
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.engine.search import actions
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.search import actions
 from extra.optimization.helpers import lin_to_feats
 from extra.optimization.pretrain_valuenet import ValueNet
 from tinygrad.nn.optim import Adam

extra/optimization/get_action_space.py

@@ -1,8 +1,8 @@
 import random
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.engine.search import actions
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.search import actions
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import tqdm
 
 tactions = set()

extra/optimization/helpers.py

@@ -1,6 +1,6 @@
 # stuff needed to unpack a kernel
 from tinygrad import Variable
-from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -9,7 +9,7 @@ inf, nan = float('inf'), float('nan')
 UOps = Ops
 
 # kernel unpacker
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 def ast_str_to_ast(ast_str:str) -> UOp: return eval(ast_str)
 def ast_str_to_lin(ast_str:str, opts=None): return Kernel(ast_str_to_ast(ast_str), opts=opts)
 def kern_str_to_lin(kern_str:str, opts=None):
@@ -101,7 +101,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
   return ret
 
 from tinygrad.device import Device, Buffer
-from tinygrad.engine.search import _ensure_buffer_alloc, _time_program
+from tinygrad.opt.search import _ensure_buffer_alloc, _time_program
 from tinygrad.helpers import to_function_name, CACHELEVEL, diskcache_get, diskcache_put
 
 def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_global_size=65536, cnt=3, disable_cache=False, clear_l2=False) -> float:  # noqa: E501

extra/optimization/pretrain_valuenet.py

@@ -1,4 +1,4 @@
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tqdm import tqdm, trange
 import math
 import random
@@ -14,7 +14,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.uop.ops import Variable
 inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 
 from extra.optimization.helpers import lin_to_feats, MAX_DIMS
 

extra/optimization/rl.py

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

extra/optimization/run_qnet.py

@@ -1,6 +1,6 @@
 from typing import List, Tuple
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.engine.search import get_kernel_actions, actions
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.search import get_kernel_actions, actions
 
 _net = None
 def beam_q_estimate(beam:List[Tuple[Kernel, float]]) -> List[Tuple[Kernel, float]]:

extra/optimization/search.py

@@ -4,8 +4,8 @@ from extra.optimization.helpers import ast_str_to_lin, time_linearizer
 from tinygrad import dtypes
 from tinygrad.helpers import BEAM, getenv
 from tinygrad.device import Device, Compiled
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.engine.search import beam_search, bufs_from_lin
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.search import beam_search, bufs_from_lin
 
 
 if __name__ == '__main__':

extra/optimization/test_net.py

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

extra/optimization/test_time_linearizer.py

@@ -1,5 +1,5 @@
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
-from tinygrad.engine.search import bufs_from_lin, get_kernel_actions
+from tinygrad.opt.search import bufs_from_lin, get_kernel_actions
 
 if __name__ == "__main__":
   ast_strs = load_worlds()

extra/profiler/index.html

@@ -1,173 +0,0 @@
-<!DOCTYPE html>
-<html>
-<head>
-  <title>tinygrad profiler</title>
-  <meta name="viewport" content="width=device-width, initial-scale=1.0">
-  <meta charset="UTF-8">
-  <link rel="icon" href="data:;base64,iVBORw0KGgo=">
-  <script src="assets/d3js.org/d3.v7.min.js" charset="utf-8"></script>
-  <style>
-  * {
-    box-sizing: border-box;
-    margin: 0;
-    padding: 0;
-  }
-  html, body {
-    color: #f0f0f5;
-    margin: 0;
-    padding: 0;
-    width: 100%;
-    height: 100%;
-    font-family: sans-serif;
-    font-optical-sizing: auto;
-    font-weight: 400;
-    font-style: normal;
-    font-variation-settings: "wdth" 100;
-    font-size: 14px;
-    overflow: hidden;
-    background-color: #08090e;
-  }
-  #root {
-    display:flex;
-    width:100%;
-    height: 100%;
-    padding: 20px;
-  }
-  #process-name {
-    background: #0f1018;
-    padding: 2px;
-    border-radius: 2px;
-  }
-  [id^="thread"] {
-    padding: 2px;
-  }
-  #table-root {
-    position: absolute;
-    width: 100%;
-    height: 300px;
-    background: #0f1018;
-    bottom: 0;
-    left: 0;
-    overflow: auto;
-  }
-  table {
-    border-collapse: collapse;
-    width: 100%;
-  }
-  table thead th {
-    position: sticky;
-    top: 0;
-    z-index: 1
-  }
-  th {
-    background: #1D1F2A;
-    cursor: pointer;
-  }
-  th, td {
-    padding: 8px 16px;
-    text-align: left;
-  }
-  th.sorted-asc::after { content: " ↑"; }
-  th.sorted-desc::after { content: " 	↓"; }
-  </style>
-</head>
-<body>
-<script>
-const colors = ["7aa2f7", "ff9e64", "f7768e", "2ac3de", "7dcfff", "1abc9c", "9ece6a", "e0af68", "bb9af7", "9d7cd8", "ff007c"];
-
-const formatTime = (ms) => {
-  if (ms<=1e3) return `${ms}us`;
-  if (ms<=1e6) return `${(ms*1e-3).toFixed(2)}ms`;
-  return `${(ms*1e-6).toFixed(2)}s`;
-}
-
-async function main() {
-  const { traceEvents } = await (await fetch("/get_profile")).json();
-  const root = createChild("div.root", document.querySelector("body"));
-  const list = createChild("div.list", root);
-  const data = [];
-  const nameColors = {}; // event names get a unique color
-  const procNames = {};
-  for (const e of traceEvents) {
-    if (e.name === "process_name") {
-      const proc = createChild(`div.proc-${e.pid}`, list);
-      createChild("p.process-name", proc).textContent = e.args.name;
-      procNames[e.pid] = e.args.name;
-    }
-    else if (e.name === "thread_name") {
-      const thread = createChild(`div.thread-${e.pid}-${e.tid}`, `proc-${e.pid}`);
-      createChild("p.thread-name", thread).textContent = e.args.name;
-    }
-    else if (e.ph === "X") {
-      const thread = document.getElementById(`thread-${e.pid}-${e.tid}`);
-      if (!(e.name in nameColors)) nameColors[e.name] = colors[data.length%(colors.length-1)];
-      data.push({ ...e, y:rect(thread).y, color:`#${nameColors[e.name]}`, proc:procNames[e.pid] });
-    }
-  }
-  // render graph
-  const svg = d3.select(root).append("svg").attr("width", "100%");
-  const { y, width } = rect(svg.node()); // global coordinates
-  const render = svg.append("g").attr("transform", `translate(0, ${y})`);
-  const timestamps = data.map(t => t.ts);
-  const st = Math.min(...timestamps);
-  const timeScale = d3.scaleLinear().domain([0, Math.max(...timestamps)-st]).range([y, width]);
-  const timeAxis = render.append("g").call(d3.axisTop(timeScale).tickFormat(formatTime));
-  list.style = `margin-top: ${rect(timeAxis.node()).bottom}px;`;
-  // rescale time based coordinates to fit screen
-  for (e of data) {
-    e.st = e.ts-st;
-    e.x = timeScale(e.st);
-    e.width = timeScale(e.dur);
-  }
-  render.selectAll("rect").data(data).join("rect").attr("fill", d => d.color).attr("x", d => d.x).attr("y", d => d.y).attr("width", d => d.width)
-    .attr("height", 20);
-  render.call(d3.brush().on("end", (e) => {
-    if (!e.selection) return renderTable({ data });
-    const [[x0, y0], [x1, y1]] = e.selection;
-    const newData = data.filter(d => d.x>=x0 && d.x<=x1 && d.y>=y0 && d.y<=y1);
-    renderTable({ data: newData });
-  }));
-  createChild("div.table-root", root);
-  renderTable({ data });
-}
-
-const rect = (e) => e.getBoundingClientRect();
-
-const createChild = (es, p) => {
-  const parts = es.split(".", 2);
-  if (typeof p === "string") p = document.getElementById(p);
-  const ret = p.appendChild(document.createElement(parts[0]));
-  if (parts.length !== 1) ret.id = parts[1];
-  return ret;
-}
-
-const columnNames = {"name":"Name", "st":"Start Time", "dur":"Duration", "proc":"Process"};
-const tableState = {data:null, sortBy:null, asc:true};
-function renderTable(newState) {
-  const { data, sortBy, asc } = Object.assign(tableState, newState);
-  const root = document.getElementById("table-root");
-  root.innerHTML = "";
-  const table = createChild("table", root);
-  const thead = createChild("tr", createChild("thead", table));
-  for (const [k,v] of Object.entries(columnNames)) {
-    const th = createChild(`th.${k}`, thead);
-    th.innerText = v;
-    th.onclick = (e) => renderTable(k === sortBy ? { asc:!asc } : { sortBy:k, asc:true });
-  }
-  if (sortBy != null) {
-    data.sort((a, b) => asc ? a[sortBy]-b[sortBy] : b[sortBy]-a[sortBy]); // inplace sort
-    document.getElementById(sortBy).className = asc ? "sorted-asc" : "sorted-desc";
-  }
-  const tbody = createChild("tbody", table);
-  for (const d of data) {
-    const row = createChild("tr", tbody);
-    for (const k of Object.keys(columnNames)) {
-      let formatted = typeof d[k] === "string" ? d[k] : formatTime(d[k]);
-      createChild("td", row).innerText = formatted;
-    }
-  }
-}
-main()
-</script>
-</body>
-</html>

extra/replay_pkl.py

@@ -6,8 +6,8 @@ from tinygrad.helpers import getenv, BEAM
 from tinygrad.engine.jit import TinyJit
 from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item
 from tinygrad.renderer import ProgramSpec
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.heuristic import hand_coded_optimizations
 import numpy as np
 
 def move_jit_captured_to_dev(captured, device="DSP"):

setup.py

@@ -26,7 +26,7 @@ setup(name='tinygrad',
       long_description_content_type='text/markdown',
       packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.runtime.autogen.am', 'tinygrad.codegen', 'tinygrad.nn',
                   'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support',
-                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop'],
+                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt'],
       package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'perfetto.html', 'assets/**/*', 'js/*']},
       classifiers=[
         "Programming Language :: Python :: 3",

test/external/external_benchmark_hcopt.py

@@ -1,7 +1,7 @@
 import random
 from tinygrad.helpers import getenv
-from tinygrad.engine.search import beam_search, bufs_from_lin
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.search import beam_search, bufs_from_lin
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
 
 def optimize_kernel(k):

test/external/external_benchmark_schedule.py

@@ -3,10 +3,11 @@ from extra.models.resnet import ResNet50
 from tinygrad import Tensor, nn
 from tinygrad.helpers import Profiling, Timing, getenv, BEAM, NOOPT, DEBUG, Context, ansilen
 from tinygrad.uop.ops import Ops
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites, rewrites_for_linearizer
-from tinygrad.engine.search import beam_search, bufs_from_lin
+from tinygrad.opt.search import beam_search, bufs_from_lin
+from tinygrad.uop.spec import type_verify
 
 if __name__ == "__main__":
   mdl = ResNet50()
@@ -56,4 +57,6 @@ if __name__ == "__main__":
             uops_line = []
             for u in rewritten_uops:
               uops_line.append(apply_rewrites(u, rewrites_for_linearizer))
+          with Timing("***** model verify in    "):
+            for u in uops_line: type_verify(u.arg.lst)
           print(sum(len(u.arg.lst) for u in uops_line))

test/external/external_benchmark_sdxl_softmax.py

@@ -7,7 +7,7 @@ if __name__ == "__main__":
   GlobalCounters.reset()
   t.softmax(-1, dtype="half", _single_kernel=True).realize()
 
-  from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+  from tinygrad.opt.kernel import Kernel, Opt, OptOps
   from tinygrad.helpers import get_single_element
   GlobalCounters.reset()
   si = get_single_element(t.softmax(-1, dtype="half", _single_kernel=True).schedule())

test/external/external_debug_metal_sd_conv.py

@@ -1,8 +1,8 @@
 # ruff: noqa: E501
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import CompiledRunner
-from tinygrad.engine.search import bufs_from_lin
+from tinygrad.opt.search import bufs_from_lin
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View

test/external/external_fuzz_ampt.py

@@ -29,7 +29,7 @@ class AMPTFuzzer:
         self.d.vram[pte['paddr']] = pattern # Mark this page
         assert pte['valid'] == 1
 
-        # If page has contigous fragment, all range should be this valid memory
+        # If page has contiguous fragment, all range should be this valid memory
         frags_cnt = pte['fragment']
         contig_range = (1 << (frags_cnt + 12))
         start_vaddr = _vaddr & ~(contig_range - 1)

test/external/external_test_am.py

@@ -91,6 +91,17 @@ class TestAMPageTable(unittest.TestCase):
           assert pte['paddr'] == 0
           assert pte['valid'] == 0
 
+  def test_map_notaligned(self):
+    mm0 = self.d[0].mm
+
+    for (va1,sz1),(va2,sz2) in [((0x10000, (0x1000)), (0x11000, (2 << 20)))]:
+      exteranl_va1 = va1 + AMMemoryManager.va_allocator.base
+      exteranl_va2 = va2 + AMMemoryManager.va_allocator.base
+      mm0.map_range(vaddr=exteranl_va1, size=sz1, paddrs=[(va1, sz1)])
+      mm0.map_range(vaddr=exteranl_va2, size=sz2, paddrs=[(va2, sz2)])
+      mm0.unmap_range(va2, sz2)
+      mm0.unmap_range(va1, sz1)
+
   def test_double_map(self):
     mm0 = self.d[0].mm
 

test/external/external_test_hcq_fuzz_failures.py

@@ -4,10 +4,10 @@ os.environ["VALIDATE_HCQ"]="1"
 
 import unittest, random
 import numpy as np
-from tinygrad.codegen.kernel import Kernel, KernelOptError
+from tinygrad.opt.kernel import Kernel, KernelOptError
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from test.external.fuzz_linearizer import compare_linearizer, compare_states, get_fuzz_rawbuf_like
 

test/external/external_test_hip_compile.py

@@ -3,7 +3,7 @@ from tinygrad.runtime.support.hip_comgr import compile_hip
 from tinygrad import Tensor
 from tinygrad.device import Device
 from tinygrad.engine.schedule import create_schedule
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 
 class TestHIPCompileSpeed(unittest.TestCase):
   @unittest.skipIf(Device.DEFAULT != "HIP", "only run on HIP")

test/external/external_test_nv.py

@@ -2,12 +2,12 @@ import unittest, struct, array, ctypes
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.helpers import to_mv
 from tinygrad.runtime.ops_nv import NVDevice, HWQueue
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from test.test_linearizer_failures import helper_test_lin
 from tinygrad.engine.realize import get_runner, CompiledRunner
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
 
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.uop.ops import LazyOp, Ops, ReduceOps, BufferOps, MemBuffer
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View

test/external/external_test_onnx_backend.py

@@ -44,6 +44,9 @@ class TinygradBackend(Backend):
 
 backend_test = onnx.backend.test.BackendTest(TinygradBackend, __name__)
 
+# BUG: segfaults
+backend_test.exclude('test_MaxPool1d_stride_padding_dilation_cpu')
+
 # BUG: buggy onnx tests
 backend_test.exclude('test_adam_multiple_cpu')
 

test/external/external_test_train_gpt2.py

@@ -2,11 +2,11 @@
 import unittest
 
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 
 from test.external.fuzz_linearizer import run_linearizer
 

test/external/external_test_valid_remove.py

@@ -3,11 +3,11 @@ import unittest
 
 from tinygrad import Device
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 
 class TestOpenpilotValidhack(unittest.TestCase):
   def test_valid_removal(self):

test/external/external_uop_gc.py

@@ -1,7 +1,7 @@
 import gc
 from tinygrad import Tensor, UOp, Device
 from tinygrad.shape.shapetracker import views_to_indexed_uops
-from tinygrad.engine.realize import method_cache, get_kernel
+from tinygrad.engine.realize import method_cache, get_program
 
 def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
 def print_uops():
@@ -14,12 +14,10 @@ def two_plus_two(): Tensor([2])+Tensor([2])
 def two_plus_two_schedule(): (Tensor([2])+Tensor([2])).schedule()
 def two_plus_two_kernel():
   si = (Tensor([2])+Tensor([2])).schedule()[-1]
-  get_kernel(Device.default.renderer, si.ast)
+  get_program(Device.default.renderer, si.ast)
 def two_plus_two_linearize():
   si = (Tensor([2])+Tensor([2])).schedule()[-1]
-  k = get_kernel(Device.default.renderer, si.ast)
-  k.get_optimized_ast()
-  #k.linearize()
+  get_program(Device.default.renderer, si.ast)
 def two_plus_two_realize(): (Tensor([2])+Tensor([2])).realize()
 def two_plus_two_item(): (Tensor([2])+Tensor([2])).item()
 def gradient_test():
@@ -36,7 +34,7 @@ def kernel_matmul():
   y = Tensor([[2.0,0,-2.0]], requires_grad=True)
   z = y.matmul(x)
   si = z.schedule()[-1]
-  get_kernel(Device.default.renderer, si.ast)
+  get_program(Device.default.renderer, si.ast)
 def realized_matmul():
   x = Tensor.eye(3, requires_grad=True)
   y = Tensor([[2.0,0,-2.0]], requires_grad=True)

test/external/fuzz_linearizer.py

@@ -20,9 +20,9 @@ if os.getenv("VALIDATE_HCQ", 0) != 0:
 
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.kernel import Opt, OptOps
-from tinygrad.engine.search import get_kernel_actions, bufs_from_lin
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.opt.search import get_kernel_actions, bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing
 from tinygrad.uop.ops import UOp, Ops

test/external/process_replay/process_replay.py

@@ -3,8 +3,8 @@
 import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools
 from typing import Callable, Any
 from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm, to_function_name
-from tinygrad.engine.grouper import get_kernelize_map
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.engine.kernelize import get_kernelize_map
+from tinygrad.opt.kernel import Kernel
 from tinygrad.uop.ops import UOp, Ops
 
 # *** process replay settings

test/external/speed_beam_v_hcopt.py

@@ -1,7 +1,7 @@
 from tinygrad import Device
 from tinygrad.helpers import getenv, DEBUG, BEAM
-from tinygrad.engine.search import beam_search, bufs_from_lin
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.search import beam_search, bufs_from_lin
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
 
 if __name__ == "__main__":

test/external/speed_compare_amd_am.py

@@ -2,8 +2,8 @@ from tinygrad import Device, dtypes
 from tinygrad.helpers import getenv, colorize_float, DEBUG
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
-from tinygrad.codegen.heuristic import hand_coded_optimizations
-from tinygrad.engine.search import bufs_from_lin
+from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.opt.search import bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.runtime.ops_amd import AMDDevice

test/external/speed_compare_cuda_nv.py

@@ -2,8 +2,8 @@ from tinygrad import Device, dtypes
 from tinygrad.helpers import getenv, colorize_float
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
-from tinygrad.codegen.heuristic import hand_coded_optimizations
-from tinygrad.engine.search import bufs_from_lin
+from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.opt.search import bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.tensor import _to_np_dtype
 import numpy as np

test/external/speed_compare_cuda_ptx.py

@@ -1,10 +1,10 @@
 import itertools
 from tinygrad import Device
 from tinygrad.engine.realize import CompiledRunner
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import getenv, colorize_float
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.engine.search import bufs_from_lin
+from tinygrad.opt.search import bufs_from_lin
 from tinygrad.runtime.ops_cuda import PTXCompiler, PTXRenderer, CUDACompiler
 
 if __name__ == "__main__":

test/external/speed_v_theoretical.py

@@ -87,11 +87,11 @@ class TestKernelSpeed(unittest.TestCase):
 
   # NOTE: tiny7 was slower than tiny12
   # TODO: why are convs so slow?!?
-  def test_conv_3x3_256_32_32_256_256(self): self._test_conv_3x3(256, 32, 32, 256, 256, nv_tflops=27, amd_tflops=20)
+  def test_conv_3x3_256_32_32_256_256(self): self._test_conv_3x3(256, 32, 32, 256, 256, nv_tflops=27, amd_tflops=14)
 
   # theoretical is nv_tflops=165, amd_tflops=123
   def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=115, amd_tflops=65)
-  def test_gemm_8192(self): self._test_matmul(8192, nv_tflops=125, amd_tflops=65)
+  def test_gemm_8192(self): self._test_matmul(8192, nv_tflops=125, amd_tflops=60)
 
   # theoretical is nv_gbs=1008, amd_gbs=960
   def test_gemv_16384_4096(self): self._test_matmul(16384, 4096, 1, nv_gbs=840, amd_gbs=750)

test/external/verify_kernel.py

@@ -3,7 +3,7 @@ from collections import defaultdict
 from extra.optimization.helpers import kern_str_to_lin, time_linearizer
 from test.external.fuzz_linearizer import compare_linearizer
 from tinygrad.helpers import colored
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 
 # Use this with the LOGKERNS options to verify that all executed kernels are valid and evaluate to the same ground truth results
 

test/helpers.py

@@ -6,7 +6,7 @@ from tinygrad.uop.ops import UOp, Ops, sint, graph_rewrite
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.engine.realize import Runner
-from tinygrad.engine.grouper import view_left
+from tinygrad.engine.kernelize import view_left
 from tinygrad.dtype import ConstType, DType
 from tinygrad.nn.state import get_parameters
 from tinygrad.helpers import T, unwrap, CI
@@ -59,8 +59,9 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
     bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
     allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + buf_dt.fmt, *data)))
   g = UOp(Ops.DEFINE_GLOBAL, uop.dtype.ptr(), arg=0, src=())
-  lst = full_rewrite(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), PythonRenderer)
-  prog = PythonProgram("run", PythonCompiler().compile(PythonRenderer().render(lst)))
+  opts = PythonRenderer()
+  lst = full_rewrite(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), opts)
+  prog = PythonProgram("run", PythonCompiler().compile(opts.render(lst)))
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
   return out_buf.cast(uop.dtype.fmt).tolist()[0]
 

test/models/test_onnx.py

@@ -105,6 +105,7 @@ class TestOnnxModel(unittest.TestCase):
       fetch("https://github.com/onnx/models/raw/main/validated/vision/classification/efficientnet-lite4/model/efficientnet-lite4-11.onnx"),
       input_name, input_new)
 
+  @unittest.skip("TODO: FIX THIS IT CAUSES SEGFAULT")
   def test_shufflenet(self):
     input_name, input_new = "gpu_0/data_0", False
     self._test_model(

test/test_arange.py

@@ -3,9 +3,9 @@ import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import CI, Context, getenv
 from tinygrad.engine.realize import run_schedule
-from tinygrad.codegen.kernel import Opt, OptOps, Kernel, KernelOptError
+from tinygrad.opt.kernel import Opt, OptOps, Kernel, KernelOptError
 from tinygrad.engine.realize import CompiledRunner, ExecItem
-from tinygrad.engine.search import get_kernel_actions
+from tinygrad.opt.search import get_kernel_actions
 from tinygrad.uop.ops import Ops
 
 class TestArange(unittest.TestCase):

test/test_dtype.py

@@ -1,15 +1,15 @@
-import unittest, operator, subprocess, math
+import unittest, math
 import numpy as np
 import torch
 from typing import Any, List
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, DEBUG, CI
-from tinygrad.dtype import DType, DTYPES_DICT, ImageDType, PtrDType, least_upper_float, least_upper_dtype, truncate_fp16, truncate_bf16, to_dtype
-from tinygrad.dtype import truncate, fp8_to_float, float_to_fp8
+from tinygrad.dtype import DType, DTYPES_DICT, ImageDType, PtrDType, least_upper_dtype, to_dtype, fp8_to_float, float_to_fp8
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.tensor import _to_np_dtype
 from hypothesis import assume, given, settings, strategies as strat
 from test.helpers import rand_for_dtype
+from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
 import ml_dtypes
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")
@@ -17,12 +17,7 @@ pytestmark = pytest.mark.filterwarnings("ignore")
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
 
-core_dtypes = list(DTYPES_DICT.values())
 if Device.DEFAULT == "CPU": core_dtypes.remove(dtypes.bfloat16)  # NOTE: this is for teenygrad, don't remove
-dtype_ints = [dt for dt in core_dtypes if dtypes.is_int(dt) and is_dtype_supported(dt)]
-dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and is_dtype_supported(dt)]
-FP8E4M3_MAX = 448.0
-FP8E5M2_MAX = 57344.0
 
 def get_available_cast_dtypes(dtype: DType) -> List[DType]:
   if not is_dtype_supported(dtype): return []
@@ -39,14 +34,6 @@ def _test_to_np(a:Tensor, np_dtype, target):
   except AssertionError as e:
     raise AssertionError(f"\ntensor {a.numpy()} does not match target {target} with np_dtype {np_dtype}") from e
 
-def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float=1e-7):
-  if DEBUG >= 2: print(tensor.numpy())
-  try:
-    assert tensor.dtype == target_dtype
-    np.testing.assert_allclose(tensor.numpy(), target, rtol={dtypes.float16:1e-3, dtypes.bfloat16:1e-2}.get(target_dtype, tol_target_dtype))
-  except AssertionError as e:
-    raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e
-
 def _test_op(fxn, target_dtype:DType, target):
   _assert_eq(fxn(), target_dtype, target)
 def _test_cast(a:Tensor, target_dtype:DType):
@@ -413,530 +400,6 @@ class TestEqStrDType(unittest.TestCase):
     self.assertEqual(str(dtypes.imagef((1,2,4))), "dtypes.imagef((1, 2, 4))")
     self.assertEqual(str(dtypes.float32.ptr(16)), "dtypes.float.ptr(16)")
 
-class TestHelpers(unittest.TestCase):
-  signed_ints = (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64)
-  uints = (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
-  floats = (dtypes.float16, dtypes.float32, dtypes.float64)
-
-  @given(strat.sampled_from(signed_ints+uints), strat.integers(min_value=1, max_value=8))
-  def test_is_int(self, dtype, amt):
-    assert dtypes.is_int(dtype.vec(amt) if amt > 1 else dtype)
-    assert not dtypes.is_float(dtype.vec(amt) if amt > 1 else dtype)
-
-  @given(strat.sampled_from(uints), strat.integers(min_value=1, max_value=8))
-  def test_is_unsigned_uints(self, dtype, amt):
-    assert dtypes.is_unsigned(dtype.vec(amt) if amt > 1 else dtype)
-
-  @given(strat.sampled_from(signed_ints), strat.integers(min_value=1, max_value=8))
-  def test_is_unsigned_signed_ints(self, dtype, amt):
-    assert not dtypes.is_unsigned(dtype.vec(amt) if amt > 1 else dtype)
-
-  @given(strat.sampled_from(floats), strat.integers(min_value=1, max_value=8))
-  def test_is_float(self, dtype, amt):
-    assert dtypes.is_float(dtype.vec(amt) if amt > 1 else dtype)
-    assert not dtypes.is_int(dtype.vec(amt) if amt > 1 else dtype)
-    assert not dtypes.is_unsigned(dtype.vec(amt) if amt > 1 else dtype)
-
-  def test_bf16_is_float(self):
-    assert dtypes.is_float(dtypes.bfloat16)
-
-  def test_fp8s_are_float(self):
-    assert dtypes.is_float(dtypes.fp8e4m3)
-    assert dtypes.is_float(dtypes.fp8e5m2)
-
-  @given(strat.sampled_from([d for d in DTYPES_DICT.values() if dtypes.is_float(d) or dtypes.is_int(d)]), strat.integers(min_value=2, max_value=8))
-  def test_scalar(self, dtype, amt):
-    assert dtype.vec(amt).scalar() == dtype
-
-  def test_from_py(self):
-    assert dtypes.from_py(True) == dtypes.bool
-    assert dtypes.from_py(2) == dtypes.default_int
-    assert dtypes.from_py(3.0) == dtypes.default_float
-    assert dtypes.from_py([]) == dtypes.default_float
-    assert dtypes.from_py(()) == dtypes.default_float
-    assert dtypes.from_py([True]) == dtypes.bool
-    assert dtypes.from_py([True, 2]) == dtypes.default_int
-    assert dtypes.from_py([True, 3.0]) == dtypes.default_float
-    assert dtypes.from_py([2, 3.0]) == dtypes.default_float
-    assert dtypes.from_py([True, 2, 3.0]) == dtypes.default_float
-    with self.assertRaises(RuntimeError): dtypes.from_py(None)
-    with self.assertRaises(RuntimeError): dtypes.from_py([None])
-    with self.assertRaises(RuntimeError): dtypes.from_py({})
-    with self.assertRaises(RuntimeError): dtypes.from_py(set())
-
-  def test_dtype_range(self):
-    for dt in core_dtypes:
-      if dtypes.is_float(dt):
-        np.testing.assert_equal(dtypes.min(dt), -math.inf)
-        np.testing.assert_equal(dtypes.max(dt), math.inf)
-        np.testing.assert_equal(dt.min, -math.inf)
-        np.testing.assert_equal(dt.max, math.inf)
-      elif dtypes.is_int(dt):
-        info = np.iinfo(_to_np_dtype(dt))
-        np.testing.assert_equal(dtypes.min(dt), info.min)
-        np.testing.assert_equal(dtypes.max(dt), info.max)
-        np.testing.assert_equal(dt.min, info.min)
-        np.testing.assert_equal(dt.max, info.max)
-      else:
-        assert dt == dtypes.bool, dt
-        np.testing.assert_equal(dtypes.min(dt), False)
-        np.testing.assert_equal(dtypes.max(dt), True)
-        np.testing.assert_equal(dt.min, False)
-        np.testing.assert_equal(dt.max, True)
-
-  def test_truncate_fp16(self):
-    self.assertEqual(truncate_fp16(1), 1)
-    self.assertEqual(truncate_fp16(65504), 65504)
-    self.assertEqual(truncate_fp16(65519.999), 65504)
-    self.assertEqual(truncate_fp16(65520), math.inf)
-
-  def test_truncate_bf16(self):
-    self.assertEqual(truncate_bf16(1), 1)
-    self.assertAlmostEqual(truncate_bf16(1.1), 1.09375, places=7)
-    for a in [1234, 23456, -777.777]:
-      self.assertEqual(truncate_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
-    # TODO: torch bfloat 1.1 gives 1.1015625 instead of 1.09375
-    max_bf16 = torch.finfo(torch.bfloat16).max
-    self.assertEqual(truncate_bf16(max_bf16), max_bf16)
-    self.assertEqual(truncate_bf16(min_bf16:=-max_bf16), min_bf16)
-    self.assertEqual(truncate_bf16(max_bf16 * 1.00001), math.inf)
-    self.assertEqual(truncate_bf16(min_bf16 * 1.00001), -math.inf)
-
-  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
-  def test_truncate_fp8e4m3(self, x):
-    if x > FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), FP8E4M3_MAX)
-    elif x < -FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), -FP8E4M3_MAX)
-    else: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), ml_dtypes.float8_e4m3fn(x))
-
-  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
-  def test_truncate_fp8e5m2(self, x):
-    if x > FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), FP8E5M2_MAX)
-    elif x < -FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), -FP8E5M2_MAX)
-    else: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), ml_dtypes.float8_e5m2(x))
-
-class TestTypeSpec(unittest.TestCase):
-  def setUp(self):
-    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
-  def tearDown(self):
-    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
-
-  def test_set_dtype_default(self):
-    for default_int in [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64]:
-      dtypes.default_int = default_int
-      assert dtypes.default_int == default_int
-
-    for default_float in [*dtypes.fp8s, dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
-      dtypes.default_float = default_float
-      assert dtypes.default_float == default_float
-
-  @unittest.skip("this test is slow and spawning whole pythons")
-  def test_env_set_default_float(self):
-    # check default
-    subprocess.run(['python3 -c "from tinygrad import dtypes; assert dtypes.default_float == dtypes.float"'],
-                    shell=True, check=True)
-    # check change
-    subprocess.run(['DEFAULT_FLOAT=HALF python3 -c "from tinygrad import dtypes; assert dtypes.default_float == dtypes.half"'],
-                    shell=True, check=True)
-    # check invalid
-    with self.assertRaises(subprocess.CalledProcessError):
-      subprocess.run(['DEFAULT_FLOAT=INT32 python3 -c "from tinygrad import dtypes"'],
-                      shell=True, check=True)
-
-    with self.assertRaises(subprocess.CalledProcessError):
-      subprocess.run(['DEFAULT_FLOAT=TYPO python3 -c "from tinygrad import dtypes"'],
-                      shell=True, check=True)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.int8), f"no int8 on {Device.DEFAULT}")
-  def test_dtype_str_arg(self):
-    n = np.random.normal(0, 1, (10, 10)).astype(np.float32)
-    tested = 0
-    for dtype_str, dtype in [
-      ("bool", dtypes.bool), ("int8", dtypes.int8), ("int", dtypes.int), ("uint32", dtypes.uint32), ("float32", dtypes.float32)]:
-      np.testing.assert_equal(Tensor(n, dtype=dtype_str).numpy(), Tensor(n, dtype=dtype).numpy())
-      np.testing.assert_equal(Tensor(n).cast(dtype_str).numpy(), Tensor(n).cast(dtype).numpy())
-      if dtype.itemsize == 4:
-        np.testing.assert_equal(Tensor(n).bitcast(dtype_str).numpy(), Tensor(n).bitcast(dtype).numpy())
-        tested += 1
-    assert tested == 3
-
-    with self.assertRaises(AttributeError): Tensor([1, 2, 3], dtype="nonexistdtype")
-    with self.assertRaises(AttributeError): Tensor([1, 2, 3], dtype="")
-
-    np.testing.assert_equal(Tensor(n).sum(dtype="int16").numpy(), Tensor(n).sum(dtype=dtypes.int16).numpy())
-
-  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
-  def test_creation(self, default_int, default_float):
-    dtypes.default_int, dtypes.default_float = default_int, default_float
-    _assert_eq(Tensor(True), dtypes.bool, True)
-    _assert_eq(Tensor(None), dtypes.default_float, [])
-    _assert_eq(Tensor(2), dtypes.default_int, 2)
-    _assert_eq(Tensor(2.34), dtypes.default_float, 2.34)
-    _assert_eq(Tensor([]), dtypes.default_float, [])
-    _assert_eq(Tensor([1]), dtypes.default_int, [1])
-    _assert_eq(Tensor([1.1]), dtypes.default_float, [1.1])
-
-    _assert_eq(Tensor.eye(0), dtypes.default_float, np.eye(0))
-    _assert_eq(Tensor.eye(3), dtypes.default_float, np.eye(3))
-    if is_dtype_supported(dtypes.int64):
-      _assert_eq(Tensor.eye(3, dtype=dtypes.int64), dtypes.int64, np.eye(3))
-    if is_dtype_supported(dtypes.float16):
-      _assert_eq(Tensor.eye(3, dtype=dtypes.float16), dtypes.float16, np.eye(3))
-
-  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
-  def test_full(self, default_int, default_float):
-    dtypes.default_int, dtypes.default_float = default_int, default_float
-
-    _assert_eq(Tensor.zeros((2, 3)), dtypes.default_float, np.zeros((2, 3)))
-    if is_dtype_supported(dtypes.int64):
-      _assert_eq(Tensor.zeros((2, 3), dtype=dtypes.int64), dtypes.int64, np.zeros((2, 3)))
-    if is_dtype_supported(dtypes.float16):
-      _assert_eq(Tensor.zeros((2, 3), dtype=dtypes.float16), dtypes.float16, np.zeros((2, 3)))
-
-    _assert_eq(Tensor.ones((2, 3)), dtypes.default_float, np.ones((2, 3)))
-    if is_dtype_supported(dtypes.int64):
-      _assert_eq(Tensor.ones((2, 3), dtype=dtypes.int64), dtypes.int64, np.ones((2, 3)))
-    if is_dtype_supported(dtypes.float16):
-      _assert_eq(Tensor.ones((2, 3), dtype=dtypes.float16), dtypes.float16, np.ones((2, 3)))
-
-    _assert_eq(Tensor.full((2, 3), 3.0), dtypes.default_float, np.full((2, 3), 3.0))
-    _assert_eq(Tensor.full((2, 3), 3), dtypes.default_int, np.full((2, 3), 3))
-    _assert_eq(Tensor.full((2, 3), True), dtypes.bool, np.full((2, 3), True))
-    if is_dtype_supported(dtypes.int64):
-      _assert_eq(Tensor.full((2, 3), 3, dtype=dtypes.int64), dtypes.int64, np.full((2, 3), 3))
-      _assert_eq(Tensor.full((2, 3), 3.0, dtype=dtypes.int64), dtypes.int64, np.full((2, 3), 3))
-    if is_dtype_supported(dtypes.float16):
-      _assert_eq(Tensor.full((2, 3), 3, dtype=dtypes.float16), dtypes.float16, np.full((2, 3), 3))
-      _assert_eq(Tensor.full((2, 3), 3.0, dtype=dtypes.float16), dtypes.float16, np.full((2, 3), 3))
-
-  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
-  def test_reduce_0d_default(self, default_int, default_float):
-    dtypes.default_int, dtypes.default_float = default_int, default_float
-    _assert_eq(Tensor.ones((2,3,0)).sum(2), dtypes.default_float, np.zeros((2, 3)))
-    # TODO: what should this one be?
-    # _assert_eq(Tensor.ones((2,3,0), dtype=dtypes.default_int).sum(2), dtypes.default_int, np.zeros((2, 3)))
-    _assert_eq(Tensor.ones((2,3,0), dtype=dtypes.int32).sum(2), dtypes.int32, np.zeros((2, 3)))
-
-  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
-  def test_arange(self, default_int, default_float):
-    dtypes.default_int, dtypes.default_float = default_int, default_float
-
-    _assert_eq(Tensor.arange(5), dtypes.default_int, np.arange(5))
-    _assert_eq(Tensor.arange(120), dtypes.default_int, np.arange(120))
-    _assert_eq(Tensor.arange(5.0), dtypes.default_float, np.arange(5))
-    if is_dtype_supported(dtypes.int16):
-      _assert_eq(Tensor.arange(5, dtype=dtypes.int16), dtypes.int16, np.arange(5))
-    if is_dtype_supported(dtypes.int64):
-      _assert_eq(Tensor.arange(5, dtype=dtypes.int64), dtypes.int64, np.arange(5))
-    if is_dtype_supported(dtypes.float16):
-      _assert_eq(Tensor.arange(5, dtype=dtypes.float16), dtypes.float16, np.arange(5))
-    _assert_eq(Tensor.arange(3, 9, 0.7), dtypes.default_float, np.arange(3, 9, 0.7), 1e-6 if Device.DEFAULT == "WEBGPU" else 1e-7)
-    _assert_eq(Tensor.arange(3, 8.5, 3), dtypes.default_float, np.arange(3, 8.5, 3))
-    # stop-start and step have different signs
-    _assert_eq(Tensor.arange(3, 5, -2), dtypes.default_int, np.arange(3, 5, -2))
-    _assert_eq(Tensor.arange(5.0, 3.0), dtypes.default_float, np.arange(5.0, 3.0))
-
-  @given(strat.sampled_from(core_dtypes), strat.sampled_from([operator.gt, operator.ge, operator.le, operator.lt, operator.eq, operator.ne]))
-  def test_bool_ops(self, dtype, op):
-    assert op(Tensor.ones(4, 4, dtype=dtype), Tensor.ones(4, 4, dtype=dtype)).dtype == dtypes.bool
-
-  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
-  def test_functions_return_index(self, dtype, default_int, default_float):
-    dtypes.default_int, dtypes.default_float = default_int, default_float
-    assert Tensor([0, 1], dtype=dtype).argmax().dtype == dtypes.int32
-    assert Tensor([0, 1], dtype=dtype).argmin().dtype == dtypes.int32
-    assert Tensor([0, 1], dtype=dtype).multinomial().dtype == dtypes.int32
-
-  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints))
-  def test_tensor_indexing_returns_same_dtype(self, data_dtype, indices_dtype):
-    X_data =  Tensor.ones(60000, 1, 28, 28, dtype=data_dtype)
-    indices =  Tensor.randint(512, high=X_data.shape[0]).cast(indices_dtype)
-    assert X_data[indices].dtype == X_data.dtype
-
-  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints))
-  def test_gather_returns_same_dtype(self, data_dtype, indices_dtype):
-    X_data = Tensor([[1, 0], [0, 1]], dtype=data_dtype)
-    indices = Tensor([[0, 0], [1, 0]], dtype=indices_dtype)
-    assert X_data.gather(0, indices).dtype == X_data.dtype
-    assert X_data.gather(1, indices).dtype == X_data.dtype
-
-  @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
-  def test_attention_returns_same_dtype(self, data_dtype, default_float):
-    dtypes.default_float = default_float
-    query = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
-    key = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
-    value = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
-    mask = (Tensor.rand(32, 8, 128, 128) < 0.5)
-    assert query.scaled_dot_product_attention(key, value, is_causal=True).dtype == data_dtype
-    assert query.scaled_dot_product_attention(key, value, is_causal=True, dropout_p=0.3).dtype == data_dtype
-    assert query.scaled_dot_product_attention(key, value, is_causal=False).dtype == data_dtype
-    assert query.scaled_dot_product_attention(key, value, attn_mask=mask).dtype == data_dtype
-
-class TestTypePromotion(unittest.TestCase):
-  @given(strat.sampled_from(core_dtypes))
-  def test_self_promo_to_self(self, dtype):
-    assert least_upper_dtype(dtype) == dtype
-    assert least_upper_dtype(dtype, dtype) == dtype
-    assert least_upper_dtype(dtype, dtype, dtype) == dtype
-
-  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
-  def test_promo_resulted_higher_than_inputs(self, dtype1, dtype2):
-    result = least_upper_dtype(dtype1, dtype2)
-    assert not (result < dtype1) and not (result < dtype2)
-
-  def test_dtype_promo(self):
-    assert least_upper_dtype(dtypes.bool, dtypes.int8) == dtypes.int8
-    assert least_upper_dtype(dtypes.int8, dtypes.uint8) == dtypes.int16
-    assert least_upper_dtype(dtypes.uint8, dtypes.int16) == dtypes.int16
-    assert least_upper_dtype(dtypes.int16, dtypes.uint16) == dtypes.int32
-    assert least_upper_dtype(dtypes.uint16, dtypes.int32) == dtypes.int32
-    assert least_upper_dtype(dtypes.int32, dtypes.uint32) == dtypes.int64
-    assert least_upper_dtype(dtypes.uint32, dtypes.int64) == dtypes.int64
-    # similar to jax but we don't use weak type
-    assert least_upper_dtype(dtypes.int64, dtypes.uint64) == dtypes.float16
-    assert least_upper_dtype(dtypes.float16, dtypes.float32) == dtypes.float32
-    assert least_upper_dtype(dtypes.float32, dtypes.float64) == dtypes.float64
-
-    assert least_upper_dtype(dtypes.bool, dtypes.float32) == dtypes.float32
-    assert least_upper_dtype(dtypes.bool, dtypes.float64) == dtypes.float64
-    assert least_upper_dtype(dtypes.float16, dtypes.int64) == dtypes.float16
-    assert least_upper_dtype(dtypes.float16, dtypes.uint64) == dtypes.float16
-    assert least_upper_dtype(dtypes.fp8e4m3, dtypes.fp8e5m2) == dtypes.half
-
-class TestAutoCastType(unittest.TestCase):
-  def setUp(self):
-    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
-  def tearDown(self):
-    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
-
-  @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
-  def test_least_upper_float_input_is_float(self, input_dtype, default_float):
-    dtypes.default_float = default_float
-    self.assertEqual(least_upper_float(input_dtype), input_dtype)
-
-  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
-  def test_least_upper_float_input_is_int(self, input_dtype, default_float):
-    dtypes.default_float = default_float
-    self.assertEqual(least_upper_float(input_dtype), default_float)
-
-  @given(strat.sampled_from([d for d in core_dtypes if dtypes.is_int(d) and is_dtype_supported(d)]))
-  def test_int_to_float_unary_func(self, dtype):
-    for func in [
-      lambda t: t.exp(),
-      lambda t: t.exp2(),
-      lambda t: t.log(),
-      lambda t: t.log2(),
-      lambda t: t.sqrt(),
-      lambda t: t.rsqrt(),
-      lambda t: t.sin(),
-      lambda t: t.cos(),
-      lambda t: t.tan(),
-      lambda t: t.sigmoid(),
-    ]:
-      a = [2, 3, 4]
-      # float16 can have larger precision errors
-      np.testing.assert_allclose(func(Tensor(a, dtype=dtype)).numpy(), func(torch.tensor(a)), rtol=1e-3, atol=1e-3)
-
-  @given(strat.sampled_from(core_dtypes))
-  def test_broadcast_scalar(self, dt):
-    assert (Tensor.ones(4, 4, dtype=dt) + 2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
-    assert (Tensor.ones(4, 4, dtype=dt) + 2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
-    assert (Tensor.ones(4, 4, dtype=dt) + True).dtype == dt
-
-  @given(strat.sampled_from(dtype_floats))
-  def test_int_div_int(self, default_float):
-    dtypes.default_float = default_float
-    self.assertEqual(Tensor([1]).div(Tensor([2])).dtype, default_float)
-
-  def test_sum(self):
-    assert (Tensor([0, 1], dtype=dtypes.bool)).sum().dtype == dtypes.int32
-    assert (Tensor([0, 1], dtype=dtypes.int8)).sum().dtype == dtypes.int32
-    assert (Tensor([0, 1], dtype=dtypes.int16)).sum().dtype == dtypes.int32
-    assert (Tensor([0, 1], dtype=dtypes.int32)).sum().dtype == dtypes.int32
-    assert (Tensor([0, 1], dtype=dtypes.int64)).sum().dtype == dtypes.int64
-    assert (Tensor([0, 1], dtype=dtypes.uint8)).sum().dtype == dtypes.uint32
-    assert (Tensor([0, 1], dtype=dtypes.uint16)).sum().dtype == dtypes.uint32
-    assert (Tensor([0, 1], dtype=dtypes.uint32)).sum().dtype == dtypes.uint32
-    assert (Tensor([0, 1], dtype=dtypes.uint64)).sum().dtype == dtypes.uint64
-    assert (Tensor([0, 1], dtype=dtypes.float16)).sum().dtype == dtypes.float16
-    #assert (Tensor([0, 1], dtype=dtypes.bfloat16)).sum().dtype == dtypes.bfloat16
-    assert (Tensor([0, 1], dtype=dtypes.float32)).sum().dtype == dtypes.float32
-    assert (Tensor([0, 1], dtype=dtypes.float64)).sum().dtype == dtypes.float64
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16")
-  def test_sum_dtype_arg(self):
-    t = Tensor([40000, 40000], dtype=dtypes.float16)
-    # default float16 sum returns in float16, overflowed in this case
-    assert t.sum().dtype == dtypes.float16
-    assert math.isinf(t.sum().numpy().item())
-    # specifiying dtype and it's not downcasted
-    assert t.sum(dtype=dtypes.float32).dtype == dtypes.float32
-    np.testing.assert_allclose(t.sum(dtype=dtypes.float32).numpy(), 80000)
-
-  def test_prod_dtype_arg(self):
-    t = Tensor([100, 200], dtype=dtypes.int32)
-    assert t.prod().dtype == dtypes.int32
-    np.testing.assert_allclose(t.prod().numpy(), 20000)
-    assert t.prod(dtype=dtypes.float32).dtype == dtypes.float32
-    np.testing.assert_allclose(t.prod(dtype=dtypes.float32).numpy(), 20000)
-
-  def test_mean(self):
-    assert (Tensor([0, 1], dtype=dtypes.bool)).mean().dtype == dtypes.float32
-    assert (Tensor([0, 1], dtype=dtypes.int8)).mean().dtype == dtypes.float32
-    assert (Tensor([0, 1], dtype=dtypes.int16)).mean().dtype == dtypes.float32
-    assert (Tensor([0, 1], dtype=dtypes.int32)).mean().dtype == dtypes.float32
-    assert (Tensor([0, 1], dtype=dtypes.int64)).mean().dtype == dtypes.float32
-    assert (Tensor([0, 1], dtype=dtypes.uint8)).mean().dtype == dtypes.float32
-    assert (Tensor([0, 1], dtype=dtypes.uint16)).mean().dtype == dtypes.float32
-    assert (Tensor([0, 1], dtype=dtypes.uint32)).mean().dtype == dtypes.float32
-    assert (Tensor([0, 1], dtype=dtypes.uint64)).mean().dtype == dtypes.float32
-    assert (Tensor([0, 1], dtype=dtypes.float16)).mean().dtype == dtypes.float16
-    #assert (Tensor([0, 1], dtype=dtypes.bfloat16)).mean().dtype == dtypes.bfloat16
-    assert (Tensor([0, 1], dtype=dtypes.float32)).mean().dtype == dtypes.float32
-    assert (Tensor([0, 1], dtype=dtypes.float64)).mean().dtype == dtypes.float64
-
-  def test_cumsum(self):
-    assert (Tensor([0, 1], dtype=dtypes.bool)).cumsum(0).dtype == dtypes.int32
-    assert (Tensor([0, 1], dtype=dtypes.int8)).cumsum(0).dtype == dtypes.int32
-    assert (Tensor([0, 1], dtype=dtypes.int16)).cumsum(0).dtype == dtypes.int32
-    assert (Tensor([0, 1], dtype=dtypes.int32)).cumsum(0).dtype == dtypes.int32
-    assert (Tensor([0, 1], dtype=dtypes.int64)).cumsum(0).dtype == dtypes.int64
-    assert (Tensor([0, 1], dtype=dtypes.uint8)).cumsum(0).dtype == dtypes.uint32
-    assert (Tensor([0, 1], dtype=dtypes.uint16)).cumsum(0).dtype == dtypes.uint32
-    assert (Tensor([0, 1], dtype=dtypes.uint32)).cumsum(0).dtype == dtypes.uint32
-    assert (Tensor([0, 1], dtype=dtypes.uint64)).cumsum(0).dtype == dtypes.uint64
-    assert (Tensor([0, 1], dtype=dtypes.float16)).cumsum(0).dtype == dtypes.float16
-    #assert (Tensor([0, 1], dtype=dtypes.bfloat16)).cumsum(0).dtype == dtypes.bfloat16
-    assert (Tensor([0, 1], dtype=dtypes.float32)).cumsum(0).dtype == dtypes.float32
-    assert (Tensor([0, 1], dtype=dtypes.float64)).cumsum(0).dtype == dtypes.float64
-
-  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
-  def test_matmul(self, dt1, dt2, acc_dt):
-    t1 = Tensor([0, 1], dtype=dt1)
-    t2 = Tensor([0, 1], dtype=dt2)
-    self.assertEqual(t1.matmul(t2).dtype, least_upper_dtype(t1.dtype, t2.dtype))
-    # if dtype is specified, return in dtype
-    self.assertEqual(t1.matmul(t2, dtype=acc_dt).dtype, acc_dt)
-
-  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
-  def test_linear(self, dt1, dt2, dt3, acc_dt):
-    x = Tensor([0, 1], dtype=dt1)
-    w = Tensor([0, 1], dtype=dt2)
-    b = Tensor([0, 1], dtype=dt3)
-    self.assertEqual(x.linear(w).dtype, least_upper_dtype(x.dtype, w.dtype))
-    self.assertEqual(x.linear(w, b).dtype, least_upper_dtype(least_upper_dtype(x.dtype, w.dtype), b.dtype))
-    # if dtype is specified, return in dtype
-    self.assertEqual(x.linear(w, dtype=acc_dt).dtype, acc_dt)
-    self.assertEqual(x.linear(w, b, dtype=acc_dt).dtype, acc_dt)
-
-  @staticmethod
-  def check_where_alternate_input_other(input_, other, data_type):
-    assert (Tensor([True, False]).where(input_, other)).dtype == data_type
-    assert (Tensor([True, False]).where(other, input_)).dtype == data_type
-
-  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
-  def test_where_no_scalar(self, dt1, dt2):
-    self.check_where_alternate_input_other(Tensor(2, dtype=dt1), Tensor(3, dtype=dt2), least_upper_dtype(dt1, dt2))
-
-  @given(strat.sampled_from(core_dtypes))
-  def test_where_one_scalar(self, dt):
-    t = Tensor(2, dtype=dt)
-    self.check_where_alternate_input_other(t, 3.2, (dt if dtypes.is_float(dt) else dtypes.default_float))
-    self.check_where_alternate_input_other(t, 3, (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int))
-    self.check_where_alternate_input_other(t, True, dt)
-
-  def test_where_two_scalars(self):
-    self.check_where_alternate_input_other(3.1, 3.2, dtypes.default_float)
-    self.check_where_alternate_input_other(3.1, 3, dtypes.default_float)
-    self.check_where_alternate_input_other(3.1, True, dtypes.default_float)
-    self.check_where_alternate_input_other(3, 2, dtypes.default_int)
-    self.check_where_alternate_input_other(3, True, dtypes.default_int)
-    self.check_where_alternate_input_other(False, True, dtypes.bool)
-
-  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
-  def test_maximum(self, dt1, dt2):
-    assert Tensor([0, 1, 2], dtype=dt1).maximum(Tensor([2, 0, 5], dtype=dt2)).dtype == least_upper_dtype(dt1, dt2)
-
-  @given(strat.sampled_from(core_dtypes))
-  def test_maximum_const(self, dt):
-    assert Tensor([1, 2], dtype=dt).maximum(3.1).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
-    assert Tensor([1, 2], dtype=dt).maximum(3).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
-    assert Tensor([1, 2], dtype=dt).maximum(True).dtype == dt
-
-  def test_div(self):
-    assert (Tensor([1, 2], dtype=dtypes.int32) / Tensor([2, 2], dtype=dtypes.int32)).dtype == dtypes.default_float
-    assert (Tensor([1, 2], dtype=dtypes.int16) / Tensor([2, 2], dtype=dtypes.int32)).dtype == dtypes.default_float
-    assert (Tensor([1, 2], dtype=dtypes.float32) / Tensor([2, 2], dtype=dtypes.float16)).dtype == dtypes.float32
-    assert (Tensor([1, 2], dtype=dtypes.int32) / Tensor([2, 2], dtype=dtypes.float16)).dtype == dtypes.float16
-
-  def test_div_const(self):
-    assert (Tensor([1, 2], dtype=dtypes.int32) / 2).dtype == dtypes.default_float
-    assert (Tensor([1, 2], dtype=dtypes.int32) / 2.0).dtype == dtypes.default_float
-    assert (Tensor([1, 2], dtype=dtypes.float16) / 2).dtype == dtypes.float16
-    assert (Tensor([1, 2], dtype=dtypes.float16) / 2.0).dtype == dtypes.float16
-
-  def test_gradient_dtype(self):
-    old_default_float = dtypes.default_float
-
-    for default_dtype in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
-      if not is_dtype_supported(default_dtype): continue
-      dtypes.default_float = default_dtype
-      for dtype in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
-        if not is_dtype_supported(dtype): continue
-        if DEBUG >= 2:
-          print(f"testing {default_dtype=}, {dtype=}")
-        a = Tensor([1, 2, 3], dtype=dtype, requires_grad=True)
-        b = (a * 5).sum()
-        b.backward()  # if there is dtype mismatch, lazy should assert
-        assert a.grad.dtype == a.dtype
-        np.testing.assert_allclose(a.grad.numpy(), [5, 5, 5])
-
-    dtypes.default_float = old_default_float
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  def test_backward_sum_acc_dtype(self):
-    # test acc of sum in the backward is upcasted to float
-    t = Tensor([5, -5], dtype=dtypes.half, requires_grad=True)
-    t.reshape(2, 1).expand(2, 10001).max().backward()
-    np.testing.assert_allclose(t.grad.numpy(), [1, 0])
-
-  @unittest.skipIf(Device.DEFAULT == "PYTHON", "very slow")
-  @unittest.skipIf(CI and Device.DEFAULT == "AMD", "very slow")
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Binding size is larger than the maximum storage buffer binding size")
-  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  def test_mean_half_precision_underflow(self):
-    N = 10000
-    x = 0.001
-    t = Tensor([[x]], dtype=dtypes.half, requires_grad=True).expand(N, N).contiguous()
-    np.testing.assert_allclose(t.mean(axis=1).numpy(), np.array([x] * N, dtype=np.float16), rtol=1e-3)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  def test_mean_half_precision_overflow(self):
-    N = 256
-    t = Tensor([60000] * N*N, dtype=dtypes.half, requires_grad=True).reshape(N, N)
-    np.testing.assert_allclose(t.mean().numpy(), 60000)
-    t.square().mean().backward()
-    np.testing.assert_allclose(t.grad.numpy().flatten(), [60000 * 2 / (N*N)] * N*N)
-
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Precision error")
-  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  def test_softmax_dtype(self):
-    data = [1, 2, 3]
-    t = Tensor(data, dtype=dtypes.half)
-    tt = torch.tensor(data, dtype=torch.half)
-
-    out = t.softmax(0)
-    self.assertEqual(out.dtype, dtypes.half)
-    np.testing.assert_allclose(out.numpy(), tt.softmax(0).numpy(), rtol=1e-3)
-    out = t.softmax(0, dtype=dtypes.float)
-    self.assertEqual(out.dtype, dtypes.float)
-    np.testing.assert_allclose(out.numpy(), tt.softmax(0, dtype=torch.float).numpy(), rtol=1e-3)
-    out = t.log_softmax(0)
-    self.assertEqual(out.dtype, dtypes.half)
-    np.testing.assert_allclose(out.numpy(), tt.log_softmax(0).numpy(), rtol=1e-3)
-    out = t.log_softmax(0, dtype=dtypes.float)
-    self.assertEqual(out.dtype, dtypes.float)
-    np.testing.assert_allclose(out.numpy(), tt.log_softmax(0, dtype=torch.float).numpy(), rtol=1e-3)
-
 class TestImplicitFunctionTypeChange(unittest.TestCase):
   def test_functions(self):
     result = []
@@ -980,5 +443,13 @@ class TestToDtype(unittest.TestCase):
     self.assertIsInstance(res, DType)
     self.assertEqual(res, dtypes.int32)
 
+@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
+class TestOpsBFloat16(unittest.TestCase):
+  def test_cast(self):
+    # TODO: helper_test_op breaks in unrelated part
+    # TODO: wrong output with GPU=1 / PYTHON=1 on mac
+    data = [60000.0, 70000.0, 80000.0]
+    np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy())
+
 if __name__ == '__main__':
   unittest.main()

test/test_gc.py

@@ -86,10 +86,10 @@ class TestGC(unittest.TestCase):
     a.realize()
     real_buf = a.uop.buffer
     # after the Tensor UOp is deleted there shouldn't be any references on the Buffer
-    self.assertEqual(real_buf.lb_refcount, 1)
+    self.assertEqual(real_buf.uop_refcount, 1)
     self.assertEqual(bufs_allocated()-init, 1)
     del a.uop
-    self.assertEqual(real_buf.lb_refcount, 0)
+    self.assertEqual(real_buf.uop_refcount, 0)
     self.assertEqual(bufs_allocated()-init, 1) # keep the buffer alive
     del real_buf
     self.assertEqual(bufs_allocated()-init, 0)
@@ -99,14 +99,14 @@ class TestGC(unittest.TestCase):
     a = Tensor.full((4,), 1.).contiguous()
     a.realize()
     real_buf = a.uop.buffer
-    self.assertEqual(real_buf.lb_refcount, 1)
+    self.assertEqual(real_buf.uop_refcount, 1)
     a.assign(Tensor.full((4,), 2.))
     self.assertIs(a.uop.src[0].buffer, real_buf)
     # NOTE: this is still 1, we don't count the ASSIGN
-    self.assertEqual(real_buf.lb_refcount, 1)
+    self.assertEqual(real_buf.uop_refcount, 1)
     a.realize()
     del a
-    self.assertEqual(real_buf.lb_refcount, 0) # no UOps for this Buffer
+    self.assertEqual(real_buf.uop_refcount, 0) # no UOps for this Buffer
     self.assertEqual(bufs_allocated()-init, 1) # Buffer is alive
 
 if __name__ == '__main__':

test/test_hcq.py

@@ -5,7 +5,7 @@ from tinygrad.device import Buffer, BufferSpec
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.engine.realize import get_runner, CompiledRunner
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad import Variable
 
 MOCKGPU = getenv("MOCKGPU")

test/test_linearizer.py

@@ -4,7 +4,7 @@ import unittest
 from dataclasses import replace
 
 from test.helpers import ast_const
-from tinygrad.codegen.kernel import Opt, OptOps, KernelOptError, Kernel
+from tinygrad.opt.kernel import Opt, OptOps, KernelOptError, Kernel
 from tinygrad.codegen.lowerer import get_grouped_dims
 from tinygrad.uop.ops import UOp, Ops, GroupOp
 from tinygrad.device import Device, Buffer, is_dtype_supported
@@ -12,7 +12,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import prod, Context, getenv, CI, flatten, dedup, AMX
 from tinygrad.dtype import DType, dtypes
 
@@ -1136,7 +1136,7 @@ class TestLinearizer(unittest.TestCase):
         np.testing.assert_allclose(result, golden_result, atol=0.1, rtol=0.2)
 
       # check that get_kernel_actions produces all 9 options
-      from tinygrad.engine.search import get_kernel_actions
+      from tinygrad.opt.search import get_kernel_actions
       tc_actions = [k for i, k in get_kernel_actions(Kernel(realized_ast), False).items() if k.applied_opts[0].op == OptOps.TC]
 
       available_tc = len([x for x in Device[Device.DEFAULT].renderer.tensor_cores if x.dtype_in == tc.dtype_in and x.dtype_out == tc.dtype_out])
@@ -2027,7 +2027,13 @@ class TestKernelOpts(unittest.TestCase):
         [Opt(OptOps.UNROLL, 0, 2), Opt(OptOps.UPCAST, 0, 4)],
         [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 2), Opt(OptOps.UPCAST, 1, 4)],
         [Opt(OptOps.UNROLL, 0, 2), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)],
-        # [Opt(OptOps.GROUP, 0, 2)] # doesn't work because group_for_reduce dims become early locals (conflicting with TC)
+        [Opt(OptOps.GROUP, 0, 2)],
+        [Opt(OptOps.GROUPTOP, 0, 4)],
+        [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.GROUP, 0, 2)],
+        [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUP, 0, 2)],
+        [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.GROUP, 0, 2)],
+        [Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUP, 0, 2)],
+        [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.UNROLL, 0, 2), Opt(OptOps.UPCAST, 1, 2)],
       ], apply_tc=True, atol=atol, rtol=rtol)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")

test/test_linearizer_dumb.py

@@ -8,8 +8,8 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.helpers import getenv
 from tinygrad.shape.shapetracker import ShapeTracker, View
-from tinygrad.engine.search import Opt, OptOps
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.search import Opt, OptOps
+from tinygrad.opt.kernel import Kernel
 
 class TestLinearizerDumb(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")

test/test_linearizer_failures.py

@@ -1,10 +1,10 @@
 # ruff: noqa: E501
 import unittest, random
 import numpy as np
-from tinygrad.codegen.kernel import Kernel, KernelOptError
+from tinygrad.opt.kernel import Kernel, KernelOptError
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.helpers import CI, Context
 from test.external.fuzz_linearizer import compare_linearizer

test/test_linearizer_overflows.py

@@ -2,8 +2,8 @@
 import unittest
 from tinygrad import dtypes, Device
 from tinygrad.helpers import CI
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.engine.search import Opt, OptOps, bufs_from_lin
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.search import Opt, OptOps, bufs_from_lin
 from extra.optimization.helpers import time_linearizer
 
 # stuff needed to unpack a kernel

test/test_memory_planner.py

@@ -18,13 +18,13 @@ def check_assign(buffers:list[list[Buffer]|tuple[Buffer, ...]]):
   first_appearance, last_appearance = {}, {}
   for i,u in enumerate(buffers):
     for buf in u:
-      if buf.is_allocated() or buf.base.is_allocated() or buf.lb_refcount > 0: continue
+      if buf.is_allocated() or buf.base.is_allocated() or buf.uop_refcount > 0: continue
       if buf.base not in first_appearance: first_appearance[buf.base] = i
       last_appearance[buf.base] = i
 
   for i,u in enumerate(buffers):
     for buf in u:
-      if buf.is_allocated() or buf.base.is_allocated() or buf.lb_refcount > 0: continue
+      if buf.is_allocated() or buf.base.is_allocated() or buf.uop_refcount > 0: continue
       cur, base = assigned.get(buf, buf), assigned.get(buf.base, buf.base)
       if buf._base is not None:
         assert cur.base == base.base and cur.offset == buf.offset + base.offset, f"failed: {buf} {cur} {base} {buf.offset} {base.offset}"

test/test_multitensor.py

@@ -681,13 +681,15 @@ class TestMultiTensor(unittest.TestCase):
     self.assertEqual(d0_rand, d1_rand_flip)
     self.assertEqual(d1_rand, d0_rand_flip)
 
-  def test_rand_like_on_shard(self):
-    t = Tensor.empty((16, 16)).shard(devices_2)
+  def test_rand_like_on_shard(self, axis=None):
+    t = Tensor.empty((16, 16)).shard(devices_2, axis=axis)
     t2 = Tensor.rand_like(t)
     self.assertEqual(t.shape, t2.shape)
     self.assertEqual(t.device, t2.device)
     self.assertEqual(t.dtype, t2.dtype)
     self.assertEqual(t.uop.axis, t2.uop.axis)
+    t2.realize()
+  def test_rand_like_on_shard_axis(self): self.test_rand_like_on_shard(0)
 
   def test_rand_like_from_alu(self):
     a = Tensor.ones(4, 4).shard(devices_4, axis=0)

test/test_ops.py

@@ -1093,8 +1093,8 @@ class TestOps(unittest.TestCase):
   def test_sort(self):
     for dim in [-1, 0, 1]:
       for descending in [True, False]:
-        helper_test_op([(8,45,6)], lambda x: x.sort(dim, descending).values, lambda x: x.sort(dim, descending)[0], forward_only=True)
-        helper_test_op([(8,45,6)], lambda x: x.sort(dim, descending).indices.type(torch.int32), lambda x: x.sort(dim, descending)[1],
+        helper_test_op([(8,8,6)], lambda x: x.sort(dim, descending).values, lambda x: x.sort(dim, descending)[0], forward_only=True)
+        helper_test_op([(8,8,6)], lambda x: x.sort(dim, descending).indices.type(torch.int32), lambda x: x.sort(dim, descending)[1],
                        forward_only=True)
     # repeated values
     helper_test_op(None, lambda x: x.sort(stable=True).values, lambda x: x.sort()[0], forward_only=True, vals=[[0, 1] * 9])
@@ -1110,12 +1110,12 @@ class TestOps(unittest.TestCase):
     for dim in [0, 1, -1]:
       for largest in [True, False]:
         for sorted_ in [True]: # TODO support False
-          helper_test_op([(10,12,6)],
-                          lambda x: x.topk(5, dim, largest, sorted_).values,
-                          lambda x: x.topk(5, dim, largest, sorted_)[0], forward_only=True)
-          helper_test_op([(10,12,6)],
-                          lambda x: x.topk(5, dim, largest, sorted_).indices.type(torch.int32),
-                          lambda x: x.topk(5, dim, largest, sorted_)[1], forward_only=True)
+          helper_test_op([(6,5,4)],
+                          lambda x: x.topk(4, dim, largest, sorted_).values,
+                          lambda x: x.topk(4, dim, largest, sorted_)[0], forward_only=True)
+          helper_test_op([(5,5,4)],
+                          lambda x: x.topk(4, dim, largest, sorted_).indices.type(torch.int32),
+                          lambda x: x.topk(4, dim, largest, sorted_)[1], forward_only=True)
     # repeated values
     value, indices = Tensor([1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0]).topk(3)
     np.testing.assert_equal(value.numpy(), [1, 1, 1])
@@ -1981,106 +1981,106 @@ class TestOps(unittest.TestCase):
 
   def test_simple_conv2d(self):
     helper_test_op([(1,4,9,9), (4,4,3,3)],
-      lambda x,w: torch.nn.functional.conv2d(x,w).relu(),
-      lambda x,w: Tensor.conv2d(x,w).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv2d(x,w),
+      lambda x,w: Tensor.conv2d(x,w), grad_rtol=1e-5)
 
   def test_simple_conv2d_bias(self):
     helper_test_op([(1,4,9,9), (4,4,3,3), (4,)],
-      lambda x,w,b: torch.nn.functional.conv2d(x,w,b).relu(),
-      lambda x,w,b: Tensor.conv2d(x,w,b).relu(), grad_rtol=1e-5)
+      lambda x,w,b: torch.nn.functional.conv2d(x,w,b),
+      lambda x,w,b: Tensor.conv2d(x,w,b), grad_rtol=1e-5)
 
   @unittest.skipIf(IMAGE>0, "no conv3d on images")
   def test_simple_conv3d(self):
     helper_test_op([(1,4,9,9,9), (4,4,3,3,3)],
-      lambda x,w: torch.nn.functional.conv3d(x,w).relu(),
-      lambda x,w: Tensor.conv2d(x,w).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv3d(x,w),
+      lambda x,w: Tensor.conv2d(x,w), grad_rtol=1e-5)
 
   @unittest.skipIf(IMAGE>0, "no conv3d on images")
   def test_padded_conv3d(self):
     helper_test_op([(1,4,5,5,5), (4,4,3,3,3)],
-      lambda x,w: torch.nn.functional.conv3d(x,w,padding=1).relu(),
-      lambda x,w: Tensor.conv2d(x,w,padding=[1,1,1,1,1,1]).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv3d(x,w,padding=1),
+      lambda x,w: Tensor.conv2d(x,w,padding=[1,1,1,1,1,1]), grad_rtol=1e-5)
 
   def test_simple_conv2d_m4(self):
     helper_test_op([(1,16,18,18), (16,16,3,3)],
-      lambda x,w: torch.nn.functional.conv2d(x,w).relu(),
-      lambda x,w: Tensor.conv2d(x,w).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv2d(x,w),
+      lambda x,w: Tensor.conv2d(x,w), atol=1e-05, grad_rtol=1e-5)
 
   def test_simple_conv2d_1x1(self):
     helper_test_op([(1,4,9,9), (4,4,1,1)],
-      lambda x,w: torch.nn.functional.conv2d(x,w).relu(),
-      lambda x,w: Tensor.conv2d(x,w).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv2d(x,w),
+      lambda x,w: Tensor.conv2d(x,w), grad_rtol=1e-5)
 
   def test_simple_conv2d_1x1_m4(self):
     helper_test_op([(1,16,32,32), (16,16,1,1)],
-      lambda x,w: torch.nn.functional.conv2d(x,w).relu(),
-      lambda x,w: Tensor.conv2d(x,w).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv2d(x,w),
+      lambda x,w: Tensor.conv2d(x,w), grad_rtol=1e-5)
 
   def test_nested_conv2d(self):
     helper_test_op([(1,32,9,9), (32,32,3,3), (32,32,3,3)],
-      lambda x,w1,w2: torch.nn.functional.conv2d(torch.nn.functional.conv2d(x,w1).relu(), w2).relu(),
-      lambda x,w1,w2: x.conv2d(w1).relu().conv2d(w2).relu())
+      lambda x,w1,w2: torch.nn.functional.conv2d(torch.nn.functional.conv2d(x,w1).relu(), w2),
+      lambda x,w1,w2: x.conv2d(w1).relu().conv2d(w2))
 
   # expect reduce nodes == 3
   def test_simple_conv2d_nhwc(self):
     # weights (from tf): filter_height x filter_width x in_channels x out_channels
     helper_test_op([(2,9,9,10), (3,3,10,20)],
-      lambda x,w: torch.nn.functional.conv2d(x.permute(0,3,1,2),w.permute(3,2,0,1)).relu(),
-      lambda x,w: Tensor.conv2d(x.permute(0,3,1,2),w.permute(3,2,0,1)).relu(), atol=1e-5, grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv2d(x.permute(0,3,1,2),w.permute(3,2,0,1)),
+      lambda x,w: Tensor.conv2d(x.permute(0,3,1,2),w.permute(3,2,0,1)), atol=1e-5, grad_rtol=1e-5)
 
   def test_simple_conv2d_batched(self):
     helper_test_op([(2,4,9,9), (4,4,3,3)],
-      lambda x,w: torch.nn.functional.conv2d(x,w).relu(),
-      lambda x,w: Tensor.conv2d(x,w).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv2d(x,w),
+      lambda x,w: Tensor.conv2d(x,w), grad_rtol=1e-5)
 
   # conv transpose
 
   def test_simple_conv_transpose2d(self):
     helper_test_op([(2,4,9,9), (4,4,3,3)],
-      lambda x,w: torch.nn.functional.conv_transpose2d(x,w).relu(),
-      lambda x,w: Tensor.conv_transpose2d(x,w).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv_transpose2d(x,w),
+      lambda x,w: Tensor.conv_transpose2d(x,w), grad_rtol=1e-5)
 
   def test_bias_conv_transpose2d(self):
     helper_test_op([(2,4,9,9), (4,4,3,3), (4,)],
-      lambda x,w,b: torch.nn.functional.conv_transpose2d(x,w,b).relu(),
-      lambda x,w,b: Tensor.conv_transpose2d(x,w,b).relu(), grad_rtol=1e-5)
+      lambda x,w,b: torch.nn.functional.conv_transpose2d(x,w,b),
+      lambda x,w,b: Tensor.conv_transpose2d(x,w,b), grad_rtol=1e-5)
 
   def test_grouped_conv_transpose2d(self):
     helper_test_op([(2,4,9,9), (4,4,3,3)],
-      lambda x,w: torch.nn.functional.conv_transpose2d(x,w,groups=2).relu(),
-      lambda x,w: Tensor.conv_transpose2d(x,w,groups=2).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv_transpose2d(x,w,groups=2),
+      lambda x,w: Tensor.conv_transpose2d(x,w,groups=2), grad_rtol=1e-5)
 
   def test_padded_conv_transpose2d(self):
     for padding in [(1,2), (2,1), 2, 1, 0]:
       helper_test_op([(2,4,9,9), (4,4,3,3)],
-        lambda x,w: torch.nn.functional.conv_transpose2d(x,w,padding=padding).relu(),
-        lambda x,w: Tensor.conv_transpose2d(x,w,padding=padding).relu(), grad_rtol=1e-5)
+        lambda x,w: torch.nn.functional.conv_transpose2d(x,w,padding=padding),
+        lambda x,w: Tensor.conv_transpose2d(x,w,padding=padding), grad_rtol=1e-5)
     self.helper_test_exception([(2,16,2,2), (32,16,3,3)], lambda x,w: torch.nn.functional.conv_transpose2d(x,w,padding=(1,1,1)),
                    lambda x,w: Tensor.conv_transpose2d(x,w,padding=(1,1,1)), expected=(RuntimeError, ValueError))
 
   def test_dilated_conv_transpose2d(self):
     for dilation in [(1,2), (2,1), 2, 1]:
       helper_test_op([(2,4,9,9), (4,4,3,3)],
-        lambda x,w: torch.nn.functional.conv_transpose2d(x,w,dilation=dilation).relu(),
-        lambda x,w: Tensor.conv_transpose2d(x,w,dilation=dilation).relu(), grad_rtol=1e-5)
+        lambda x,w: torch.nn.functional.conv_transpose2d(x,w,dilation=dilation),
+        lambda x,w: Tensor.conv_transpose2d(x,w,dilation=dilation), grad_rtol=1e-5)
 
   def test_strided_conv_transpose2d(self):
     for stride in [(2,1), (1,2), 1]:
       helper_test_op([(2,4,4,5), (4,4,3,3)],
-        lambda x,w: torch.nn.functional.conv_transpose2d(x,w, stride=stride).relu(),
-        lambda x,w: Tensor.conv_transpose2d(x,w,stride=stride).relu(), atol=1e-5, grad_rtol=1e-5)
+        lambda x,w: torch.nn.functional.conv_transpose2d(x,w, stride=stride),
+        lambda x,w: Tensor.conv_transpose2d(x,w,stride=stride), atol=1e-5, grad_rtol=1e-5)
 
   def test_output_padded_conv_transpose2d(self):
     for output_padding, stride in [((1,1), (2,3)), ((2,1), (3,2))]:
       helper_test_op([(2,4,6,5), (4,4,3,3),(4,)],
-        lambda x,w,b: torch.nn.functional.conv_transpose2d(x,w,b,output_padding=output_padding,stride=stride).relu(),
-        lambda x,w,b: Tensor.conv_transpose2d(x,w,b,output_padding=output_padding,stride=stride).relu(), grad_rtol=1e-5)
+        lambda x,w,b: torch.nn.functional.conv_transpose2d(x,w,b,output_padding=output_padding,stride=stride),
+        lambda x,w,b: Tensor.conv_transpose2d(x,w,b,output_padding=output_padding,stride=stride), grad_rtol=1e-5)
 
   @unittest.skipIf(IMAGE>0, "no conv3d on images")
   def test_simple_conv_transpose3d(self):
     helper_test_op([(2,4,9,9,9), (4,4,3,3,3)],
-      lambda x,w: torch.nn.functional.conv_transpose3d(x,w).relu(),
-      lambda x,w: Tensor.conv_transpose2d(x,w).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv_transpose3d(x,w),
+      lambda x,w: Tensor.conv_transpose2d(x,w), grad_rtol=1e-5)
 
   @unittest.skipIf((IMAGE>0), "no conv1d on images")
   def test_conv1d(self):
@@ -2090,8 +2090,8 @@ class TestOps(unittest.TestCase):
           for groups in [1,3] if cin == 3 and H == 5 else [1]:
             with self.subTest(batch_size=bs, channels=cin, groups=groups, height=H):
               helper_test_op([(bs,cin,11), (6,cin//groups,H)],
-                lambda x,w: torch.nn.functional.conv1d(x,w,groups=groups).relu(),
-                lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), grad_rtol=1e-5)
+                lambda x,w: torch.nn.functional.conv1d(x,w,groups=groups),
+                lambda x,w: Tensor.conv2d(x,w,groups=groups), grad_rtol=1e-5)
 
   @unittest.skipIf(IMAGE>0, "no conv1d on images")
   def test_simple_padding_conv1d(self):
@@ -2101,15 +2101,15 @@ class TestOps(unittest.TestCase):
     H = 5
     p = (1,1)
     helper_test_op([(bs,cin,11), (6,cin//groups,H)],
-      lambda x,w: torch.nn.functional.conv1d(torch.nn.functional.pad(x, p),w).relu(),
-      lambda x,w: Tensor.conv2d(x,w,padding=p).relu())
+      lambda x,w: torch.nn.functional.conv1d(torch.nn.functional.pad(x, p),w),
+      lambda x,w: Tensor.conv2d(x,w,padding=p))
 
   @unittest.skipIf(IMAGE>0, "no conv1d on images")
   def test_strided_conv1d_simple(self):
     bs, H = 2, 3
     helper_test_op([(bs,1,5), (1,1,H)],
-      lambda x,w: torch.nn.functional.conv1d(x,w,stride=2).relu(),
-      lambda x,w: Tensor.conv2d(x,w,stride=2).relu())
+      lambda x,w: torch.nn.functional.conv1d(x,w,stride=2),
+      lambda x,w: Tensor.conv2d(x,w,stride=2))
 
   @unittest.skipIf(IMAGE>0, "no conv1d on images")
   def test_asymmetric_padding_conv1d(self):
@@ -2118,8 +2118,8 @@ class TestOps(unittest.TestCase):
         for n in [3,4]:
           for k in [2]:
             helper_test_op([(1,1,n), (1,1,k)],
-              lambda x,w: torch.nn.functional.conv1d(torch.nn.functional.pad(x, p),w).relu(),
-              lambda x,w: Tensor.conv2d(x,w,padding=p).relu())
+              lambda x,w: torch.nn.functional.conv1d(torch.nn.functional.pad(x, p),w),
+              lambda x,w: Tensor.conv2d(x,w,padding=p))
 
   def _test_conv2d(self, bs=1, cin=1, cout=6):
     for H in [2,3]:
@@ -2127,8 +2127,8 @@ class TestOps(unittest.TestCase):
         for groups in [1,3] if cin == 3 and cout == 6 and H == 3 and W == 3 else [1]:
           with self.subTest(batch_size=bs, channels=cin, groups=groups, height=H, width=W):
             helper_test_op([(bs,cin,5,7), (cout,cin//groups,H,W)],
-              lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
-              lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), grad_rtol=1e-5)
+              lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups),
+              lambda x,w: Tensor.conv2d(x,w,groups=groups), grad_rtol=1e-5)
   def test_conv2d(self): self._test_conv2d(bs=1, cin=3)
   def test_conv2d_bs_4_cin_3(self): self._test_conv2d(bs=4, cin=3, cout=2)
   def test_conv2d_bs_1_cin_1(self): self._test_conv2d(bs=1, cin=1)
@@ -2152,9 +2152,9 @@ class TestOps(unittest.TestCase):
     H = 5
     W = 2
     helper_test_op([(bs,cin,64,64), (6,cin//groups,H,W)],
-      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
-      # needed to relax tolerance on NVIDIA
-      lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), atol=1e-4, grad_atol=1e-4, grad_rtol=1e-4)
+      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups),
+      # needed to relax tolerance for larger input
+      lambda x,w: Tensor.conv2d(x,w,groups=groups), atol=1e-4, grad_atol=3e-4, grad_rtol=1e-4)
 
   def test_simple_grouped_conv2d(self):
     bs = 1
@@ -2162,8 +2162,8 @@ class TestOps(unittest.TestCase):
     rcout = 1
     cin = 2
     helper_test_op([(bs,groups*cin,1,1), (groups*rcout,cin,1,1)],
-      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
-      lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups),
+      lambda x,w: Tensor.conv2d(x,w,groups=groups), grad_rtol=1e-5)
 
   def test_medium_grouped_conv2d(self):
     bs = 1
@@ -2171,8 +2171,8 @@ class TestOps(unittest.TestCase):
     rcout = 2
     cin = 2
     helper_test_op([(bs,groups*cin,1,1), (groups*rcout,cin,1,1)],
-      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
-      lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups),
+      lambda x,w: Tensor.conv2d(x,w,groups=groups), grad_rtol=1e-5)
 
   def test_depthwise_conv2d(self):
     bs = 1
@@ -2180,8 +2180,8 @@ class TestOps(unittest.TestCase):
     rcout = 1
     cin = 1
     helper_test_op([(bs,groups*cin,32,32), (groups*rcout,cin,1,1)],
-      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
-      lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups),
+      lambda x,w: Tensor.conv2d(x,w,groups=groups), grad_rtol=1e-5)
 
   def test_grouped_conv2d(self):
     bs = 4
@@ -2189,8 +2189,8 @@ class TestOps(unittest.TestCase):
     rcout = 7
     cin = 3
     helper_test_op([(bs,groups*cin,5,5), (groups*rcout,cin,3,3)],
-      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
-      lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups),
+      lambda x,w: Tensor.conv2d(x,w,groups=groups), grad_rtol=1e-5)
 
   def test_fancy_conv2d(self):
     bs = 2
@@ -2199,14 +2199,14 @@ class TestOps(unittest.TestCase):
     groups = 3
     H,W = 3,3
     helper_test_op([(bs,cin,11,28), (groups*cout,cin//groups,H,W)],
-      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
-      lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), grad_rtol=1e-5)
+      lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups),
+      lambda x,w: Tensor.conv2d(x,w,groups=groups), grad_rtol=1e-5)
 
   def test_strided_conv2d_simple(self):
     bs,H,W = 2,3,1
     helper_test_op([(bs,1,5,1), (1,1,H,W)],
-      lambda x,w: torch.nn.functional.conv2d(x,w,stride=2).relu(),
-      lambda x,w: Tensor.conv2d(x,w,stride=2).relu())
+      lambda x,w: torch.nn.functional.conv2d(x,w,stride=2),
+      lambda x,w: Tensor.conv2d(x,w,stride=2))
 
   @unittest.skipIf(Device.DEFAULT != "LLVM", "DEVECTORIZE=0 only for LLVM")
   def test_strided_conv2d_simple_vec(self):
@@ -2218,27 +2218,27 @@ class TestOps(unittest.TestCase):
     H,W = 3,3
     with self.subTest(stride := 2):
       helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
-        lambda x,w: torch.nn.functional.conv2d(x,w,stride=2).relu(),
-        lambda x,w: Tensor.conv2d(x,w,stride=stride).relu())
+        lambda x,w: torch.nn.functional.conv2d(x,w,stride=2),
+        lambda x,w: Tensor.conv2d(x,w,stride=stride))
     with self.subTest(stride := (2,1)):
       helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
-        lambda x,w: torch.nn.functional.conv2d(x,w,stride=stride).relu(),
-        lambda x,w: Tensor.conv2d(x,w,stride=(2,1)).relu())
+        lambda x,w: torch.nn.functional.conv2d(x,w,stride=stride),
+        lambda x,w: Tensor.conv2d(x,w,stride=(2,1)))
 
   def test_negative_padding_conv2d(self):
     n,k = 10, 3
     helper_test_op([(1,1,n,n), (1,1,k,k)],
-      lambda x,w: torch.nn.functional.conv2d(x[:, :, 1:-1, 1:-1],w).relu(),
-      lambda x,w: Tensor.conv2d(x,w,padding=-1).relu())
+      lambda x,w: torch.nn.functional.conv2d(x[:, :, 1:-1, 1:-1],w),
+      lambda x,w: Tensor.conv2d(x,w,padding=-1))
     helper_test_op([(1,1,n,n), (1,1,k,k)],
-      lambda x,w: torch.nn.functional.conv2d(x[:, :, 1:, 1:],w).relu(),
-      lambda x,w: Tensor.conv2d(x,w,padding=(-1,0,-1,0)).relu())
+      lambda x,w: torch.nn.functional.conv2d(x[:, :, 1:, 1:],w),
+      lambda x,w: Tensor.conv2d(x,w,padding=(-1,0,-1,0)))
 
   def test_simple_padding_conv2d(self):
     p = (1,1,1,1)
     helper_test_op(None,
-      lambda x,w: torch.nn.functional.conv2d(torch.nn.functional.pad(x, p),w).relu(),
-      lambda x,w: Tensor.conv2d(x,w,padding=p).relu(), vals=[[[[[2.,3.]]]], [[[[1.]]]]])
+      lambda x,w: torch.nn.functional.conv2d(torch.nn.functional.pad(x, p),w),
+      lambda x,w: Tensor.conv2d(x,w,padding=p), vals=[[[[[2.,3.]]]], [[[[1.]]]]])
 
   def test_asymmetric_padding_conv2d(self):
     for p in [(0,1,0,1), (2,1,2,1), (2,0,2,1)]:
@@ -2246,35 +2246,35 @@ class TestOps(unittest.TestCase):
         for n in [3,4]:
           for k in [2]:
             helper_test_op([(1,1,n,n), (1,1,k,k)],
-              lambda x,w: torch.nn.functional.conv2d(torch.nn.functional.pad(x, p),w).relu(),
-              lambda x,w: Tensor.conv2d(x,w,padding=p).relu())
+              lambda x,w: torch.nn.functional.conv2d(torch.nn.functional.pad(x, p),w),
+              lambda x,w: Tensor.conv2d(x,w,padding=p))
             helper_test_op([(1,1,n,n), (1,1,k,k)],
-              lambda x,w: torch.nn.functional.conv2d(torch.nn.functional.pad(x, p),w).relu(),
-              lambda x,w: Tensor.conv2d(x,w,padding=p).relu())
+              lambda x,w: torch.nn.functional.conv2d(torch.nn.functional.pad(x, p),w),
+              lambda x,w: Tensor.conv2d(x,w,padding=p))
 
   def test_padded_conv2d_p21(self):
     bs,cin,H,W,padding = 4, 3, 3, 3, (2,1)
     helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
-      lambda x,w: torch.nn.functional.conv2d(x,w,padding=padding).relu(),
-      lambda x,w: Tensor.conv2d(x,w,padding=padding).relu())
+      lambda x,w: torch.nn.functional.conv2d(x,w,padding=padding),
+      lambda x,w: Tensor.conv2d(x,w,padding=padding))
 
   def test_padded_conv2d_p22(self):
     bs,cin,H,W,padding = 4, 3, 3, 3, (2,2)
     helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
-      lambda x,w: torch.nn.functional.conv2d(x,w,padding=padding).relu(),
-      lambda x,w: Tensor.conv2d(x,w,padding=padding).relu())
+      lambda x,w: torch.nn.functional.conv2d(x,w,padding=padding),
+      lambda x,w: Tensor.conv2d(x,w,padding=padding))
 
   def test_padded_conv2d_1x1(self):
     bs,cin,H,W,padding = 4, 3, 1, 1, 2
     helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
-      lambda x,w: torch.nn.functional.conv2d(x,w,padding=padding).relu(),
-      lambda x,w: Tensor.conv2d(x,w,padding=padding).relu())
+      lambda x,w: torch.nn.functional.conv2d(x,w,padding=padding),
+      lambda x,w: Tensor.conv2d(x,w,padding=padding))
 
   def test_padded_conv2d_bs1(self):
     bs,cin,H,W,padding = 1, 3, 3, 3, 1
     helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
-      lambda x,w: torch.nn.functional.conv2d(x,w,padding=padding).relu(),
-      lambda x,w: Tensor.conv2d(x,w,padding=padding).relu())
+      lambda x,w: torch.nn.functional.conv2d(x,w,padding=padding),
+      lambda x,w: Tensor.conv2d(x,w,padding=padding))
 
   def test_padding_add(self):
     helper_test_op([(64,64), (60,60)],
@@ -2288,8 +2288,8 @@ class TestOps(unittest.TestCase):
     for d in [2, (2,1)]:
       with self.subTest(dilation := d):
         helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
-          lambda x,w: torch.nn.functional.conv2d(x,w,dilation=dilation).relu(),
-          lambda x,w: Tensor.conv2d(x,w,dilation=dilation).relu())
+          lambda x,w: torch.nn.functional.conv2d(x,w,dilation=dilation),
+          lambda x,w: Tensor.conv2d(x,w,dilation=dilation))
 
   def test_max_pool2d_simple(self):
     ksz = (2,2)
@@ -3028,14 +3028,6 @@ class TestOpsUint8(unittest.TestCase):
       lambda x: x.type(torch.uint8).min(),
       lambda x: x.cast(dtypes.uint8).min(), forward_only=True, vals=[[0, 128, 255, 64, 32, 16]])
 
-@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
-class TestOpsBFloat16(unittest.TestCase):
-  def test_cast(self):
-    # TODO: helper_test_op breaks in unrelated part
-    # TODO: wrong output with GPU=1 / PYTHON=1 on mac
-    data = [60000.0, 70000.0, 80000.0]
-    np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy())
-
 if __name__ == '__main__':
   np.random.seed(1337)
   unittest.main(verbosity=2)

test/test_opt_gemm.py

@@ -2,7 +2,7 @@ import numpy as np
 import unittest
 from tinygrad import Tensor
 from tinygrad.helpers import get_single_element
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem
 
 class TestOptGemm(unittest.TestCase):

test/test_quantize_onnx.py

@@ -4,9 +4,9 @@ import unittest
 from dataclasses import replace
 from tinygrad import Tensor, Context, Device, dtypes
 from tinygrad.uop.ops import Ops, UOp # noqa: F401 # pylint: disable=unused-import
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item
-from tinygrad.engine.search import bufs_from_lin
+from tinygrad.opt.search import bufs_from_lin
 from tinygrad.shape.shapetracker import ShapeTracker, View # noqa: F401 # pylint: disable=unused-import
 
 N = 512

test/test_remote.py

@@ -0,0 +1,65 @@
+import numpy as np, unittest, string
+from hypothesis import given, strategies as st
+from tinygrad import Device, Tensor, TinyJit
+from tinygrad.runtime.ops_remote import RemoteDevice, parse_hosts
+from tinygrad.helpers import LazySeq, all_same
+
+def multihost_env(devices):
+  def same_hosts(devices): return all_same([h for h,_ in devices])
+  return isinstance(devices, list) and len(devices) >= 12 and not same_hosts(devices[0:12]) and same_hosts(devices[0:6]) and same_hosts(devices[6:12])
+
+@unittest.skipUnless(Device.DEFAULT == "REMOTE" and multihost_env(RemoteDevice.devices), "Requires special environment")
+class TestRemoteMultiHost(unittest.TestCase):
+  def test_mutlihost_transfer(self):
+    a = Tensor.arange(0, 16, device='REMOTE:0').contiguous().realize()
+    b = a.to('REMOTE:6').contiguous().realize()
+    np.testing.assert_equal(b.numpy(), np.arange(0, 16))
+
+  # NOTE: remote graph currently throws GraphException on host mismatch, this just checks that it is being handled, not that jit graph is being used
+  def test_multihost_matmul_jit(self):
+    @TinyJit
+    def do(a:Tensor, b:Tensor): return (a @ b).contiguous().realize()
+    ds = ('REMOTE:0', 'REMOTE:1', 'REMOTE:6', 'REMOTE:7')
+    for _ in range(3):
+      na, nb = np.random.rand(128, 128).astype(np.float32), np.random.rand(128, 128).astype(np.float32)
+      a, b = Tensor(na).shard(ds, 0).contiguous().realize(), Tensor(nb).shard(ds, 0).contiguous().realize()
+      nc = na @ nb
+      c = do(a, b)
+      np.testing.assert_allclose(nc, c.numpy(), rtol=3e-2, atol=1e-4) # tolerances from extra/gemm/simple_matmul.py
+
+class TestParseHosts(unittest.TestCase):
+  def assert_seq(self, result:LazySeq, host:str):
+    self.assertIsInstance(result, LazySeq)
+    for i in [0, 1, 5, 10]: self.assertEqual(result[i], (host, i))
+
+  @given(st.sampled_from(["", "localhost", "192.168.1.1:8080", "host"]))
+  def test_single_host_no_count(self, host:str):
+    self.assert_seq(parse_hosts(host), host)
+
+  @given(host=st.sampled_from(["localhost", "host", "192.168.1.1:8080"]), count=st.integers(0, 10))
+  def test_single_host_with_count(self, host:str, count:int):
+    self.assertEqual(parse_hosts(f"{host}*{count}"), [(host, i) for i in range(count)])
+
+  def test_multiple_hosts_with_counts_simple(self):
+    self.assertEqual(parse_hosts("host1*2,host2*3"), [("host1", i) for i in range(2)] + [("host2", i) for i in range(3)])
+
+  @given(st.lists(st.tuples(st.text(alphabet=string.ascii_letters + string.digits + ".-:"), st.integers(1, 16)), min_size=1))
+  def test_multiple_hosts_with_counts_sampled(self, host_count_pairs):
+    hosts_str = ",".join(f"{host}*{count}" for host, count in host_count_pairs)
+    expected = [(host, i) for host, count in host_count_pairs for i in range(count)]
+    self.assertEqual(parse_hosts(hosts_str), expected)
+
+  @given(st.sampled_from(["host1*2,host2", "a*1,b", "x*3,y*2,z"]))
+  def test_mixed_hosts_fails(self, hosts):
+    with self.assertRaises(AssertionError): parse_hosts(hosts)
+
+  @given(st.sampled_from(["host*abc", "test*xyz", "a*1.5"]))
+  def test_invalid_count_fails(self, hosts):
+    with self.assertRaises(ValueError): parse_hosts(hosts)
+
+  @given(st.sampled_from(["host*2*3", "a*1*2*3", "test*x*y"]))
+  def test_multiple_asterisks_fails(self, hosts):
+    with self.assertRaises(ValueError): parse_hosts(hosts)
+
+if __name__ == '__main__':
+  unittest.main()

test/test_schedule.py

@@ -15,7 +15,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
-from tinygrad.engine.grouper import view_left, view_right, sym, get_kernelize_map, Kernel, create_ast, merge_views, create_kernels
+from tinygrad.engine.kernelize import merge_views, get_kernelize_map, Kernel
 from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars
 from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule
 
@@ -66,8 +66,8 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float, **kwargs):
     np.testing.assert_allclose(img.grad.numpy(), ref_img.grad.detach().numpy(), atol=1e-6 if dtype == dtypes.float else 1e-2)
     np.testing.assert_allclose(w.grad.numpy(), ref_w.grad.detach().numpy(), atol=1e-6 if dtype == dtypes.float else 1e-2)
 
-@track_rewrites(named=True)
-def schedule_graph_rewrite(big_sink:UOp): return graph_rewrite(big_sink, merge_views+sym, {})
+@track_rewrites(name=True)
+def schedule_graph_rewrite(big_sink:UOp): return get_kernelize_map(big_sink)[big_sink]
 
 class TestSchedule(unittest.TestCase):
   def test_arange_avgpool2d(self, kcount=2):
@@ -182,20 +182,6 @@ class TestSchedule(unittest.TestCase):
     with Context(DONT_GROUP_REDUCES=1):
       check_schedule(x, 3, [Tensor._device_rng_counters[x.device]])
 
-  @unittest.skip("TODO: do not divide by zero given x.idiv(VALID)")
-  def test_rand_handcoded(self):
-    Tensor.manual_seed(0)
-    x = Tensor.rand(32)
-    # pre-realize shared seed
-    Tensor._device_rng_counters[x.device].realize()
-    # run custom kernelized kernel
-    sched_sink = graph_rewrite(x.uop, create_kernels, ctx={u:None for u in x.uop.toposort() if u.op is Ops.COPY}, bottom_up=True)
-    y = Tensor(graph_rewrite(sched_sink, create_ast, bottom_up=True))
-    run_schedule(check_schedule(y, 1))
-    # compare against reference
-    run_schedule(check_schedule(x, 3))
-    np.testing.assert_allclose(y.numpy(), x.numpy())
-
   def test_empty_is_not_realized(self):
     a = Tensor.empty(10)
     child = a+2
@@ -722,7 +708,6 @@ class TestSchedule(unittest.TestCase):
     c = Tensor.arange(4).realize().uop
     kernel = UOp(Ops.KERNEL, src=(a, b, c.base), arg=Kernel(UOp.sink(c.r(Ops.ADD, (0,))+1, c.r(Ops.ADD, (0,))*2)))
     assert all(s.op is Ops.BUFFER for s in kernel.src), f"views are not allowed here {kernel}"
-    kernel = graph_rewrite(kernel, create_ast)
     run_schedule(check_schedule(UOp.sink(a.assign(kernel), b.assign(kernel)), 1))
     self.assertEqual(a.buffer.numpy(), [7])
     self.assertEqual(b.buffer.numpy(), [12])
@@ -1999,7 +1984,7 @@ class TestIndexing(unittest.TestCase):
   def test_recursive_swizzle(self):
     a = Tensor([1,2,3,4]).realize()
     for _ in range(24): a = a + a
-    new_uop = swizzle_rewrite(a.uop.reshape((4, 1)))
+    new_uop = a.reshape(4,1).realize().uop
     self.assertEqual(new_uop.st, ShapeTracker.from_shape((4,)).reshape((4, 1)))
     self.assertEqual(swizzle_cnt(new_uop), 0)
 
@@ -2021,10 +2006,8 @@ class TestIndexing(unittest.TestCase):
     self.assertEqual(ast.shape, (32, 1))
     self.assertEqual(a.uop.shape, (32,))
 
-@track_rewrites(named=True)
-def swizzle_rewrite(u:UOp) -> UOp: return graph_rewrite(graph_rewrite(u, view_left), view_right)
 def swizzle_cnt(u:UOp) -> int:
-  return len([x for x in u.toposort() if x.op is Ops.VIEW and len(x.src) != 0 and x.src[0].op not in {Ops.BUFFER, Ops.DEFINE_GLOBAL}])
+  return len([x for x in u.toposort() if x.op is Ops.VIEW and len(x.src) != 0 and x.src[0].op not in {Ops.BUFFER, Ops.DEFINE_GLOBAL, Ops.ASSIGN}])
 
 class TestSwizzle(unittest.TestCase):
   def test_swizzle_simple(self):
@@ -2360,13 +2343,12 @@ class TestCopyFolding(unittest.TestCase):
     self.assertListEqual(b.tolist(), [0, 0, 0])
 
   def test_alu_after_copy(self):
-    a = Tensor.ones((4,)).to("CPU").uop
-    b = Tensor.empty(4, device="CPU").uop
+    a = Tensor.ones((4,)).to("CPU")
+    b = Tensor.empty(4, device="CPU")
     add = a+b
-    add = schedule_graph_rewrite(add)
-    assert all_same([x.device for x in add.src]), f"ALU has different devices! {[x.device for x in add.src]}"
+    add.kernelize()
+    assert all_same([x.device for x in add.uop.src]), f"ALU has different devices! {[x.device for x in add.src]}"
 
-  @unittest.skip("this is just clone now")
   def test_copy_to_same_device(self):
     a = Tensor.empty(4).uop
     b = a.copy_to_device(a.device)
@@ -2376,7 +2358,6 @@ class TestCopyFolding(unittest.TestCase):
     # in the scheduler because buffer already has shape (4,)
     self.assertIs(b, a.base)
 
-  @unittest.skip("this is just clone now")
   def test_copy_to_same_device_alt(self):
     a = Tensor.empty(4, 4).uop
     b = a.copy_to_device(a.device)
@@ -2437,29 +2418,6 @@ class TestCopyFolding(unittest.TestCase):
     b.realize()
     self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
 
-class TestTensorUOpSpec(unittest.TestCase):
-  def test_const_must_be_unmasked(self):
-    a = Tensor.ones((4, 4)).pad((2, 2))
-    unsafe_push_views = PatternMatcher([
-      (UPat.cvar("root").view(name="view"), lambda root,view: root.replace(src=tuple(x.view(view.st) for x in root.src))),
-    ])
-    a.uop = graph_rewrite(a.uop.sink(), merge_views+merge_views+unsafe_push_views)
-    with self.assertRaisesRegex(RuntimeError, "UOp verification failed"):
-      a.schedule()
-
-  def test_expanded_const_ok(self):
-    a = Tensor.ones((4, 4))
-    t = graph_rewrite(a.uop.sink(), merge_views+merge_views)
-    create_schedule_with_vars(t)
-
-  # NOTE: changing symbolic CONST VIEWs is not allowed
-  @unittest.expectedFailure
-  def test_symbolic_shape_ok(self):
-    a = Tensor.ones(4)
-    vi = UOp.variable("i", 1, 10).bind(4)
-    a.uop = graph_rewrite(a.reshape(vi).sum().uop, merge_views+merge_views)
-    a.schedule()
-
 class TestBufferUOp(unittest.TestCase):
   # BUFFER has a ShapeTracker of shape=(n,) and stride=(1,)
   def test_buffer_has_buffer(self):

test/test_search.py

@@ -1,10 +1,9 @@
 import unittest
 
-from tinygrad.codegen.kernel import Opt, OptOps
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Opt, OptOps, Kernel
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import bufs_from_lin, actions, beam_search
-from tinygrad.device import Device, Buffer
+from tinygrad.opt.search import bufs_from_lin, actions, beam_search
+from tinygrad.device import Device
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import Context, GlobalCounters
@@ -13,45 +12,6 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from extra.optimization.helpers import time_linearizer
 
-class TestTimeLinearizer(unittest.TestCase):
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WebGPU timestamps are low precision, tm is 0")
-  def test_reasonable_time(self):
-    a = Tensor([1,2,3,4]).realize()
-    si = (a+1).schedule()[0]
-    # create fresh empty buffers
-    rawbufs = [Buffer(b.device, b.size, b.dtype).allocate() for b in si.bufs]
-    tm = time_linearizer(Kernel(si.ast), rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
-    assert tm > 0 and tm != float('inf')
-
-  def test_bufs_from_lin(self):
-    a = Tensor([1,2,3,4]).realize()
-    si = (a+1).schedule()[0]
-    rawbufs = bufs_from_lin(lin:=Kernel(si.ast))
-    assert len(rawbufs) == len(lin.membufs) == 2
-    assert all(r is not None for r in rawbufs)
-    assert all(isinstance(r, Buffer) for r in rawbufs)
-    assert all(r.size > 0 for r in rawbufs)
-
-  def test_bufs_from_lin_alt(self):
-    a = Tensor.randn(4, 4).realize()
-    b = a+a[0]
-    si = b.schedule()[0]
-    rawbufs = bufs_from_lin(k:=Kernel(si.ast))
-    assert len(rawbufs) == len(k.membufs) == 2
-    assert all(r is not None for r in rawbufs)
-    assert all(isinstance(r, Buffer) for r in rawbufs)
-    assert all(r.size > 0 for r in rawbufs)
-
-  # Ensure that the kernel count is not incremented by time_linearizer when clearing l2
-  def test_kernel_count(self):
-    ast = Tensor.zeros(16).contiguous().kernelize().uop.src[1].arg.ast
-    lin = Kernel(ast)
-    bufs = bufs_from_lin(lin)
-
-    kernel_count = GlobalCounters.kernel_count
-    time_linearizer(lin, bufs, allow_test_size=False, cnt=2, disable_cache=True, clear_l2=True)
-    assert GlobalCounters.kernel_count == kernel_count, "kernel count was incremented by time_linearizer"
-
 class TestBEAM(unittest.TestCase):
   def test_dynamic_beam(self):
     # TODO: make this infra globally usable
@@ -79,7 +39,7 @@ class TestBEAM(unittest.TestCase):
     a = Tensor.rand(4, 3)
     b = Tensor.rand(3)
     realized_ast, _ = helper_realized_ast(a @ b)
-    from tinygrad.engine.search import get_kernel_actions
+    from tinygrad.opt.search import get_kernel_actions
     lins = get_kernel_actions(Kernel(realized_ast), False).values()
 
     # ensure amt=0 are not duplicated
@@ -97,7 +57,7 @@ class TestBEAM(unittest.TestCase):
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
   def test_search_over_shape(self):
     from test.test_linearizer import helper_realized_ast
-    from tinygrad.engine.search import get_kernel_actions
+    from tinygrad.opt.search import get_kernel_actions
 
     dtype_pairs = [(tc.dtype_in, tc.dtype_out) for tc in Device[Device.DEFAULT].renderer.tensor_cores]
     multi_shape_dtype_pairs = [dts for dts in dtype_pairs if dtype_pairs.count(dts) > 1]
@@ -114,7 +74,7 @@ class TestBEAM(unittest.TestCase):
 
   def test_get_kernel_actions_preserves_actions_state(self):
     from test.test_linearizer import helper_realized_ast
-    from tinygrad.engine.search import get_kernel_actions
+    from tinygrad.opt.search import get_kernel_actions
     a = Tensor.rand(16, 16)
     b = Tensor.rand(16, 16)
     realized_ast, _ = helper_realized_ast(a @ b)

test/test_uops.py

@@ -4,18 +4,18 @@ import numpy as np
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View # noqa F401
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.helpers import CI, DEBUG, getenv, Context, Timing
+from tinygrad.helpers import CI, DEBUG, getenv, Timing
 from tinygrad.dtype import dtypes, DType
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu # noqa F401
 from tinygrad.uop.spec import spec
 from tinygrad.renderer import ProgramSpec
-from tinygrad.engine.grouper import fix_kernel_ops
-from tinygrad.engine.realize import CompiledRunner, get_kernel
+from tinygrad.engine.kernelize import fix_kernel_ops
+from tinygrad.engine.realize import CompiledRunner
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 
 def to_uops_list(u:list[UOp], opts=None, skip_check=False) -> list[UOp]: return full_rewrite(UOp.sink(*u), opts)
 
@@ -361,15 +361,16 @@ class TestAssembly(unittest.TestCase):
     self.assertIn(Ops.MUL, ops)
 
   def test_division_power_of_two(self):
-    g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
-    c = UOp(Ops.CONST, dtypes.uint, (), 2)
-    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
-    a = UOp(Ops.IDIV, dtypes.uint, (l, c))
-    uops = to_uops_list([a], opts=Device[Device.DEFAULT].renderer)
-    Device[Device.DEFAULT].renderer.render(uops)
-    ops = [x.op for x in uops]
-    self.assertIn(Ops.SHR, ops)
-    self.assertNotIn(Ops.IDIV, ops)
+    for dt in (dtypes.int32, dtypes.uint32):
+      g = UOp(Ops.DEFINE_GLOBAL, dt.ptr(), (), 0)
+      c = UOp(Ops.CONST, dt, (), 2)
+      l = UOp(Ops.LOAD, dt, (g.index(c),))
+      a = UOp(Ops.IDIV, dt, (l, c))
+      uops = to_uops_list([a], opts=Device[Device.DEFAULT].renderer)
+      Device[Device.DEFAULT].renderer.render(uops)
+      ops = [x.op for x in uops]
+      self.assertIn(Ops.SHR, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
+      self.assertNotIn(Ops.IDIV, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
 
   def test_fast_idiv_and_mod(self):
     g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
@@ -460,13 +461,6 @@ class TestUOpStr(unittest.TestCase):
     assert len(str(a)) < 10_000, "exponential string growth"
     assert str(eval(str(a))) == str(a)
 
-    t = Tensor.arange(10)
-    t = t + t * Tensor.rand(10)
-    # nice big complicated uop
-    with Context(NOOPT=1):
-      sink = UOp(Ops.SINK, dtypes.void, (get_kernel(Device[Device.DEFAULT].renderer, t.schedule()[-1].ast).linearize().uops[-1],))
-    self.assertEqual(sink, eval(str(sink)))
-
   def test_vectorized_str(self):
     vec = UOp(Ops.VECTORIZE, dtypes.int.vec(4), tuple(UOp.const(dtypes.int, x) for x in range(4)))
     assert str(eval(str(vec))) == str(vec)
@@ -522,7 +516,7 @@ class TestIndexingOrdering(unittest.TestCase):
 class TestUPatHelpers(unittest.TestCase):
   def test_location(self):
     self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "symbolic.py")
-    self.assertEqual(fix_kernel_ops.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "grouper.py")
+    self.assertEqual(fix_kernel_ops.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "kernelize.py")
     self.assertEqual(spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
     test_upat = UPat(Ops.CONST, dtypes.bool)
     self.assertEqual(test_upat.location[0].split("/")[-1], __file__.replace("\\", "/").split("/")[-1])

test/test_uops_stats.py

@@ -6,7 +6,7 @@ from tinygrad.renderer import Estimates
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps, KernelOptError
+from tinygrad.opt.kernel import Kernel, Opt, OptOps, KernelOptError
 from tinygrad.device import Device
 
 def flops_mem(uops, ignore_indexing=False):

test/test_winograd.py

@@ -3,8 +3,8 @@ import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, Context, nn
 from tinygrad.uop.ops import Ops
 from tinygrad.helpers import Timing, CI, Profiling, WINO, DEBUG, getenv
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 
 class TestWinogradClose(unittest.TestCase):
   def test_close(self):

test/unit/test_block_reorder.py

@@ -54,7 +54,7 @@ class TestBlockReorder(unittest.TestCase):
     sink = c.store(sum(loads)).sink()
 
     # determine golden order
-    golden = block_reorder(sink.toposort())
+    golden = block_reorder(list(sink.toposort()))
 
     # render for test
     print(self._test_render(golden))

test/unit/test_dtype_spec.py

@@ -0,0 +1,552 @@
+import unittest, math, operator, subprocess
+from tinygrad.tensor import Tensor, dtypes, Device
+from tinygrad.dtype import DType, DTYPES_DICT, truncate, truncate_fp16, truncate_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
+from tinygrad.device import is_dtype_supported
+from tinygrad.helpers import getenv, CI, DEBUG
+from hypothesis import given, settings, strategies as strat
+import numpy as np
+import torch
+import ml_dtypes
+
+settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+settings.load_profile("my_profile")
+
+core_dtypes = list(DTYPES_DICT.values())
+dtype_ints = [dt for dt in core_dtypes if dtypes.is_int(dt) and is_dtype_supported(dt)]
+dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and is_dtype_supported(dt)]
+
+FP8E4M3_MAX = 448.0
+FP8E5M2_MAX = 57344.0
+
+def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float=1e-7):
+  if DEBUG >= 2: print(tensor.numpy())
+  try:
+    assert tensor.dtype == target_dtype
+    np.testing.assert_allclose(tensor.numpy(), target, rtol={dtypes.float16:1e-3, dtypes.bfloat16:1e-2}.get(target_dtype, tol_target_dtype))
+  except AssertionError as e:
+    raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e
+
+class TestHelpers(unittest.TestCase):
+  signed_ints = (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64)
+  uints = (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
+  floats = (dtypes.float16, dtypes.float32, dtypes.float64)
+
+  @given(strat.sampled_from(signed_ints+uints), strat.integers(min_value=1, max_value=8))
+  def test_is_int(self, dtype, amt):
+    assert dtypes.is_int(dtype.vec(amt) if amt > 1 else dtype)
+    assert not dtypes.is_float(dtype.vec(amt) if amt > 1 else dtype)
+
+  @given(strat.sampled_from(uints), strat.integers(min_value=1, max_value=8))
+  def test_is_unsigned_uints(self, dtype, amt):
+    assert dtypes.is_unsigned(dtype.vec(amt) if amt > 1 else dtype)
+
+  @given(strat.sampled_from(signed_ints), strat.integers(min_value=1, max_value=8))
+  def test_is_unsigned_signed_ints(self, dtype, amt):
+    assert not dtypes.is_unsigned(dtype.vec(amt) if amt > 1 else dtype)
+
+  @given(strat.sampled_from(floats), strat.integers(min_value=1, max_value=8))
+  def test_is_float(self, dtype, amt):
+    assert dtypes.is_float(dtype.vec(amt) if amt > 1 else dtype)
+    assert not dtypes.is_int(dtype.vec(amt) if amt > 1 else dtype)
+    assert not dtypes.is_unsigned(dtype.vec(amt) if amt > 1 else dtype)
+
+  def test_bf16_is_float(self):
+    assert dtypes.is_float(dtypes.bfloat16)
+
+  def test_fp8s_are_float(self):
+    assert dtypes.is_float(dtypes.fp8e4m3)
+    assert dtypes.is_float(dtypes.fp8e5m2)
+
+  @given(strat.sampled_from([d for d in DTYPES_DICT.values() if dtypes.is_float(d) or dtypes.is_int(d)]), strat.integers(min_value=2, max_value=8))
+  def test_scalar(self, dtype, amt):
+    assert dtype.vec(amt).scalar() == dtype
+
+  def test_from_py(self):
+    assert dtypes.from_py(True) == dtypes.bool
+    assert dtypes.from_py(2) == dtypes.default_int
+    assert dtypes.from_py(3.0) == dtypes.default_float
+    assert dtypes.from_py([]) == dtypes.default_float
+    assert dtypes.from_py(()) == dtypes.default_float
+    assert dtypes.from_py([True]) == dtypes.bool
+    assert dtypes.from_py([True, 2]) == dtypes.default_int
+    assert dtypes.from_py([True, 3.0]) == dtypes.default_float
+    assert dtypes.from_py([2, 3.0]) == dtypes.default_float
+    assert dtypes.from_py([True, 2, 3.0]) == dtypes.default_float
+    with self.assertRaises(RuntimeError): dtypes.from_py(None)
+    with self.assertRaises(RuntimeError): dtypes.from_py([None])
+    with self.assertRaises(RuntimeError): dtypes.from_py({})
+    with self.assertRaises(RuntimeError): dtypes.from_py(set())
+
+  def test_dtype_range(self):
+    for dt in core_dtypes:
+      if dtypes.is_float(dt):
+        np.testing.assert_equal(dtypes.min(dt), -math.inf)
+        np.testing.assert_equal(dtypes.max(dt), math.inf)
+        np.testing.assert_equal(dt.min, -math.inf)
+        np.testing.assert_equal(dt.max, math.inf)
+      elif dtypes.is_int(dt):
+        info = np.iinfo(_to_np_dtype(dt))
+        np.testing.assert_equal(dtypes.min(dt), info.min)
+        np.testing.assert_equal(dtypes.max(dt), info.max)
+        np.testing.assert_equal(dt.min, info.min)
+        np.testing.assert_equal(dt.max, info.max)
+      else:
+        assert dt == dtypes.bool, dt
+        np.testing.assert_equal(dtypes.min(dt), False)
+        np.testing.assert_equal(dtypes.max(dt), True)
+        np.testing.assert_equal(dt.min, False)
+        np.testing.assert_equal(dt.max, True)
+
+  def test_truncate_fp16(self):
+    self.assertEqual(truncate_fp16(1), 1)
+    self.assertEqual(truncate_fp16(65504), 65504)
+    self.assertEqual(truncate_fp16(65519.999), 65504)
+    self.assertEqual(truncate_fp16(65520), math.inf)
+
+  def test_truncate_bf16(self):
+    self.assertEqual(truncate_bf16(1), 1)
+    self.assertAlmostEqual(truncate_bf16(1.1), 1.09375, places=7)
+    for a in [1234, 23456, -777.777]:
+      self.assertEqual(truncate_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
+    # TODO: torch bfloat 1.1 gives 1.1015625 instead of 1.09375
+    max_bf16 = torch.finfo(torch.bfloat16).max
+    self.assertEqual(truncate_bf16(max_bf16), max_bf16)
+    self.assertEqual(truncate_bf16(min_bf16:=-max_bf16), min_bf16)
+    self.assertEqual(truncate_bf16(max_bf16 * 1.00001), math.inf)
+    self.assertEqual(truncate_bf16(min_bf16 * 1.00001), -math.inf)
+
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
+  def test_truncate_fp8e4m3(self, x):
+    if x > FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), FP8E4M3_MAX)
+    elif x < -FP8E4M3_MAX: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), -FP8E4M3_MAX)
+    else: np.testing.assert_equal(truncate[dtypes.fp8e4m3](x), ml_dtypes.float8_e4m3fn(x))
+
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=True, allow_infinity=True))
+  def test_truncate_fp8e5m2(self, x):
+    if x > FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), FP8E5M2_MAX)
+    elif x < -FP8E5M2_MAX: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), -FP8E5M2_MAX)
+    else: np.testing.assert_equal(truncate[dtypes.fp8e5m2](x), ml_dtypes.float8_e5m2(x))
+
+class TestTypeSpec(unittest.TestCase):
+  def setUp(self):
+    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
+  def tearDown(self):
+    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
+
+  def test_set_dtype_default(self):
+    for default_int in [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64]:
+      dtypes.default_int = default_int
+      assert dtypes.default_int == default_int
+
+    for default_float in [*dtypes.fp8s, dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
+      dtypes.default_float = default_float
+      assert dtypes.default_float == default_float
+
+  @unittest.skip("this test is slow and spawning whole pythons")
+  def test_env_set_default_float(self):
+    # check default
+    subprocess.run(['python3 -c "from tinygrad import dtypes; assert dtypes.default_float == dtypes.float"'],
+                    shell=True, check=True)
+    # check change
+    subprocess.run(['DEFAULT_FLOAT=HALF python3 -c "from tinygrad import dtypes; assert dtypes.default_float == dtypes.half"'],
+                    shell=True, check=True)
+    # check invalid
+    with self.assertRaises(subprocess.CalledProcessError):
+      subprocess.run(['DEFAULT_FLOAT=INT32 python3 -c "from tinygrad import dtypes"'],
+                      shell=True, check=True)
+
+    with self.assertRaises(subprocess.CalledProcessError):
+      subprocess.run(['DEFAULT_FLOAT=TYPO python3 -c "from tinygrad import dtypes"'],
+                      shell=True, check=True)
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.int8), f"no int8 on {Device.DEFAULT}")
+  def test_dtype_str_arg(self):
+    n = np.random.normal(0, 1, (10, 10)).astype(np.float32)
+    tested = 0
+    for dtype_str, dtype in [
+      ("bool", dtypes.bool), ("int8", dtypes.int8), ("int", dtypes.int), ("uint32", dtypes.uint32), ("float32", dtypes.float32)]:
+      np.testing.assert_equal(Tensor(n, dtype=dtype_str).numpy(), Tensor(n, dtype=dtype).numpy())
+      np.testing.assert_equal(Tensor(n).cast(dtype_str).numpy(), Tensor(n).cast(dtype).numpy())
+      if dtype.itemsize == 4:
+        np.testing.assert_equal(Tensor(n).bitcast(dtype_str).numpy(), Tensor(n).bitcast(dtype).numpy())
+        tested += 1
+    assert tested == 3
+
+    with self.assertRaises(AttributeError): Tensor([1, 2, 3], dtype="nonexistdtype")
+    with self.assertRaises(AttributeError): Tensor([1, 2, 3], dtype="")
+
+    np.testing.assert_equal(Tensor(n).sum(dtype="int16").numpy(), Tensor(n).sum(dtype=dtypes.int16).numpy())
+
+  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
+  def test_creation(self, default_int, default_float):
+    dtypes.default_int, dtypes.default_float = default_int, default_float
+    _assert_eq(Tensor(True), dtypes.bool, True)
+    _assert_eq(Tensor(None), dtypes.default_float, [])
+    _assert_eq(Tensor(2), dtypes.default_int, 2)
+    _assert_eq(Tensor(2.34), dtypes.default_float, 2.34)
+    _assert_eq(Tensor([]), dtypes.default_float, [])
+    _assert_eq(Tensor([1]), dtypes.default_int, [1])
+    _assert_eq(Tensor([1.1]), dtypes.default_float, [1.1])
+
+    _assert_eq(Tensor.eye(0), dtypes.default_float, np.eye(0))
+    _assert_eq(Tensor.eye(3), dtypes.default_float, np.eye(3))
+    if is_dtype_supported(dtypes.int64):
+      _assert_eq(Tensor.eye(3, dtype=dtypes.int64), dtypes.int64, np.eye(3))
+    if is_dtype_supported(dtypes.float16):
+      _assert_eq(Tensor.eye(3, dtype=dtypes.float16), dtypes.float16, np.eye(3))
+
+  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
+  def test_full(self, default_int, default_float):
+    dtypes.default_int, dtypes.default_float = default_int, default_float
+
+    _assert_eq(Tensor.zeros((2, 3)), dtypes.default_float, np.zeros((2, 3)))
+    if is_dtype_supported(dtypes.int64):
+      _assert_eq(Tensor.zeros((2, 3), dtype=dtypes.int64), dtypes.int64, np.zeros((2, 3)))
+    if is_dtype_supported(dtypes.float16):
+      _assert_eq(Tensor.zeros((2, 3), dtype=dtypes.float16), dtypes.float16, np.zeros((2, 3)))
+
+    _assert_eq(Tensor.ones((2, 3)), dtypes.default_float, np.ones((2, 3)))
+    if is_dtype_supported(dtypes.int64):
+      _assert_eq(Tensor.ones((2, 3), dtype=dtypes.int64), dtypes.int64, np.ones((2, 3)))
+    if is_dtype_supported(dtypes.float16):
+      _assert_eq(Tensor.ones((2, 3), dtype=dtypes.float16), dtypes.float16, np.ones((2, 3)))
+
+    _assert_eq(Tensor.full((2, 3), 3.0), dtypes.default_float, np.full((2, 3), 3.0))
+    _assert_eq(Tensor.full((2, 3), 3), dtypes.default_int, np.full((2, 3), 3))
+    _assert_eq(Tensor.full((2, 3), True), dtypes.bool, np.full((2, 3), True))
+    if is_dtype_supported(dtypes.int64):
+      _assert_eq(Tensor.full((2, 3), 3, dtype=dtypes.int64), dtypes.int64, np.full((2, 3), 3))
+      _assert_eq(Tensor.full((2, 3), 3.0, dtype=dtypes.int64), dtypes.int64, np.full((2, 3), 3))
+    if is_dtype_supported(dtypes.float16):
+      _assert_eq(Tensor.full((2, 3), 3, dtype=dtypes.float16), dtypes.float16, np.full((2, 3), 3))
+      _assert_eq(Tensor.full((2, 3), 3.0, dtype=dtypes.float16), dtypes.float16, np.full((2, 3), 3))
+
+  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
+  def test_reduce_0d_default(self, default_int, default_float):
+    dtypes.default_int, dtypes.default_float = default_int, default_float
+    _assert_eq(Tensor.ones((2,3,0)).sum(2), dtypes.default_float, np.zeros((2, 3)))
+    # TODO: what should this one be?
+    # _assert_eq(Tensor.ones((2,3,0), dtype=dtypes.default_int).sum(2), dtypes.default_int, np.zeros((2, 3)))
+    _assert_eq(Tensor.ones((2,3,0), dtype=dtypes.int32).sum(2), dtypes.int32, np.zeros((2, 3)))
+
+  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
+  def test_arange(self, default_int, default_float):
+    dtypes.default_int, dtypes.default_float = default_int, default_float
+
+    _assert_eq(Tensor.arange(5), dtypes.default_int, np.arange(5))
+    _assert_eq(Tensor.arange(120), dtypes.default_int, np.arange(120))
+    _assert_eq(Tensor.arange(5.0), dtypes.default_float, np.arange(5))
+    if is_dtype_supported(dtypes.int16):
+      _assert_eq(Tensor.arange(5, dtype=dtypes.int16), dtypes.int16, np.arange(5))
+    if is_dtype_supported(dtypes.int64):
+      _assert_eq(Tensor.arange(5, dtype=dtypes.int64), dtypes.int64, np.arange(5))
+    if is_dtype_supported(dtypes.float16):
+      _assert_eq(Tensor.arange(5, dtype=dtypes.float16), dtypes.float16, np.arange(5))
+    _assert_eq(Tensor.arange(3, 9, 0.7), dtypes.default_float, np.arange(3, 9, 0.7), 1e-6 if Device.DEFAULT == "WEBGPU" else 1e-7)
+    _assert_eq(Tensor.arange(3, 8.5, 3), dtypes.default_float, np.arange(3, 8.5, 3))
+    # stop-start and step have different signs
+    _assert_eq(Tensor.arange(3, 5, -2), dtypes.default_int, np.arange(3, 5, -2))
+    _assert_eq(Tensor.arange(5.0, 3.0), dtypes.default_float, np.arange(5.0, 3.0))
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from([operator.gt, operator.ge, operator.le, operator.lt, operator.eq, operator.ne]))
+  def test_bool_ops(self, dtype, op):
+    assert op(Tensor.ones(4, 4, dtype=dtype), Tensor.ones(4, 4, dtype=dtype)).dtype == dtypes.bool
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
+  def test_functions_return_index(self, dtype, default_int, default_float):
+    dtypes.default_int, dtypes.default_float = default_int, default_float
+    assert Tensor([0, 1], dtype=dtype).argmax().dtype == dtypes.int32
+    assert Tensor([0, 1], dtype=dtype).argmin().dtype == dtypes.int32
+    assert Tensor([0, 1], dtype=dtype).multinomial().dtype == dtypes.int32
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints))
+  def test_tensor_indexing_returns_same_dtype(self, data_dtype, indices_dtype):
+    X_data =  Tensor.ones(60000, 1, 28, 28, dtype=data_dtype)
+    indices =  Tensor.randint(512, high=X_data.shape[0]).cast(indices_dtype)
+    assert X_data[indices].dtype == X_data.dtype
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints))
+  def test_gather_returns_same_dtype(self, data_dtype, indices_dtype):
+    X_data = Tensor([[1, 0], [0, 1]], dtype=data_dtype)
+    indices = Tensor([[0, 0], [1, 0]], dtype=indices_dtype)
+    assert X_data.gather(0, indices).dtype == X_data.dtype
+    assert X_data.gather(1, indices).dtype == X_data.dtype
+
+  @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
+  def test_attention_returns_same_dtype(self, data_dtype, default_float):
+    dtypes.default_float = default_float
+    query = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
+    key = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
+    value = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
+    mask = (Tensor.rand(32, 8, 128, 128) < 0.5)
+    assert query.scaled_dot_product_attention(key, value, is_causal=True).dtype == data_dtype
+    assert query.scaled_dot_product_attention(key, value, is_causal=True, dropout_p=0.3).dtype == data_dtype
+    assert query.scaled_dot_product_attention(key, value, is_causal=False).dtype == data_dtype
+    assert query.scaled_dot_product_attention(key, value, attn_mask=mask).dtype == data_dtype
+
+class TestTypePromotion(unittest.TestCase):
+  @given(strat.sampled_from(core_dtypes))
+  def test_self_promo_to_self(self, dtype):
+    assert least_upper_dtype(dtype) == dtype
+    assert least_upper_dtype(dtype, dtype) == dtype
+    assert least_upper_dtype(dtype, dtype, dtype) == dtype
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_promo_resulted_higher_than_inputs(self, dtype1, dtype2):
+    result = least_upper_dtype(dtype1, dtype2)
+    assert not (result < dtype1) and not (result < dtype2)
+
+  def test_dtype_promo(self):
+    assert least_upper_dtype(dtypes.bool, dtypes.int8) == dtypes.int8
+    assert least_upper_dtype(dtypes.int8, dtypes.uint8) == dtypes.int16
+    assert least_upper_dtype(dtypes.uint8, dtypes.int16) == dtypes.int16
+    assert least_upper_dtype(dtypes.int16, dtypes.uint16) == dtypes.int32
+    assert least_upper_dtype(dtypes.uint16, dtypes.int32) == dtypes.int32
+    assert least_upper_dtype(dtypes.int32, dtypes.uint32) == dtypes.int64
+    assert least_upper_dtype(dtypes.uint32, dtypes.int64) == dtypes.int64
+    # similar to jax but we don't use weak type
+    assert least_upper_dtype(dtypes.int64, dtypes.uint64) == dtypes.float16
+    assert least_upper_dtype(dtypes.float16, dtypes.float32) == dtypes.float32
+    assert least_upper_dtype(dtypes.float32, dtypes.float64) == dtypes.float64
+
+    assert least_upper_dtype(dtypes.bool, dtypes.float32) == dtypes.float32
+    assert least_upper_dtype(dtypes.bool, dtypes.float64) == dtypes.float64
+    assert least_upper_dtype(dtypes.float16, dtypes.int64) == dtypes.float16
+    assert least_upper_dtype(dtypes.float16, dtypes.uint64) == dtypes.float16
+    assert least_upper_dtype(dtypes.fp8e4m3, dtypes.fp8e5m2) == dtypes.half
+
+class TestAutoCastType(unittest.TestCase):
+  def setUp(self):
+    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
+  def tearDown(self):
+    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
+
+  @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
+  def test_least_upper_float_input_is_float(self, input_dtype, default_float):
+    dtypes.default_float = default_float
+    self.assertEqual(least_upper_float(input_dtype), input_dtype)
+
+  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
+  def test_least_upper_float_input_is_int(self, input_dtype, default_float):
+    dtypes.default_float = default_float
+    self.assertEqual(least_upper_float(input_dtype), default_float)
+
+  @given(strat.sampled_from([d for d in core_dtypes if dtypes.is_int(d) and is_dtype_supported(d)]))
+  def test_int_to_float_unary_func(self, dtype):
+    for func in [
+      lambda t: t.exp(),
+      lambda t: t.exp2(),
+      lambda t: t.log(),
+      lambda t: t.log2(),
+      lambda t: t.sqrt(),
+      lambda t: t.rsqrt(),
+      lambda t: t.sin(),
+      lambda t: t.cos(),
+      lambda t: t.tan(),
+      lambda t: t.sigmoid(),
+    ]:
+      a = [2, 3, 4]
+      # float16 can have larger precision errors
+      np.testing.assert_allclose(func(Tensor(a, dtype=dtype)).numpy(), func(torch.tensor(a)), rtol=1e-3, atol=1e-3)
+
+  @given(strat.sampled_from(core_dtypes))
+  def test_broadcast_scalar(self, dt):
+    assert (Tensor.ones(4, 4, dtype=dt) + 2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
+    assert (Tensor.ones(4, 4, dtype=dt) + 2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
+    assert (Tensor.ones(4, 4, dtype=dt) + True).dtype == dt
+
+  @given(strat.sampled_from(dtype_floats))
+  def test_int_div_int(self, default_float):
+    dtypes.default_float = default_float
+    self.assertEqual(Tensor([1]).div(Tensor([2])).dtype, default_float)
+
+  def test_sum(self):
+    assert (Tensor([0, 1], dtype=dtypes.bool)).sum().dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int8)).sum().dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int16)).sum().dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int32)).sum().dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int64)).sum().dtype == dtypes.int64
+    assert (Tensor([0, 1], dtype=dtypes.uint8)).sum().dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint16)).sum().dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint32)).sum().dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint64)).sum().dtype == dtypes.uint64
+    assert (Tensor([0, 1], dtype=dtypes.float16)).sum().dtype == dtypes.float16
+    #assert (Tensor([0, 1], dtype=dtypes.bfloat16)).sum().dtype == dtypes.bfloat16
+    assert (Tensor([0, 1], dtype=dtypes.float32)).sum().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.float64)).sum().dtype == dtypes.float64
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16")
+  def test_sum_dtype_arg(self):
+    t = Tensor([40000, 40000], dtype=dtypes.float16)
+    # default float16 sum returns in float16, overflowed in this case
+    assert t.sum().dtype == dtypes.float16
+    assert math.isinf(t.sum().numpy().item())
+    # specifiying dtype and it's not downcasted
+    assert t.sum(dtype=dtypes.float32).dtype == dtypes.float32
+    np.testing.assert_allclose(t.sum(dtype=dtypes.float32).numpy(), 80000)
+
+  def test_prod_dtype_arg(self):
+    t = Tensor([100, 200], dtype=dtypes.int32)
+    assert t.prod().dtype == dtypes.int32
+    np.testing.assert_allclose(t.prod().numpy(), 20000)
+    assert t.prod(dtype=dtypes.float32).dtype == dtypes.float32
+    np.testing.assert_allclose(t.prod(dtype=dtypes.float32).numpy(), 20000)
+
+  def test_mean(self):
+    assert (Tensor([0, 1], dtype=dtypes.bool)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.int8)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.int16)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.int32)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.int64)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.uint8)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.uint16)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.uint32)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.uint64)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.float16)).mean().dtype == dtypes.float16
+    #assert (Tensor([0, 1], dtype=dtypes.bfloat16)).mean().dtype == dtypes.bfloat16
+    assert (Tensor([0, 1], dtype=dtypes.float32)).mean().dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.float64)).mean().dtype == dtypes.float64
+
+  def test_cumsum(self):
+    assert (Tensor([0, 1], dtype=dtypes.bool)).cumsum(0).dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int8)).cumsum(0).dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int16)).cumsum(0).dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int32)).cumsum(0).dtype == dtypes.int32
+    assert (Tensor([0, 1], dtype=dtypes.int64)).cumsum(0).dtype == dtypes.int64
+    assert (Tensor([0, 1], dtype=dtypes.uint8)).cumsum(0).dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint16)).cumsum(0).dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint32)).cumsum(0).dtype == dtypes.uint32
+    assert (Tensor([0, 1], dtype=dtypes.uint64)).cumsum(0).dtype == dtypes.uint64
+    assert (Tensor([0, 1], dtype=dtypes.float16)).cumsum(0).dtype == dtypes.float16
+    #assert (Tensor([0, 1], dtype=dtypes.bfloat16)).cumsum(0).dtype == dtypes.bfloat16
+    assert (Tensor([0, 1], dtype=dtypes.float32)).cumsum(0).dtype == dtypes.float32
+    assert (Tensor([0, 1], dtype=dtypes.float64)).cumsum(0).dtype == dtypes.float64
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_matmul(self, dt1, dt2, acc_dt):
+    t1 = Tensor([0, 1], dtype=dt1)
+    t2 = Tensor([0, 1], dtype=dt2)
+    self.assertEqual(t1.matmul(t2).dtype, least_upper_dtype(t1.dtype, t2.dtype))
+    # if dtype is specified, return in dtype
+    self.assertEqual(t1.matmul(t2, dtype=acc_dt).dtype, acc_dt)
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_linear(self, dt1, dt2, dt3, acc_dt):
+    x = Tensor([0, 1], dtype=dt1)
+    w = Tensor([0, 1], dtype=dt2)
+    b = Tensor([0, 1], dtype=dt3)
+    self.assertEqual(x.linear(w).dtype, least_upper_dtype(x.dtype, w.dtype))
+    self.assertEqual(x.linear(w, b).dtype, least_upper_dtype(least_upper_dtype(x.dtype, w.dtype), b.dtype))
+    # if dtype is specified, return in dtype
+    self.assertEqual(x.linear(w, dtype=acc_dt).dtype, acc_dt)
+    self.assertEqual(x.linear(w, b, dtype=acc_dt).dtype, acc_dt)
+
+  @staticmethod
+  def check_where_alternate_input_other(input_, other, data_type):
+    assert (Tensor([True, False]).where(input_, other)).dtype == data_type
+    assert (Tensor([True, False]).where(other, input_)).dtype == data_type
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_where_no_scalar(self, dt1, dt2):
+    self.check_where_alternate_input_other(Tensor(2, dtype=dt1), Tensor(3, dtype=dt2), least_upper_dtype(dt1, dt2))
+
+  @given(strat.sampled_from(core_dtypes))
+  def test_where_one_scalar(self, dt):
+    t = Tensor(2, dtype=dt)
+    self.check_where_alternate_input_other(t, 3.2, (dt if dtypes.is_float(dt) else dtypes.default_float))
+    self.check_where_alternate_input_other(t, 3, (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int))
+    self.check_where_alternate_input_other(t, True, dt)
+
+  def test_where_two_scalars(self):
+    self.check_where_alternate_input_other(3.1, 3.2, dtypes.default_float)
+    self.check_where_alternate_input_other(3.1, 3, dtypes.default_float)
+    self.check_where_alternate_input_other(3.1, True, dtypes.default_float)
+    self.check_where_alternate_input_other(3, 2, dtypes.default_int)
+    self.check_where_alternate_input_other(3, True, dtypes.default_int)
+    self.check_where_alternate_input_other(False, True, dtypes.bool)
+
+  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
+  def test_maximum(self, dt1, dt2):
+    assert Tensor([0, 1, 2], dtype=dt1).maximum(Tensor([2, 0, 5], dtype=dt2)).dtype == least_upper_dtype(dt1, dt2)
+
+  @given(strat.sampled_from(core_dtypes))
+  def test_maximum_const(self, dt):
+    assert Tensor([1, 2], dtype=dt).maximum(3.1).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
+    assert Tensor([1, 2], dtype=dt).maximum(3).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
+    assert Tensor([1, 2], dtype=dt).maximum(True).dtype == dt
+
+  def test_div(self):
+    assert (Tensor([1, 2], dtype=dtypes.int32) / Tensor([2, 2], dtype=dtypes.int32)).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.int16) / Tensor([2, 2], dtype=dtypes.int32)).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.float32) / Tensor([2, 2], dtype=dtypes.float16)).dtype == dtypes.float32
+    assert (Tensor([1, 2], dtype=dtypes.int32) / Tensor([2, 2], dtype=dtypes.float16)).dtype == dtypes.float16
+
+  def test_div_const(self):
+    assert (Tensor([1, 2], dtype=dtypes.int32) / 2).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.int32) / 2.0).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.float16) / 2).dtype == dtypes.float16
+    assert (Tensor([1, 2], dtype=dtypes.float16) / 2.0).dtype == dtypes.float16
+
+  def test_gradient_dtype(self):
+    old_default_float = dtypes.default_float
+
+    for default_dtype in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
+      if not is_dtype_supported(default_dtype): continue
+      dtypes.default_float = default_dtype
+      for dtype in [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
+        if not is_dtype_supported(dtype): continue
+        if DEBUG >= 2:
+          print(f"testing {default_dtype=}, {dtype=}")
+        a = Tensor([1, 2, 3], dtype=dtype, requires_grad=True)
+        b = (a * 5).sum()
+        b.backward()  # if there is dtype mismatch, lazy should assert
+        assert a.grad.dtype == a.dtype
+        np.testing.assert_allclose(a.grad.numpy(), [5, 5, 5])
+
+    dtypes.default_float = old_default_float
+
+  @unittest.skipIf(CI, "TODO: broken RuntimeError: Attempting to relocate against an undefined symbol 'fmaxf'")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
+  def test_backward_sum_acc_dtype(self):
+    # test acc of sum in the backward is upcasted to float
+    t = Tensor([5, -5], dtype=dtypes.half, requires_grad=True)
+    t.reshape(2, 1).expand(2, 10001).max().backward()
+    np.testing.assert_allclose(t.grad.numpy(), [1, 0])
+
+  @unittest.skipIf(Device.DEFAULT == "PYTHON", "very slow")
+  @unittest.skipIf(CI and Device.DEFAULT == "AMD", "very slow")
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Binding size is larger than the maximum storage buffer binding size")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
+  def test_mean_half_precision_underflow(self):
+    N = 10000
+    x = 0.001
+    t = Tensor([[x]], dtype=dtypes.half, requires_grad=True).expand(N, N).contiguous()
+    np.testing.assert_allclose(t.mean(axis=1).numpy(), np.array([x] * N, dtype=np.float16), rtol=1e-3)
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
+  def test_mean_half_precision_overflow(self):
+    N = 256
+    t = Tensor([60000] * N*N, dtype=dtypes.half, requires_grad=True).reshape(N, N)
+    np.testing.assert_allclose(t.mean().numpy(), 60000)
+    t.square().mean().backward()
+    np.testing.assert_allclose(t.grad.numpy().flatten(), [60000 * 2 / (N*N)] * N*N)
+
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Precision error")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
+  def test_softmax_dtype(self):
+    data = [1, 2, 3]
+    t = Tensor(data, dtype=dtypes.half)
+    tt = torch.tensor(data, dtype=torch.half)
+
+    out = t.softmax(0)
+    self.assertEqual(out.dtype, dtypes.half)
+    np.testing.assert_allclose(out.numpy(), tt.softmax(0).numpy(), rtol=1e-3)
+    out = t.softmax(0, dtype=dtypes.float)
+    self.assertEqual(out.dtype, dtypes.float)
+    np.testing.assert_allclose(out.numpy(), tt.softmax(0, dtype=torch.float).numpy(), rtol=1e-3)
+    out = t.log_softmax(0)
+    self.assertEqual(out.dtype, dtypes.half)
+    np.testing.assert_allclose(out.numpy(), tt.log_softmax(0).numpy(), rtol=1e-3)
+    out = t.log_softmax(0, dtype=dtypes.float)
+    self.assertEqual(out.dtype, dtypes.float)
+    np.testing.assert_allclose(out.numpy(), tt.log_softmax(0, dtype=torch.float).numpy(), rtol=1e-3)

test/unit/test_graph_rewrite.py

@@ -1,7 +1,7 @@
 import unittest, math
 from tinygrad import dtypes
 from tinygrad.helpers import all_same
-from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu
+from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, UPat
 from tinygrad.codegen import full_rewrite_to_sink
 
 # Helper function to apply the graph rewrite
@@ -284,5 +284,11 @@ class TestSubstitute(unittest.TestCase):
     # the srcs are rewritten but we keep tag
     self.assertIs(ret, (b+4).replace(tag=1))
 
+class TestRecurse(unittest.TestCase):
+  def test_no_inf_loop(self):
+    a = UOp.variable('a', 0, 10)
+    pm = PatternMatcher([(UPat(Ops.DEFINE_VAR, name="x"), lambda x: x)])
+    graph_rewrite(a, pm)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_helpers.py

@@ -1,4 +1,4 @@
-import gzip, unittest
+import ctypes, gzip, unittest
 from tinygrad import Variable
 from tinygrad.helpers import Context, ContextVar, argfix
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
@@ -182,7 +182,7 @@ class TestMemoryview(unittest.TestCase):
   def test_from_mv_to_mv(self):
     base = memoryview(bytearray(b"\x11\x22\x33"*40))
     ct = from_mv(base)
-    mv = to_mv(ct, len(base))
+    mv = to_mv(ctypes.addressof(ct), len(base))
     mv[0] = 2
     assert base[0] == 2
 

test/unit/test_linearizer_rewrite.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, Context, Device
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 
 class TestLinearizerRewrite(unittest.TestCase):
   def test_reduction(self):

test/unit/test_pattern_matcher.py

@@ -5,23 +5,23 @@ from tinygrad.uop.ops import PatternMatcher, UPat
 
 class TestPatternMatcher(unittest.TestCase):
   def test_simple_match(self):
-    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float), lambda x: x)])
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float), lambda x: x.rtag())])
     c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
     c2 = UOp(Ops.CONST, dtypes.int, arg=1)
-    self.assertEqual(matcher.rewrite(c1), c1)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), None)
 
   def test_upat_any(self):
     def test(a, x=None, y=None, z=None):
       #print(x,y,z)
-      if y is not None: return a+y
+      if y is not None: return (a+y).rtag()
     matcher = PatternMatcher([
       (UPat.var("a")+UPat.any(UPat.var("x"), UPat.var("y"), UPat.var("z")), test),
     ])
     v1 = UOp.variable("a", 0, 10)
     v2 = UOp.variable("b", 0, 10)
     c1 = v1+v2
-    self.assertEqual(matcher.rewrite(c1), c1)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
 
   def test_minimum_len(self):
     matcher = PatternMatcher([
@@ -60,43 +60,43 @@ class TestPatternMatcher(unittest.TestCase):
     self.assertEqual(len(ctx), 1)
 
   def test_uop(self):
-    matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x)])
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.rtag())])
     c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
     c2 = UOp(Ops.ADD, dtypes.float, (c1, c1))
-    self.assertEqual(matcher.rewrite(c1), c1)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), None)
 
   def test_uop_set(self):
-    matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x)])
+    matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x.rtag())])
     c1 = UOp(Ops.CONST, dtypes.bool, arg=False)
     c2 = UOp(Ops.CAST, dtypes.int, (c1,))
     c3 = UOp(Ops.CONST, dtypes.float, arg=1.0)
     c4 = UOp(Ops.ADD, dtypes.float, (c3, c3))
-    self.assertEqual(matcher.rewrite(c1), c1)
-    self.assertEqual(matcher.rewrite(c2), c2)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c2), c2.rtag())
     self.assertEqual(matcher.rewrite(c4), None)
 
   def test_arg(self):
     matcher = PatternMatcher([
-      (UPat(Ops.CONST, arg=0, name="x"), lambda x: x),
-      (UPat(Ops.CONST, arg=False, name="x"), lambda x: x),
-      (UPat(Ops.MAX, name="x"), lambda x: x),
+      (UPat(Ops.CONST, arg=0, name="x"), lambda x: x.rtag()),
+      (UPat(Ops.CONST, arg=False, name="x"), lambda x: x.rtag()),
+      (UPat(Ops.MAX, name="x"), lambda x: x.rtag()),
     ])
     c1 = UOp(Ops.CONST, dtypes.float, arg=0.0)
     c2 = UOp(Ops.CONST, dtypes.bool, arg=False)
     c3 = UOp(Ops.MAX, dtypes.float, (c1, c1))
     c4 = UOp(Ops.MUL, dtypes.float, (c1, c1))
     c5 = UOp(Ops.CONST, dtypes.int, arg=-1)
-    self.assertEqual(matcher.rewrite(c1), c1)
-    self.assertEqual(matcher.rewrite(c2), c2)
-    self.assertEqual(matcher.rewrite(c3), c3)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c2), c2.rtag())
+    self.assertEqual(matcher.rewrite(c3), c3.rtag())
     self.assertEqual(matcher.rewrite(c4), None)
     self.assertEqual(matcher.rewrite(c5), None)
 
   def test_filter_arg(self):
     matcher = PatternMatcher([
       (UPat(Ops.MUL, src=[UPat(Ops.CONST, name="c"), UPat(Ops.CONST, arg=2)], name="x"),
-       lambda x,c: x if c.arg in {1, -1} else None)
+       lambda x,c: x.rtag() if c.arg in {1, -1} else None)
     ])
     y1 = UOp(Ops.CONST, dtypes.int, arg=1)
     y2 = UOp(Ops.CONST, dtypes.int, arg=2)
@@ -106,45 +106,45 @@ class TestPatternMatcher(unittest.TestCase):
     c3 = UOp(Ops.MUL, dtypes.int, (y3, y2))
     c4 = UOp(Ops.MUL, dtypes.int, (y2, y1))
     c5 = UOp(Ops.MUL, dtypes.int, (y2, y3))
-    self.assertEqual(matcher.rewrite(c1), c1)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), None)
-    self.assertEqual(matcher.rewrite(c3), c3)
-    self.assertEqual(matcher.rewrite(c4), c4)
-    self.assertEqual(matcher.rewrite(c5), c5)
+    self.assertEqual(matcher.rewrite(c3), c3.rtag())
+    self.assertEqual(matcher.rewrite(c4), c4.rtag())
+    self.assertEqual(matcher.rewrite(c5), c5.rtag())
 
   def test_dup_name(self):
-    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x)])
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x.rtag())])
     y1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
     y2 = UOp(Ops.CONST, dtypes.float, arg=1.0)
     c1 = UOp(Ops.ADD, dtypes.float, (y1, y1))
     c2 = UOp(Ops.ADD, dtypes.float, (y1, y2))
-    self.assertEqual(matcher.rewrite(c1), c1)
-    self.assertEqual(matcher.rewrite(c2), c1)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c2), c1.rtag())
 
   def test_dtype(self):
-    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float32), lambda x: x)])
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float32), lambda x: x.rtag())])
     c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
     c2 = UOp(Ops.CONST, dtypes.float64, arg=1.0)
-    self.assertEqual(matcher.rewrite(c1), c1)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
     self.assertEqual(matcher.rewrite(c2), None)
 
   def test_dtype_set(self):
-    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x)])
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())])
     c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
     c2 = UOp(Ops.CONST, dtypes.float64, arg=1.0)
     c3 = UOp(Ops.CONST, dtypes.float16, arg=1.0)
     c4 = UOp(Ops.CONST, dtypes.int, arg=1)
-    self.assertEqual(matcher.rewrite(c1), c1)
-    self.assertEqual(matcher.rewrite(c2), c2)
+    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c2), c2.rtag())
     self.assertEqual(matcher.rewrite(c3), None)
     self.assertEqual(matcher.rewrite(c4), None)
 
   def test_src_one(self):
-    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x)])
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x.rtag())])
     c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
     c2 = UOp(Ops.CONST, dtypes.float, arg=2.0)
     c3 = UOp(Ops.ADD, dtypes.float, (c1,c2))
-    self.assertEqual(matcher.rewrite(c3), c3)
+    self.assertEqual(matcher.rewrite(c3), c3.rtag())
     self.assertEqual(matcher.rewrite(c2), None)
     # that CONST/ALU -> ALU/CONST rewrite is now instant
     """
@@ -157,7 +157,7 @@ class TestPatternMatcher(unittest.TestCase):
     """
 
   def test_src_permutations(self):
-    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x)])
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x.rtag())])
     c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
     c2 = UOp(Ops.CONST, dtypes.float, arg=2.0)
     c3 = UOp(Ops.ADD, dtypes.float, (c1,c2))
@@ -165,21 +165,21 @@ class TestPatternMatcher(unittest.TestCase):
     c5 = UOp(Ops.ADD, dtypes.float, (c2,c3))
     c6 = UOp(Ops.ADD, dtypes.float, (c3,c4))
     self.assertEqual(matcher.rewrite(c3), None)
-    self.assertEqual(matcher.rewrite(c4), c4)
-    self.assertEqual(matcher.rewrite(c5), c5)
+    self.assertEqual(matcher.rewrite(c4), c4.rtag())
+    self.assertEqual(matcher.rewrite(c5), c5.rtag())
     self.assertEqual(matcher.rewrite(c6), None)
 
   def test_src_repeat(self):
-    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x)])
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x.rtag())])
     c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
     c2 = UOp(Ops.CONST, dtypes.float, arg=2.0)
     c3 = UOp(Ops.ADD, dtypes.float, (c1,c2))
     c4 = UOp(Ops.ADD, dtypes.float, (c2,c3))
-    self.assertEqual(matcher.rewrite(c3), c3)
+    self.assertEqual(matcher.rewrite(c3), c3.rtag())
     self.assertEqual(matcher.rewrite(c4), None)
 
   def test_allow_len(self):
-    matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x)])
+    matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x.rtag())])
     c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
     c2 = UOp(Ops.CONST, dtypes.float, arg=2.0)
     c3 = UOp(Ops.CONST, dtypes.float, arg=3.0)
@@ -188,7 +188,7 @@ class TestPatternMatcher(unittest.TestCase):
     c6 = UOp(Ops.MULACC, dtypes.float, (c1,c2,c3))
     self.assertEqual(matcher.rewrite(c4), None)
     self.assertEqual(matcher.rewrite(c5), None)
-    self.assertEqual(matcher.rewrite(c6), c6)
+    self.assertEqual(matcher.rewrite(c6), c6.rtag())
 
   def test_deep_src_permutations(self):
     c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)

test/unit/test_rewrite_tracked_childen.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor
 from tinygrad.uop.ops import PatternMatcher, Ops, UPat, graph_rewrite, RewriteContext, UOp
-from tinygrad.engine.grouper import sym, merge_views
+from tinygrad.engine.kernelize import sym, merge_views
 
 class TestRewriteTrackedChildren(unittest.TestCase):
   @unittest.skip("track_children no longer supported")

test/unit/test_search.py

@@ -1,5 +1,10 @@
 import unittest
-from tinygrad.engine.search import get_test_global_size
+from tinygrad import Tensor, Device
+from tinygrad.opt.kernel import Kernel
+from tinygrad.device import Buffer
+from tinygrad.opt.search import get_test_global_size, bufs_from_lin
+from tinygrad.helpers import GlobalCounters
+from extra.optimization.helpers import time_linearizer
 
 class TestSearchUtil(unittest.TestCase):
   def test_get_test_global_size(self):
@@ -7,5 +12,44 @@ class TestSearchUtil(unittest.TestCase):
     self.assertEqual(get_test_global_size([65536, 1, 1], 256, {}), ([256, 1, 1], 256.0))
     self.assertEqual(get_test_global_size([77, 1, 1], 16, {}), ([9, 1, 1], 77/9))
 
+  def test_bufs_from_lin(self):
+    a = Tensor([1,2,3,4]).realize()
+    si = (a+1).schedule()[0]
+    rawbufs = bufs_from_lin(lin:=Kernel(si.ast))
+    assert len(rawbufs) == len(lin.membufs) == 2
+    assert all(r is not None for r in rawbufs)
+    assert all(isinstance(r, Buffer) for r in rawbufs)
+    assert all(r.size > 0 for r in rawbufs)
+
+  def test_bufs_from_lin_alt(self):
+    a = Tensor.randn(4, 4).realize()
+    b = a+a[0]
+    si = b.schedule()[0]
+    rawbufs = bufs_from_lin(k:=Kernel(si.ast))
+    assert len(rawbufs) == len(k.membufs) == 2
+    assert all(r is not None for r in rawbufs)
+    assert all(isinstance(r, Buffer) for r in rawbufs)
+    assert all(r.size > 0 for r in rawbufs)
+
+class TestTimeLinearizer(unittest.TestCase):
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WebGPU timestamps are low precision, tm is 0")
+  def test_reasonable_time(self):
+    a = Tensor([1,2,3,4]).realize()
+    si = (a+1).schedule()[0]
+    # create fresh empty buffers
+    rawbufs = [Buffer(b.device, b.size, b.dtype).allocate() for b in si.bufs]
+    tm = time_linearizer(Kernel(si.ast), rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
+    assert tm > 0 and tm != float('inf')
+
+  # Ensure that the kernel count is not incremented by time_linearizer when clearing l2
+  def test_kernel_count(self):
+    ast = Tensor.zeros(16).contiguous().kernelize().uop.src[1].arg.ast
+    lin = Kernel(ast)
+    bufs = bufs_from_lin(lin)
+
+    kernel_count = GlobalCounters.kernel_count
+    time_linearizer(lin, bufs, allow_test_size=False, cnt=2, disable_cache=True, clear_l2=True)
+    assert GlobalCounters.kernel_count == kernel_count, "kernel count was incremented by time_linearizer"
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_simple_schedule.py

@@ -0,0 +1,15 @@
+import unittest
+from tinygrad import Tensor
+from tinygrad.uop.ops import Ops
+
+class TestSimpleSchedule(unittest.TestCase):
+  def test_reduce_doesnt_split(self):
+    a = Tensor.empty(16,16).sum(axis=1)
+    a1 = a.reshape(4,4)
+    a2 = a.reshape(16,1,1)
+    Tensor.kernelize(a1, a2)
+    kernels = [x for x in a1.uop.sink(a2.uop).toposort() if x.op is Ops.KERNEL]
+    self.assertEqual(len(kernels), 1)
+
+if __name__ == '__main__':
+  unittest.main()

test/unit/test_tqdm.py

@@ -80,7 +80,7 @@ class TestProgressBar(unittest.TestCase):
 
           # compare bars at each iteration (only when tinytqdm bar has been updated)
           # setting high rate to make sure it does not skip
-          for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=1e9):
+          for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=10**9):
             tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
 
             if n:
@@ -103,7 +103,7 @@ class TestProgressBar(unittest.TestCase):
     total = 10
     with patch('time.perf_counter', side_effect=[0]+list(range(100))):  # one more 0 for the init call
       # compare bars at each iteration (only when tinytqdm bar has been updated)
-      for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=1e9):
+      for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=10**9):
         tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
         elapsed = n
         tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=unit_scale)
@@ -120,7 +120,7 @@ class TestProgressBar(unittest.TestCase):
     # E   ?              +  ^
     with patch('time.perf_counter', side_effect=[0, *[i*k for i in range(100)]]):  # one more 0 for the init call
       # compare bars at each iteration (only when tinytqdm bar has been updated)
-      for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=1e9):
+      for n in tinytqdm(range(total), desc="Test", total=total, unit_scale=unit_scale, rate=10**9):
         tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
         elapsed = n*k
         tqdm_output = tqdm.format_meter(n=n, total=total, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=unit_scale)
@@ -238,7 +238,7 @@ class TestProgressBar(unittest.TestCase):
 
         # compare bars at each iteration (only when tinytqdm bar has been updated)
         # setting high rate to make sure it does not skip
-        for n,g in enumerate(tinytqdm(gen, desc="Test", unit_scale=unit_scale, rate=1e9)):
+        for n,g in enumerate(tinytqdm(gen, desc="Test", unit_scale=unit_scale, rate=10**9)):
           assert g == n
           tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
           if n:

test/unit/test_uop_spec.py

@@ -2,10 +2,10 @@ from __future__ import annotations
 import unittest
 
 from tinygrad import Tensor
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.helpers import DEBUG
 from tinygrad.uop.ops import UOp, Ops, print_uops
-from tinygrad.uop.spec import type_verify, ast_spec
+from tinygrad.uop.spec import type_verify, ast_spec, tensor_uop_spec
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad import dtypes
 from tinygrad.shape.view import View
@@ -23,7 +23,7 @@ def helper_test_verify_ast(*stores:UOp) -> Kernel:
   if DEBUG >= 4: print(k.to_program().src)
   return k
 
-class TestVerifyAST(unittest.TestCase):
+class TestUOpSpec(unittest.TestCase):
   def test_tiny_add(self):
     dtype = dtypes.int
     buf_0 = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 0)
@@ -94,5 +94,11 @@ class TestVerifyAST(unittest.TestCase):
     st = UOp.store(buf.view(ShapeTracker.from_shape(())), a.cast(dtypes.float))
     helper_test_verify_ast(st)
 
+  def test_assert_masked_view_in_const(self):
+    t = Tensor(6).uop
+    a = t.replace(src=(t.src[0].replace(arg=t.st.reshape((1,)).pad(((0, 1),))),))
+    with self.assertRaisesRegex(RuntimeError, "UOp verification failed"):
+      type_verify([a], tensor_uop_spec)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_uop_symbolic.py

@@ -110,6 +110,9 @@ class TestSymbolic(unittest.TestCase):
   def test_neg(self):
     self.helper_test_variable(-Variable("a", 0, 8), -8, 0, "(a*-1)")
 
+  def test_xor_0(self):
+    self.helper_test_variable(Variable("a", 0, 8) ^ 0, 0, 8, "a")
+
   def test_add_1(self):
     self.helper_test_variable(Variable("a", 0, 8)+1, 1, 9, "(a+1)")
 
@@ -268,6 +271,15 @@ class TestSymbolic(unittest.TestCase):
     a = Variable("a", 0, 124)
     self.helper_test_variable(((a-10)//2+10)//2, 2, 33, "((((a+-10)//2)+10)//2)")
 
+  def test_div_const_div_wrong_sign_divisor(self):
+    a = Variable("a", 0, 124)
+    self.helper_test_variable(((a+10)//-2+10)//-4, -1, 14, "(((((a//2)*-1)+5)//4)*-1)")
+
+  def test_neg_mod(self):
+    a = Variable("a", 0, 124)
+    self.helper_test_variable((-a)%4, -3, 0, "((a%4)*-1)")
+    self.helper_test_variable(a%-4, 0, 3, "(a%-4)")
+
   def test_distribute_mul(self):
     self.helper_test_variable(usum([Variable("a", 0, 3), Variable("b", 0, 5)])*3, 0, 24, "((a*3)+(b*3))")
     self.helper_test_variable((1+Variable("a", 0, 3))*(-2)+12, 4, 10, "((a*-2)+10)")
@@ -398,6 +410,8 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((-Variable("idx", 0, 100)+199)//-4 + 50, 1, 26, "((idx//4)+1)")
     self.helper_test_variable((-Variable("idx", 0, 100)+200)//-4 + 50, 0, 25, "((idx+3)//4)")
     self.helper_test_variable((-Variable("idx", 0, 100)+201)//-4 + 50, 0, 25, "((idx+2)//4)")
+    self.helper_test_variable((-Variable("idx", 0, 100))//2, -50, 0, "((idx//2)*-1)")
+    self.helper_test_variable(Variable("idx", 0, 100)//-2, -50, 0, "((idx//2)*-1)")
 
   def test_sum_div_big_const(self):
     gidx0 = Variable("gidx0", 0, 24)

test/unit/test_view.py

@@ -44,6 +44,16 @@ class TestView(unittest.TestCase):
     self.assertIsNotNone(v)
     self.assertEqual(v, View.create((20,), mask=((0,0),)))
 
+  def test_add_0(self):
+    v1 = View.create((2,3,4))
+    v2 = View.create((2,0,4))
+    self.assertEqual(v2, v1+v2)
+
+  def test_add_0_masked(self):
+    v1 = View.create((2,3,4), mask=((0, 0), (0, 0), (0, 0)))
+    v2 = View.create((2,0,4))
+    self.assertEqual(v2, v1+v2)
+
 class TestMergeDims(unittest.TestCase):
   def test_contiguous(self):
     shape = (2, 3, 4)

test/unit/test_viz.py

@@ -1,7 +1,7 @@
 import unittest, decimal, json
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import TRACK_MATCH_STATS, TrackedPatternMatcher, UOp, graph_rewrite, track_rewrites, UPat, Ops
-from tinygrad.uop.symbolic import symbolic
+from tinygrad.uop.symbolic import symbolic, symbolic_simple
 from tinygrad.uop.ops import tracked_ctxs as contexts, tracked_keys as keys, _name_cnt, _substitute
 from tinygrad.device import ProfileDeviceEvent, ProfileRangeEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_metadata, get_details, uop_to_json, to_perfetto
@@ -30,30 +30,30 @@ class TestViz(unittest.TestCase):
 
   def test_viz_simple(self):
     a = UOp.variable("a", 0, 10)
-    @track_rewrites(named=True)
+    @track_rewrites(name=True)
     def test(sink): return graph_rewrite(sink, symbolic)
     test(a*1)
     ret = get_metadata(keys, contexts)
     self.assertEqual(len(ret), 1)
     key, val = ret[0]["name"], ret[0]["steps"]
-    self.assertEqual(key, "test_1")
+    self.assertEqual(key, "test n1")
     self.assertEqual(val[0]["match_count"], 1)
 
   def test_track_two_rewrites(self):
     a = UOp.variable("a", 0, 10)
-    @track_rewrites(named=True)
+    @track_rewrites(name=True)
     def test(sink): return graph_rewrite(sink, symbolic)
     test((a+a)*1)
     ret = get_metadata(keys, contexts)
     key, val = ret[0]["name"], ret[0]["steps"]
     self.assertEqual(len(ret), 1)              # one context
     self.assertEqual(len(val), 1)              # one graph_rewrite call in context
-    self.assertEqual(key, "test_1")
+    self.assertEqual(key, "test n1")
     self.assertEqual(val[0]["match_count"], 2) # two upats applied
 
   def test_track_multiple_calls_one_ctx(self):
     a = UOp.variable("a", 0, 10)
-    @track_rewrites(named=True)
+    @track_rewrites(name=True)
     def test(a, b):
       a = graph_rewrite(a, symbolic)
       b = graph_rewrite(b, symbolic)
@@ -62,12 +62,12 @@ class TestViz(unittest.TestCase):
     key, val = ret[0]["name"], ret[0]["steps"]
     self.assertEqual(len(ret), 1)              # one context
     self.assertEqual(len(val), 2)              # two graph_rewrite calls in context
-    self.assertEqual(key, "test_1")
+    self.assertEqual(key, "test n1")
     self.assertEqual(val[0]["match_count"], 1) # one rewrite for a*0
     self.assertEqual(val[1]["match_count"], 0) # no rewrites for a*5
 
   def test_track_rewrites(self):
-    @track_rewrites(named=True)
+    @track_rewrites(name=True)
     def do_rewrite(x:UOp): return graph_rewrite(x, symbolic)
     a = UOp.variable("a", 0, 10)
     b = UOp.variable("b", 0, 4)
@@ -76,10 +76,10 @@ class TestViz(unittest.TestCase):
     ret = get_metadata(keys, contexts)
     self.assertEqual(len(ret), 2)
     key, m = ret[0]["name"], ret[0]["steps"]
-    self.assertEqual(key, "do_rewrite_1")
+    self.assertEqual(key, "do_rewrite n1")
     self.assertEqual(m[0]["match_count"], 1)
     key, m = ret[1]["name"], ret[1]["steps"]
-    self.assertEqual(key, "do_rewrite_2")
+    self.assertEqual(key, "do_rewrite n2")
     self.assertEqual(m[0]["match_count"], 0)
 
   def test_track_rewrites_with_exception(self):
@@ -93,7 +93,7 @@ class TestViz(unittest.TestCase):
     self.assertEqual(len(ret), 1)
 
   def test_track_rewrites_name_fxn(self):
-    @track_rewrites(name_fxn=lambda _,ret: f"output_{ret}")
+    @track_rewrites(name=lambda _,ret: f"output_{ret}")
     def do_rewrite(x:UOp):
       x = graph_rewrite(x, symbolic)
       return x.render()
@@ -109,7 +109,7 @@ class TestViz(unittest.TestCase):
 
   @unittest.expectedFailure
   def test_name_in_positional_arg(self):
-    @track_rewrites(named=True)
+    @track_rewrites(name=True)
     def test(sink): return graph_rewrite(sink, symbolic, None, False, "name")
     test(UOp.variable("a", 0, 1))
     self.assertEqual(contexts[0].pop().name, "name")
@@ -125,7 +125,7 @@ class TestViz(unittest.TestCase):
     a = UOp.variable("a", 0, 10)
     b = UOp.variable("b", 0, 10)
     c = UOp.variable("c", 0, 10)
-    @track_rewrites(named=True)
+    @track_rewrites(name=True)
     def fxn(sink): return graph_rewrite(sink, substitute, ctx={a+b:c}, bottom_up=True)
     fxn(a+b)
     #UOp.substitute(a+b, {a+b:c})
@@ -137,7 +137,7 @@ class TestViz(unittest.TestCase):
   # NOTE: calling graph_rewrite when the function isn't decorated with track_rewrites should not VIZ
   def test_rewrite_without_context(self):
     def untracked_graph_rewrite(sink): return graph_rewrite(sink, symbolic)
-    @track_rewrites(named=True)
+    @track_rewrites(name=True)
     def tracked_graph_rewrite(sink): return graph_rewrite(sink, symbolic)
     # test
     add = UOp.const(dtypes.int, 2) + UOp.const(dtypes.int, 1)
@@ -149,7 +149,7 @@ class TestViz(unittest.TestCase):
   def test_inner_rewrite_location(self):
     # inner rewrite gets tracked in another context
     def inner_rewrite(sink): return graph_rewrite(sink, symbolic)
-    @track_rewrites(named=True)
+    @track_rewrites(name=True)
     def tracked_graph_rewrite(sink): return inner_rewrite(sink)
     # test
     add = UOp.const(dtypes.int, 2) + UOp.const(dtypes.int, 1)
@@ -160,6 +160,10 @@ class TestViz(unittest.TestCase):
     self.assertEqual(lineno, inner_rewrite.__code__.co_firstlineno)
     self.assertEqual(fp, inner_rewrite.__code__.co_filename)
 
+  def test_upat_location(self):
+    for (pat, fn) in symbolic_simple.patterns:
+      self.assertIn("symbolic.py", pat.location[0])
+
   def test_nested_rewrite(self):
     def make_float(x:UOp, y:UOp):
       if x.dtype == dtypes.float: return None
@@ -167,7 +171,7 @@ class TestViz(unittest.TestCase):
       y2 = graph_rewrite(y, inner_rewrite, name="inner_y")
       return None if (x2 is x and y2 is y) else x2+y2
     outer_rewrite = TrackedPatternMatcher([(UPat.cvar("x")+UPat.cvar("y"), make_float),])
-    @track_rewrites(named=True)
+    @track_rewrites(name=True)
     def rewrite(u:UOp): return graph_rewrite(u, outer_rewrite, name="outer")
     a = UOp.const(dtypes.int, 1)+UOp.const(dtypes.int, 2)
     rewrite(a)
@@ -182,7 +186,7 @@ class TestViz(unittest.TestCase):
     self.assertEqual([len(x.matches) for x in tracked], [1, 1, 1])
 
   def test_depth_level(self):
-    @track_rewrites(named=True)
+    @track_rewrites(name=True)
     def fxn(u:UOp): return graph_rewrite(u, l0)
     ret = fxn(UOp(Ops.CUSTOM, arg=0))
     assert ret is UOp(Ops.CUSTOM, arg=3)
@@ -207,15 +211,15 @@ class TestViz(unittest.TestCase):
     @track_rewrites()
     def test_fxn(): return graph_rewrite(test, l0)
     assert test_fxn() is test
-    self.assertEqual(keys[0], "test_fxn_1")
+    self.assertEqual(keys[0], "test_fxn n1")
 
   @unittest.skip("TODO: doesn't work")
   def test_recursion_err(self):
     inf = TrackedPatternMatcher([
-      (UPat.const(dtypes.int, 0).named("a"), lambda a: a.const_like(1)),
-      (UPat.const(dtypes.int, 1).named("b"), lambda b: b.const_like(0)),
+      (UPat.const(dtypes.int, 0).name("a"), lambda a: a.const_like(1)),
+      (UPat.const(dtypes.int, 1).name("b"), lambda b: b.const_like(0)),
     ])
-    @track_rewrites(named=True)
+    @track_rewrites(name=True)
     def func(u): return graph_rewrite(u, inf)
     with self.assertRaises(RecursionError): func(UOp.const(dtypes.int, 0))
     _ = list(get_details(keys[0], contexts[0][0]))

test/web/test_viz.js

@@ -10,23 +10,24 @@ async function main() {
   }));
 
   // ** run browser tests
-  let browser;
+  let browser, page;
   try {
     browser = await puppeteer.launch({ headless: true });
-    const page = await browser.newPage();
+    page = await browser.newPage();
     const res = await page.goto("http://localhost:8000", { waitUntil:"domcontentloaded" });
     if (res.status() !== 200) throw new Error("Failed to load page");
-    const scheduleSelector = await page.waitForSelector("ul");
+    const scheduleSelector = await page.waitForSelector("ul:nth-of-type(2)");
     scheduleSelector.click();
     await page.waitForSelector("rect");
-    const nodes = await page.evaluate(() => document.querySelectorAll("#nodes > g").length);
-    const edges = await page.evaluate(() => document.querySelectorAll("#edges > path").length);
-    if (!nodes || !edges) {
-      throw new Error("VIZ didn't render a graph")
-    }
+    await page.waitForFunction(() => {
+      const nodes = document.querySelectorAll("#nodes > g").length;
+      const edges = document.querySelectorAll("#edges > path").length;
+      return nodes > 0 && edges > 0;
+    });
   } finally {
     // ** cleanups
-    if (browser) await browser.close();
+    if (page != null) await page.close();
+    if (browser != null) await browser.close();
     proc.kill();
   }
 }

tinygrad/codegen/__init__.py

@@ -3,12 +3,13 @@ import functools
 from dataclasses import dataclass
 from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp
+from tinygrad.uop.spec import type_verify
 from tinygrad.renderer import Renderer
 
 # import all pattern matchers here
 from tinygrad.codegen.lowerer import pm_quant, pm_lowerer, get_index
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
-from tinygrad.codegen.expander import migrate_indexing, pm_store_ignore, pm_move_ignore, pm_delete_ignore, expander
+from tinygrad.codegen.expander import migrate_indexing, expander
 from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, \
   pm_reduce, ReduceContext, correct_load_store, pm_render, get_late_rewrite_patterns
 from tinygrad.codegen.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
@@ -44,12 +45,8 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   # ** expander (expand_rewrite) **
   ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic"))
 
-  # ignore (for masked stores)
-  ret.append(RewriteStep(pm_store_ignore, name="store_ignore"))
-  ret.append(RewriteStep(pm_move_ignore, name="move_ignore"))
-
-  # expand + remove surviving ignores
-  ret.append(RewriteStep(pm_delete_ignore+sym+expander, name="expander"))
+  # expand
+  ret.append(RewriteStep(sym+expander, name="expander"))
 
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce
@@ -76,4 +73,8 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
 
 def full_rewrite_to_sink(sink:UOp, opts:Renderer|None=None, linearizer:bool=False) -> UOp:
   return apply_rewrites(sink, get_rewrites_for_renderer(opts if opts is not None else Renderer(), linearizer))
-def full_rewrite(sink:UOp, opts:Renderer|None=None) -> list[UOp]: return list(full_rewrite_to_sink(sink, opts, linearizer=True).arg.lst)
+
+def full_rewrite(sink:UOp, opts:Renderer|None=None) -> list[UOp]:
+  lst = list(full_rewrite_to_sink(sink, opts, linearizer=True).arg.lst)
+  if __debug__: type_verify(lst)
+  return lst

tinygrad/codegen/devectorizer.py

@@ -4,7 +4,7 @@ from collections import defaultdict
 from dataclasses import dataclass
 from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import dtypes, ImageDType, PtrDType, promo_lattice, DType
-from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, resolve, graph_rewrite, GroupOp, identity_element
+from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, graph_rewrite, GroupOp, identity_element
 from tinygrad.uop.symbolic import split_uop, uop_given_valid, parse_valid, simplify_valid, sym, symbolic_flat
 from tinygrad.helpers import getenv, flatten, AMX, prod, partition
 from tinygrad.uop.transcendental import xexp2, xlog2, xsin, xpow, TRANSCENDENTAL_SUPPORTED_DTYPES
@@ -153,7 +153,7 @@ def magicgu(vmax:int, d:int) -> tuple[int,int]:
   assert False
 
 def fast_idiv(ctx: Renderer|None, x: UOp, d: int) -> UOp|None:
-  # idiv is truncated division, but arithmatic shift is floored division, so can only do non-negative numbers!
+  # idiv is truncated division, but arithmetic shift is floored division, so can only do non-negative numbers!
   if x.vmin<0: return None
   sign = 1 if d > 0 else -1
   m,s = magicgu(vmax := min(x.vmax, dtypes.max(x.dtype)), abs(d))
@@ -175,9 +175,10 @@ def get_late_rewrite_patterns(ops, force_transcendental=False):
   # rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
   if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
   if Ops.SHR in ops:
-    # no reason to check x>=0 for uints
+    # no reason to check x<0 for uints
     pat += [(UPat.var("x", dtypes.uints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
-    pat += [(UPat.var("x", dtypes.sints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) and resolve(x>=0,False) else None)]
+    pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(
+      c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]  # (x+(x<0).where(c-1, 0)) >> v
     if not getenv("DISABLE_FAST_IDIV"):
       pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d"), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))]
       pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("d"), lambda ctx, x, d: x - d*f if (f:=fast_idiv(ctx, x, d.arg)) is not None else None)]

tinygrad/codegen/expander.py

@@ -114,25 +114,3 @@ migrate_indexing = PatternMatcher([
   # create gate MUST BE BEFORE expander
   (UPat(Ops.STORE, name="root"), create_gate),
 ])
-
-# **** IGNORE support ****
-
-pm_store_ignore = PatternMatcher([
-  (UPat().index(UPat(), UPat(name="mask")).store(UPat()).named("store"),
-   lambda store,mask: store.replace(src=(store.src[0], UOp(Ops.IGNORE, src=(store.src[1], mask)))) if store.src[1].op is not Ops.IGNORE else None),
-])
-
-pm_move_ignore = PatternMatcher([
-  # IGNORE on SELF is nothing
-  (UPat(Ops.IGNORE, src=(UPat(name="x"), UPat(name="x"))), lambda x: x.const_like(True)),
-  # IGNORE on a CONST is nothing
-  (UPat(Ops.IGNORE, src=(UPat((Ops.CONST, Ops.VCONST), name="c"), UPat())), lambda c: c),
-  # move the IGNOREs
-  (UPat(Ops.IGNORE, src=(UPat((*GroupOp.ALU, Ops.CAST, Ops.VECTORIZE), name="alu"), UPat.var("mask")), name="ig"),
-   lambda ig,alu,mask: alu.replace(src=tuple(UOp(Ops.IGNORE, x.dtype, (x, mask)) for x in alu.src))),
-])
-
-pm_delete_ignore = PatternMatcher([
-  # IGNORE on SELF is nothing
-  (UPat(Ops.IGNORE, src=(UPat(name="x"), UPat())), lambda x: x),
-])