fix unit test dtypes

render ranges in viz, name gbufs with sizes. changes from rangeify
[bounty] Muon optim (#11414 )
2025-08-13 12:15:12 -07:00 · 2025-08-13 12:06:34 -07:00 · 2025-08-13 14:27:55 -04:00 · 2025-08-13 13:51:08 -04:00 · 2025-08-13 09:37:12 -07:00 · 2025-08-13 08:22:06 -07:00
135 changed files with 1772 additions and 1499 deletions
@@ -325,7 +325,7 @@ jobs:
      run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA Training)
@@ -576,7 +576,7 @@ jobs:
      run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD MLPerf)
@@ -611,12 +611,12 @@ jobs:
      run: BENCHMARK_LOG=openpilot_0_9_7 PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_7.txt
    - name: benchmark openpilot w IMAGE=2 0.9.7
      run: BENCHMARK_LOG=openpilot_0_9_7_image PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
-    - name: openpilot compile3 0.9.7
-      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx
-    - name: openpilot compile3 0.9.7+ tomb raider
-      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/e8bea2c78ffa92685ece511e9b554122aaf1a79d/selfdrive/modeld/models/supercombo.onnx
-    - name: openpilot dmonitoring compile3 0.9.7
-      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 0.9.9 driving_vision
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.9.9 driving_policy
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.9.9 dmonitoring
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 Space Lab policy + vision
      run: |
        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
@@ -692,11 +692,12 @@ jobs:
      run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
    - name: Run full CIFAR training w 1 GPU
      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
+    # TODO: enable
+    # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #   run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
+      run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AM Driver)
@@ -754,7 +755,7 @@ jobs:
      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
+      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NV Driver)
@@ -375,8 +375,8 @@ jobs:
        PYTHONPATH=. python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
        DEBUG=1 MIN_ASTS=1 PYTHONPATH=. python extra/optimization/get_action_space.py
-    - name: Repo line count < 15500 lines
-      run: MAX_LINE_COUNT=15500 python sz.py
+    - name: Repo line count < 16000 lines
+      run: MAX_LINE_COUNT=16000 python sz.py

  fuzzing:
    name: Fuzzing
@@ -542,8 +542,8 @@ jobs:
        run: PYTHONPATH="." GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
      - name: Test MLPerf stuff
        run: GPU=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
-      - name: Run handcode_opt
-        run: PYTHONPATH=. MODEL=resnet GPU=1 DEBUG=1 BS=4 HALF=0 python3 examples/handcode_opt.py
+      - name: Test llama 3 training
+        run: MAX_BUFFER_SIZE=0 PYTHONPATH="." DEV=NULL SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -626,7 +626,7 @@ jobs:
    - name: Test LLVM=1 DEVECTORIZE=0 for model
      run: PYTHONPATH="." LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
    - name: Test CPU=1 DEVECTORIZE=0
-      run: CPU=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+      run: CPU=1 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"

  testwebgpu:
    name: Linux (WebGPU)
@@ -977,6 +977,7 @@ jobs:
        with:
          key: macos-${{ matrix.backend }}-minimal
          deps: testing_minimal
+          pydeps: "capstone"
          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
@@ -18,11 +18,11 @@ Group UOps into kernels.

 ---

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

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

-::: tinygrad.opt.get_optimized_ast
+::: tinygrad.codegen.opt.get_optimized_ast
    options:
        members: false
        show_labels: false
@@ -47,8 +47,8 @@ Reboot after making these changes or restart the `displayservice.service` servic

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

-## tinychat
+## Building the OS image

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

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

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

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

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

    start = self.idx.tell()
    end = start + self.count * dtypes.int32.itemsize
-    self.sizes = self.idx_t[start:end].bitcast(dtypes.int32)
+    self.sizes = self.idx_t[start:end].bitcast(dtypes.int32).numpy()

    start = end
    end = start + self.count * dtypes.int64.itemsize
-    self.pointers = self.idx_t[start:end].bitcast(dtypes.int64)
+    self.pointers = self.idx_t[start:end].bitcast(dtypes.int64).numpy()

    start = end
    end = start + doc_count * dtypes.int64.itemsize
-    self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64)
+    self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64).numpy()

    # bin file
    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))

  def _index(self, idx) -> tuple[int, int]:
-    return self.pointers[idx].item(), self.sizes[idx].item()
+    return int(self.pointers[idx]), int(self.sizes[idx])

  def get(self, idx, offset:int=0, length:int|None=None):
    ptr, size = self._index(idx)
@@ -566,10 +564,13 @@ class GPTDataset:
    # check for cache
    cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
    cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
+    print(f"try loading GPTDataset from {cache_path}...")
    if cache_path.exists():
+      print("cache found, loading...")
      with open(cache_path, "rb") as f:
        self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
    else:
+      print("cache not found, building index...")
      self.doc_idx = self._build_doc_idx()
      self.sample_idx = self._build_sample_idx()
      self.shuffle_idx = self._build_shuffle_idx()
@@ -628,14 +629,20 @@ class GPTDataset:

  # https://github.com/NVIDIA/Megatron-LM/blob/94bd476bd840c2fd4c3ebfc7448c2af220f4832b/megatron/core/datasets/gpt_dataset.py#L558
  def _build_doc_idx(self):
-    doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
-    doc_idx = doc_idx.reshape(-1)
+    print(f"building doc_idx for {self.num_epochs=}, {self.indexed_dataset.count=}")
+    st = time.perf_counter()
+    # doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
+    doc_idx = np.arange(self.indexed_dataset.count).reshape(1, -1).repeat(self.num_epochs, axis=0).flatten()
    doc_idx = doc_idx.astype(np.int32)
+    at = time.perf_counter()
    if self.shuffle: self.rng.shuffle(doc_idx)
+    print(f"doc_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
    return doc_idx

  def _build_sample_idx(self):
-    sample_idx = np.empty((self.samples + 1, 2), dtype=np.int32)
+    print(f"building sample_idx for {self.samples=}, {self.seqlen=}, {self.doc_idx.shape[0]=}")
+    sample_idx_max = max(self.doc_idx.shape[0], self.indexed_dataset.sizes.max())
+    sample_idx = np.empty((self.samples + 1, 2), dtype=np.int64 if sample_idx_max > dtypes.int32.max else np.int32)

    sample_idx_idx, doc_idx_idx, doc_offset = 0, 0, 0
    sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
@@ -645,7 +652,7 @@ class GPTDataset:
      remaining_seqlen = self.seqlen + 1
      while remaining_seqlen > 0:
        doc_idx = int(self.doc_idx[doc_idx_idx])
-        doc_len = self.indexed_dataset.sizes[doc_idx].item() - doc_offset
+        doc_len = int(self.indexed_dataset.sizes[doc_idx]) - doc_offset
        remaining_seqlen -= doc_len
        if remaining_seqlen <= 0:
          doc_offset += remaining_seqlen + doc_len - 1
@@ -654,7 +661,7 @@ class GPTDataset:
          if doc_idx_idx == len(self.doc_idx) - 1:
            assert sample_idx_idx == self.samples
            doc_idx = int(self.doc_idx[doc_idx_idx])
-            doc_offset = self.indexed_dataset.sizes[doc_idx].item() - 1
+            doc_offset = int(self.indexed_dataset.sizes[doc_idx]) - 1
            break
          doc_idx_idx += 1
          doc_offset = 0
@@ -665,13 +672,18 @@ class GPTDataset:
    return sample_idx

  def _build_shuffle_idx(self):
+    print(f"building shuffle_idx for {self.samples=}")
+    st = time.perf_counter()
    shuffle_idx = np.arange(self.samples, dtype=np.int32)
+    at = time.perf_counter()
    if self.shuffle: self.rng.shuffle(shuffle_idx)
+    print(f"shuffle_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
    return shuffle_idx

 class BlendedGPTDataset:
  def __init__(self, paths:list[Path], weights:list[float], samples:int, seqlen:int, seed:int, shuffle:bool):
-    self.seed = seed
+    self.shuffle = shuffle
+    self.rng = np.random.RandomState(seed)

    # normalize weights
    total_weight = sum(weights)
@@ -683,10 +695,47 @@ class BlendedGPTDataset:

    self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]

+    # check for cache
+    cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
+    cache_path = paths[0].with_name(f"{paths[0].name}.{cache_hash}.blend_cache")
+    print(f"try loading BlendedGPTDataset from {cache_path}...")
+    if cache_path.exists():
+      print("cache found, loading...")
+      with open(cache_path, "rb") as f:
+        self.dataset_idx, self.dataset_sample_idx = pickle.load(f)
+    else:
+      print("cache not found, building index...")
+      self.dataset_idx, self.dataset_sample_idx = self._build_blend_idx()
+      # save cache
+      with open(cache_path, "wb") as f:
+        pickle.dump((self.dataset_idx, self.dataset_sample_idx), f)
+
  def get(self, idx:int):
-    tokens = self.datasets[0][idx]
+    tokens = self.datasets[self.dataset_idx[idx]][self.dataset_sample_idx[idx]]
    return tokens

+  def _build_blend_idx(self):
+    dataset_idx = np.zeros(self.samples, dtype=np.int16)
+    dataset_sample_idx = np.zeros(self.samples, dtype=np.int64)
+
+    unspent_datasets = set(range(len(self.datasets)))
+    dataset_sample_counts = [0] * len(self.datasets)
+
+    for i in tqdm(range(self.samples)):
+      error_argmax, error_max = 0, 0.0
+      for di in unspent_datasets:
+        error = self.weights[di] * max(i, 1) - dataset_sample_counts[di]
+        if error > error_max:
+          error_max = error
+          error_argmax = di
+
+      dataset_idx[i] = error_argmax
+      dataset_sample_idx[i] = dataset_sample_counts[error_argmax]
+
+      dataset_sample_counts[error_argmax] += 1
+
+    return dataset_idx, dataset_sample_idx
+
 def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
  if val:
    dataset = BlendedGPTDataset([
@@ -739,8 +788,8 @@ if __name__ == "__main__":

  def load_llama3(val):
    bs = 24
-    samples = 5760 if val else 1_200_000
-    seqlen = 512
+    samples = 5760 if val else 1_200_000 * 1152
+    seqlen = 8192

    max_, min_ = 0, math.inf
    for tokens in tqdm(batch_load_llama3(bs, samples, seqlen, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=5760, val=bool(val)), total=samples//bs):
@@ -1296,7 +1296,7 @@ def train_llama3():
  SEED               = config["SEED"]                   = getenv("SEED", 5760)
  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
  TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
-  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000)
+  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)

  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
  # trains to 7
@@ -1323,24 +1323,22 @@ def train_llama3():
    for v in get_parameters(model):
      v.shard_(device, axis=None)

-  # TODO: MP
-  # if (GPUS := getenv("GPUS", 1)) > 1:
-  #   device = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS))
-  #   for k,v in get_state_dict(model).items():
-  #     if 'scale' in k: v.shard_(device, axis=None)  # from quantized
-  #     # elif '.attention.wq' in k: v.shard_(device, axis=0)
-  #     # elif '.attention.wk' in k: v.shard_(device, axis=0)
-  #     # elif '.attention.wv' in k: v.shard_(device, axis=0)
-  #     # elif '.attention.wo' in k: v.shard_(device, axis=1)
-  #     # elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
-  #     # elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
-  #     # elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
-  #     # elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
-  #     elif 'output.weight' in k: v.shard_(device, axis=0)  # 243.32
-  #     else:
-  #       # print(k)
-  #       # attention_norm, ffn_norm, norm
-  #       v.shard_(device, axis=None)
+  if (MP := getenv("MP", 1)) > 1:
+    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+    for k,v in get_state_dict(model).items():
+      if 'scale' in k: v.shard_(device, axis=None)  # from quantized
+      elif '.attention.wq' in k: v.shard_(device, axis=0)
+      elif '.attention.wk' in k: v.shard_(device, axis=0)
+      elif '.attention.wv' in k: v.shard_(device, axis=0)
+      elif '.attention.wo' in k: v.shard_(device, axis=1)
+      elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
+      elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
+      elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
+      elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
+      elif 'output.weight' in k: v.shard_(device, axis=0)
+      else:
+        # attention_norm, ffn_norm, norm
+        v.shard_(device, axis=None)

  optim = AdamW(get_parameters(model), lr=0.0,
                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
@@ -1355,6 +1353,9 @@ def train_llama3():
      if (DP := getenv("DP", 1)) > 1:
        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
        batch = batch.shard(device, 0)
+      if (MP := getenv("MP", 1)) > 1:
+        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+        batch = batch.shard(device)
      logits:Tensor = model(batch[:, :-1], start_pos=0, temperature=math.nan)
      loss = logits.sparse_categorical_crossentropy(batch[:, 1:])
      loss.backward()
@@ -1391,11 +1392,12 @@ def train_llama3():
    t = time.perf_counter()
    GlobalCounters.reset()
    loss, lr = train_step(model, tokens, grad_acc)
+    loss = loss.float().item()
    # above as tqdm.write f-string
-    tqdm.write(f"{loss.item():.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
+    tqdm.write(f"{loss:.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
    if (fname:=getenv("LOSS_FILE", "")):
      with open(fname, "a") as f:
-        f.write(f"{i} {loss.item():.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
+        f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")

    if getenv("CKPT") and (i % 200 == 0 or i == 10):
      tqdm.write("saving checkpoint")
@@ -4,6 +4,8 @@ export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128

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

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

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

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

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

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

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

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

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

 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -5,9 +5,9 @@ set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." NV=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

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

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

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

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

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

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

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

 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -5,9 +5,9 @@ set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_red"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

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

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

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

 # model is bf16, 1.3B active, 6.9B total
 # M3 Max is 400 GB/s, so 400/2.6 = ~154 tok/s
@@ -71,8 +71,8 @@ def bbox_iou(box1, box2):
  # get the coordinates of the intersection rectangle
  inter_rect_x1 = np.maximum(b1_x1, b2_x1)
  inter_rect_y1 = np.maximum(b1_y1, b2_y1)
-  inter_rect_x2 = np.maximum(b1_x2, b2_x2)
-  inter_rect_y2 = np.maximum(b1_y2, b2_y2)
+  inter_rect_x2 = np.minimum(b1_x2, b2_x2)
+  inter_rect_y2 = np.minimum(b1_y2, b2_y2)
  #Intersection area
  inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, 99999) * np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, 99999)
  #Union Area
@@ -297,13 +297,13 @@ class Darknet:
          # Get the number of weights of batchnorm
          num_bn_biases = math.prod(bn.bias.shape)
          # Load weights
-          bn_biases = Tensor(weights[ptr:ptr + num_bn_biases])
+          bn_biases = Tensor(weights[ptr:ptr + num_bn_biases].astype(np.float32))
          ptr += num_bn_biases
-          bn_weights = Tensor(weights[ptr:ptr+num_bn_biases])
+          bn_weights = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
          ptr += num_bn_biases
-          bn_running_mean = Tensor(weights[ptr:ptr+num_bn_biases])
+          bn_running_mean = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
          ptr += num_bn_biases
-          bn_running_var = Tensor(weights[ptr:ptr+num_bn_biases])
+          bn_running_var = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
          ptr += num_bn_biases
          # Cast the loaded weights into dims of model weights
          bn_biases = bn_biases.reshape(shape=tuple(bn.bias.shape))
@@ -319,7 +319,7 @@ class Darknet:
          # load biases of the conv layer
          num_biases = math.prod(conv.bias.shape)
          # Load weights
-          conv_biases = Tensor(weights[ptr: ptr+num_biases])
+          conv_biases = Tensor(weights[ptr: ptr+num_biases].astype(np.float32))
          ptr += num_biases
          # Reshape
          conv_biases = conv_biases.reshape(shape=tuple(conv.bias.shape))
@@ -327,7 +327,7 @@ class Darknet:
          conv.bias = conv_biases
        # Load weighys for conv layers
        num_weights = math.prod(conv.weight.shape)
-        conv_weights = Tensor(weights[ptr:ptr+num_weights])
+        conv_weights = Tensor(weights[ptr:ptr+num_weights].astype(np.float32))
        ptr += num_weights
        conv_weights = conv_weights.reshape(shape=tuple(conv.weight.shape))
        conv.weight = conv_weights
@@ -371,7 +371,7 @@ class Darknet:
 if __name__ == "__main__":
  model = Darknet(fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg').read_bytes())
  print("Loading weights file (237MB). This might take a while…")
-  model.load_weights('https://pjreddie.com/media/files/yolov3.weights')
+  model.load_weights('https://github.com/shadiakiki1986/yolov3.weights/releases/download/3.0.1/yolov3.weights')
  if len(sys.argv) > 1:
    url = sys.argv[1]
  else:
@@ -1,5 +1,5 @@
 from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
-from tinygrad.opt.kernel import Ops, MemOp, UOp
+from tinygrad.codegen.opt.kernel import Ops, MemOp, UOp
 from tinygrad.uop.ops import BinaryOps, UnaryOps
 from tinygrad.dtype import DType, dtypes
 from tinygrad.helpers import DEBUG
@@ -3,7 +3,7 @@ from platform import system
 from typing import Tuple, Dict, List, Optional
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.opt.kernel import Ops, UOp
+from tinygrad.codegen.opt.kernel import Ops, UOp
 from tinygrad.helpers import CI
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage

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

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

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

@@ -5,9 +5,9 @@ from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
 from tinygrad.engine.realize import get_program

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

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

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

 sched = C.schedule()
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.device import CompilerOptions
 lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
 lin.to_program()
@@ -8,7 +8,6 @@ bert_train_params = {
  "BS": 96,
  "EVAL_BS": 96,
  "FUSE_ARANGE": 1,
-  "FUSE_ARANGE_UINT": 0,
  "BASEDIR": "/raid/datasets/wiki",
 }

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

 class MCTSNode:
@@ -181,13 +181,11 @@ class Transformer:
  def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):
    _bsz, seqlen = tokens.shape
    h = self.tok_embeddings(tokens)
-
-    self.freqs_cis = self.freqs_cis.cast(h.dtype).contiguous()
-    freqs_cis = self.freqs_cis[:, start_pos:start_pos+seqlen, :, :, :]
+    freqs_cis = self.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :]

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

    return sample(logits[:, -1, :].flatten(), temperature, top_k, top_p, alpha_f, alpha_p)
@@ -5,9 +5,9 @@ from tinygrad.nn import Linear
 from tinygrad.tensor import Tensor
 from tinygrad.nn.optim import Adam
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.opt.search import actions
+from tinygrad.codegen.opt.search import actions
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.helpers import getenv

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

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

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

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

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

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

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

 from extra.optimization.helpers import lin_to_feats, MAX_DIMS

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

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


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

 if __name__ == "__main__":
  ast_strs = load_worlds()
@@ -6,8 +6,8 @@ from tinygrad.helpers import getenv, BEAM
 from tinygrad.engine.jit import TinyJit
 from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item, get_program
 from tinygrad.renderer import ProgramSpec
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 import numpy as np

 def move_jit_captured_to_dev(captured, device="DSP"):
@@ -223,15 +223,18 @@ def max_unpool2d(self:torch.Tensor, indices:torch.Tensor, output_size):

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

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

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

@torch.library.impl("aten::convolution_overrideable", "privateuseone")
 def convolution_overrideable(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups):
@@ -368,6 +371,7 @@ from torch._decomp import get_decompositions
 decomps = [
  aten.native_batch_norm, aten.native_batch_norm_backward,
  aten.native_layer_norm_backward,
+  aten.linalg_cross,
  aten.addmm,
  aten.addcmul,
  aten.addcdiv,
@@ -135,7 +135,7 @@ class TestTorchBackend(unittest.TestCase):
    print(c.cpu())

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

@@ -203,6 +203,12 @@ class TestTorchBackend(unittest.TestCase):
    b = torch.linalg.det(a)
    np.testing.assert_equal(b.cpu().numpy(), 120.0)

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

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

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

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

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

-  from tinygrad.opt.kernel import Kernel, Opt, OptOps
+  from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
  from tinygrad.helpers import get_single_element
  GlobalCounters.reset()
  si = get_single_element(t.softmax(-1, dtype="half", _single_kernel=True).schedule())
@@ -1,8 +1,8 @@
 # ruff: noqa: E501
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.opt.search import bufs_from_lin
+from tinygrad.codegen.opt.search import bufs_from_lin
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
@@ -4,10 +4,10 @@ os.environ["VALIDATE_HCQ"]="1"

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

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

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

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

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

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

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

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

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

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

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

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

 class TestOnnxRunnerDtypes(unittest.TestCase):
  """
@@ -95,7 +95,7 @@ class TestOnnxRunnerDtypes(unittest.TestCase):
  External tensors (inputs) preserve their original dtype - user must ensure compatibility with device.
  """
  def _get_expected_dtype(self, onnx_dtype: int, is_input: bool):
-    true_dtype = data_types[onnx_dtype]
+    true_dtype = OnnxDataType(onnx_dtype).to_dtype()
    # inputs always preserve their true dtype.
    if is_input:
      return true_dtype
@@ -2,11 +2,11 @@
 import unittest

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

 from test.external.fuzz_linearizer import run_linearizer

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

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

 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.kernel import Opt, OptOps
-from tinygrad.opt.search import get_kernel_actions, bufs_from_lin
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.search import get_kernel_actions, bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing
 from tinygrad.uop.ops import UOp, Ops
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 # compare kernels created by HEAD against master
 import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools, functools, base64, codecs
+from dataclasses import replace
 from typing import Callable, Any

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

-def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
-  input_ast = ast.replace(arg=KernelInfo(opts_to_apply=p.applied_opts, name=p.name)) if ast.arg is None else ast
-  p2 = get_program(input_ast, renderer)
+def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
+  # NOTE: this always uses the opts_to_apply path
+  sink_arg = ast.arg or KernelInfo(opts_to_apply=p.applied_opts)
+  input_ast = ast.replace(arg=replace(sink_arg, name=p.name))
+  # if no renderer was provided, open the device to get it
+  if renderer is None: renderer = Device[p.device].renderer
+  p2 = get_program(input_ast, renderer=renderer)
  def to_str(ret:ProgramSpec) -> str:
    # PYTHON renderer pickles UOps, first unpickle and decode here
    if p.device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(ret.src))])
@@ -1,7 +1,7 @@
 from tinygrad import Device
 from tinygrad.helpers import getenv, DEBUG, BEAM
-from tinygrad.opt.search import beam_search, bufs_from_lin
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer

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

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

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

@@ -3,10 +3,11 @@ 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.opt.kernel import Opt, OptOps, Kernel, KernelOptError
+from tinygrad.codegen.opt.kernel import Opt, OptOps, Kernel, KernelOptError
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
-from tinygrad.opt.search import get_kernel_actions
+from tinygrad.codegen.opt.search import get_kernel_actions
 from tinygrad.uop.ops import Ops
+from tinygrad.codegen import apply_rewrites, rewrites_for_views

 class TestArange(unittest.TestCase):
  def _get_flops(self, N, opts=None):
@@ -14,10 +15,7 @@ class TestArange(unittest.TestCase):
    tt = Tensor.arange(N)
    sched = tt.schedule()
    self.assertEqual(len(sched), 1)
-    k = Kernel(sched[-1].ast)
-    if opts is not None:
-      for o in opts: k.apply_opt(o)
-    p = get_program(k.get_optimized_ast(), k.opts)
+    p = get_program(sched[-1].ast, opts=opts)
    print(p.name)
    #print(p.src)
    ExecItem(CompiledRunner(p), [tt.uop.buffer]).run()
@@ -52,11 +50,11 @@ class TestArange(unittest.TestCase):
    def test_complexity_w_local_and_padto(self): return self.test_complexity([Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.PADTO, axis=1, arg=32)])

  def test_all_opts(self, opts=None, exclude=None):
-    k = Kernel(Tensor.arange(256).schedule()[-1].ast)
+    k = Kernel(apply_rewrites(Tensor.arange(256).schedule()[-1].ast, rewrites_for_views))
    if opts is not None:
      for o in opts: k.apply_opt(o)
    all_opts_256 = [kk.applied_opts for kk in get_kernel_actions(k, include_0=False).values()]
-    k = Kernel(Tensor.arange(2560).schedule()[-1].ast)
+    k = Kernel(apply_rewrites(Tensor.arange(2560).schedule()[-1].ast, rewrites_for_views))
    if opts is not None:
      for o in opts: k.apply_opt(o)
    all_opts_2560 = [kk.applied_opts for kk in get_kernel_actions(k, include_0=False).values()]
@@ -139,10 +139,9 @@ class TestBitcastConstFolding(unittest.TestCase):
 class TestIndexingConstFolding(unittest.TestCase):
  def test_scalar_index(self):
    t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
-    # TODO: fold these
-    _check_ast_count(2, t[:,:,Tensor(1),:])
-    _check_ast_count(2, t[:,:,Tensor(1)+2,:])
-    _check_ast_count(2, t[:,:,Tensor(1),Tensor(0)])
+    _check_ast_count(1, t[:,:,Tensor(1),:])
+    _check_ast_count(1, t[:,:,Tensor(1)+2,:])
+    _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])

  @unittest.expectedFailure
  def test_const_tensor_index(self):
@@ -291,17 +290,12 @@ class TestMultiConstFolding(unittest.TestCase):
    np.testing.assert_equal((t + zero).numpy(), np.arange(16))
    np.testing.assert_equal((t * zero).numpy(), [0] * 16)
    np.testing.assert_equal((t * one).numpy(), np.arange(16))
-
-  def test_multi_todo_pow(self):
-    ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
-    t = Tensor.arange(16).float().to(ds).realize()
-    zero = Tensor.zeros(16).to(ds).realize()
-    one = Tensor.ones(16).to(ds).realize()
-
-    # TODO: fix pow folding
    _check_ast_count(0, t ** zero)
    _check_ast_count(0, t ** one)
    _check_ast_count(0, one ** t)
+    np.testing.assert_equal((t ** zero).numpy(), [1] * 16)
+    np.testing.assert_equal((t ** one).numpy(), np.arange(16))
+    np.testing.assert_equal((one ** t).numpy(), [1] * 16)

 class TestTautologicalCompare(unittest.TestCase):
  # without const folding, these would have triggered -Wtautological-compare in clang
@@ -0,0 +1,21 @@
+import unittest, io
+from tinygrad import Tensor, dtypes
+from contextlib import redirect_stdout
+from tinygrad.device import Device
+from tinygrad.helpers import OSX
+from tinygrad.engine.realize import get_program
+
+class TestDisassembly(unittest.TestCase):
+  # TODO: fails on llvm. llvm.LLVMGetHostCPUName() returns "generic"
+  @unittest.skipUnless(Device.DEFAULT in ("CPU",) and OSX, "m series cpus support fp16 arithmetic")
+  def test_float16_alu(self):
+    c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
+    s = c.schedule()[-1]
+    p = get_program(s.ast, Device[Device.DEFAULT].renderer)
+    lib = Device[Device.DEFAULT].compiler.compile(p.src)
+    out = io.StringIO()
+    with redirect_stdout(out): Device[Device.DEFAULT].compiler.disassemble(lib)
+    assert "fcvt" not in out.getvalue()
+
+if __name__ == "__main__":
+  unittest.main()
@@ -62,7 +62,6 @@ class TestNaNEdgeCases(unittest.TestCase):
 class TestEmptyTensorEdgeCases(unittest.TestCase):
  # we don't need more of these

-  @unittest.expectedFailure
  def test_sort_empty(self):
    # Sorting an empty tensor works in PyTorch and should return empty
    # values and indices. tinygrad raises an error instead.
@@ -219,7 +218,6 @@ class TestAssignIssues(unittest.TestCase):
    t.shrink(((1, 3), (1, 3))).assign(Tensor.ones(2, 2))
    np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())

-  @unittest.expectedFailure
  def test_assign_broadcast(self):
    # broadcasting during assign should behave like PyTorch
    torch_tensor = torch.zeros(3, 5)
@@ -258,12 +256,11 @@ class TestEdgeCases(unittest.TestCase):
    out = Tensor(arr).pad((1, -1, 1, -1), mode='circular')
    np.testing.assert_equal(out.numpy(), torch_out.numpy())

-  @unittest.expectedFailure
  def test_arange_float_step(self):
    # float steps should match PyTorch exactly
    torch_out = torch.arange(0, 2, 0.3).numpy()
    out = Tensor.arange(0, 2, 0.3).numpy()
-    np.testing.assert_allclose(out, torch_out)
+    np.testing.assert_allclose(out, torch_out, atol=1e-7)

  @unittest.skip("this is flaky")
  @unittest.expectedFailure
@@ -6,7 +6,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps
 from tinygrad import Variable

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

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

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

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

  def test_speed_copy_bandwidth(self):
    if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -718,8 +718,6 @@ class TestJitGraphSplit(unittest.TestCase):
  def ji_xfer(self): return {"type": "xfer"}

  def test_jit_split_simple(self):
-    if Device.DEFAULT == "REMOTE": raise unittest.SkipTest("REMOTE gpu is broken")
-
    @TinyJit
    def f(inp):
      op0 = self.compute(Device.DEFAULT, inp)
@@ -792,9 +790,9 @@ class TestJitGraphSplit(unittest.TestCase):
      multigraph=[self.ji_graph(5)],
      hcqgraph=[self.ji_graph(5)])

-  @unittest.skip("flaky")
  def test_jit_multidev_xfer(self):
    if Device.DEFAULT in {"CPU", "LLVM"}: raise unittest.SkipTest("CPU/LLVM is not a valid default device for this test (zero-copies)")
+    if Device.DEFAULT == "METAL" or REAL_DEV == "METAL": raise unittest.SkipTest("Metal is flaky, with multidevice (same as metal llama 4gpu?)")

    try: Device[f"{Device.DEFAULT}:1"]
    except Exception: raise unittest.SkipTest("no multidevice")
@@ -819,7 +817,6 @@ class TestJitGraphSplit(unittest.TestCase):
  @unittest.skipIf(getenv("MOCKGPU"), "MockGPU does not support parallel copies")
  def test_jit_multidev_copy(self):
    if Device.DEFAULT in {"CPU", "LLVM"}: raise unittest.SkipTest("CPU/LLVM is not a valid default device for this test (zero-copies)")
-    if Device.DEFAULT == "REMOTE": raise unittest.SkipTest("REMOTE gpu is broken")

    @TinyJit
    def f(inp):
@@ -62,5 +62,15 @@ class TestLinAlg(unittest.TestCase):
      orthogonality_helper(Q)
      reconstruction_helper([Q,R],a)

+  def test_newton_schulz(self):
+    coefficients = [(2, -1.5, 0.5), (2.0, -1.4, 0.2, 0.2)]#these params map to the sign function
+    sizes = [(2,2), (3,2), (2,3), (2,2,2)]
+    for coefs in coefficients:
+      for size in sizes:
+        a = Tensor.randn(size)
+        b = Tensor.newton_schulz(a, steps=20, params=coefs, eps=0.0)
+        # ns(A) = U @ Vt -> (U @ Vt) @ (U @ Vt)t = I
+        orthogonality_helper(b if size[-1] > size[-2] else b.transpose(-2, -1), tolerance=1e-1)
+
 if __name__ == "__main__":
  unittest.main()
@@ -2,7 +2,7 @@ import numpy as np
 import unittest
 from dataclasses import replace

-from tinygrad.opt.kernel import Opt, OptOps, KernelOptError, Kernel, AxisType
+from tinygrad.codegen.opt.kernel import Opt, OptOps, KernelOptError, Kernel, AxisType
 from tinygrad.codegen.gpudims import get_grouped_dims
 from tinygrad.uop.ops import UOp, Ops, GroupOp, KernelInfo
 from tinygrad.device import Device, Buffer, is_dtype_supported
@@ -10,9 +10,12 @@ 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, get_program
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import prod, Context, getenv, CI, flatten, dedup, AMX, AMD_LLVM
 from tinygrad.dtype import DType, dtypes, AddrSpace
+from tinygrad.codegen import apply_rewrites, rewrites_for_views
+
+def push_views(ast): return apply_rewrites(ast, rewrites_for_views)

 def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
  if isinstance(r, Tensor): r = [r]
@@ -22,7 +25,7 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
  # now all input buffers in s[-1] should be realized
  # create fresh buffers for the outputs
  bufs = [Buffer((x).device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)]
-  return s[-1].ast, bufs
+  return push_views(s[-1].ast), bufs

 def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0, use_tensor_cores:int=1):
  a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
@@ -121,7 +124,7 @@ class TestLinearizer(unittest.TestCase):
    with Context(FUSE_ARANGE=1):
      sink = dataset[idxs].contiguous().kernelize().uop.base.src[1].arg.ast
    real_index = dataset.numpy()[idxs.numpy()].reshape(4, 256, 1, 1)
-    helper_linearizer_ast(sink, [dataset, idxs], wanna_output=[real_index])
+    helper_linearizer_ast(push_views(sink), [dataset, idxs], wanna_output=[real_index])

  def test_two_nested_range(self):
    a = Tensor.randn(2, ).realize()
@@ -214,9 +217,7 @@ class TestLinearizer(unittest.TestCase):
    # these are of size 3 to avoid float4 coalesce
    r = a[:-1] + a[1:]

-    k = Kernel(r.schedule()[-1].ast)
-    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
-    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
    num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
    assert num_loads <= 4, "more load uops than needed"
    assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
@@ -227,9 +228,7 @@ class TestLinearizer(unittest.TestCase):
    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
    r = a.expand([2]) + b.expand([2])

-    k = Kernel(r.schedule()[-1].ast)
-    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
-    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
    assert num_ops <= 1, "more alu uops than needed"

@@ -238,10 +237,7 @@ class TestLinearizer(unittest.TestCase):
    x, w = Tensor.randn((1,1,3)).realize(), Tensor.randn((1,1,2)).realize()
    r = Tensor.conv2d(x,w,padding=1).relu()

-    k = Kernel(r.schedule()[-1].ast)
-    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
-    k.apply_opt(Opt(op=OptOps.UNROLL, axis=0, arg=0))
-    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
    accs = [u for u in uops if u.op is Ops.DEFINE_REG]
    stores = [u for u in uops if u.op is Ops.STORE]
    assert len(accs) == 0  # it's removed now
@@ -253,9 +249,7 @@ class TestLinearizer(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
  def test_upcast_with_locals_cpu(self):
    out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
-    k = Kernel(out.schedule()[-1].ast)
-    k.apply_opt(Opt(OptOps.LOCAL, axis=0, arg=4))
-    prg = get_program(k.get_optimized_ast(), k.opts)
+    prg = get_program(out.schedule()[-1].ast, opts=[Opt(OptOps.LOCAL, axis=0, arg=4)]).uops
    self.assertEqual(len(prg.src.split("for")), 5)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -265,10 +259,8 @@ class TestLinearizer(unittest.TestCase):
  def test_upcast_with_locals(self):
    x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
    r = (x@y).relu()
-    realized_ast = r.schedule()[-1].ast
    opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
-    program = get_program(realized_ast, Device[Device.DEFAULT].renderer)
+    program = get_program(r.schedule()[-1].ast, opts=opts_to_apply)

    stores = [u for u in program.uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]

@@ -282,10 +274,7 @@ class TestLinearizer(unittest.TestCase):
  def test_zero_fold(self):
    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
    r = Tensor.stack(a, b)
-
-    k = Kernel(r.schedule()[-1].ast)
-    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
-    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
    assert num_ops == 0, "more alu uops than needed"

@@ -295,16 +284,14 @@ class TestLinearizer(unittest.TestCase):
      if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype):
        a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
        realized_ast = a.schedule()[-1].ast
-        realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple()))
-        program = get_program(realized_ast, Device[Device.DEFAULT].renderer)
+        program = get_program(realized_ast, opts=[])
        local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
        assert local[0].dtype.base == acc_dtype

  def test_arg_acc_dtype(self):
    def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
      realized_ast = c.schedule()[-1].ast
-      realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple()))
-      program = get_program(realized_ast, Device[Device.DEFAULT].renderer)
+      program = get_program(realized_ast, opts=[])
      local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
      self.assertEqual(local[0].dtype.base, expected_dtype)

@@ -341,7 +328,7 @@ class TestLinearizer(unittest.TestCase):
      a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in)
      r = a.matmul(b, dtype=tc.dtype_out)
      sched = r.schedule()
-      realized_ast = sched[-1].ast
+      realized_ast = push_views(sched[-1].ast)
      kernel = Kernel(realized_ast)
      kernel.apply_tensor_cores(1, axis=0, tc_select=-1, tc_opt=2)
      prg = get_program(kernel.get_optimized_ast(), kernel.opts)
@@ -423,7 +410,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.opt.search import get_kernel_actions
+      from tinygrad.codegen.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])
@@ -762,11 +749,7 @@ class TestFloat4(unittest.TestCase):
    c = a + b

    s = c.schedule()[0]
-    k = Kernel(s.ast)
-    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=4))
-    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=2))
-    uops = get_program(k.get_optimized_ast(), k.opts).uops
-
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]).uops
    assert TestFloat4.count_float4(uops) == (4, 2)

  @unittest.skipUnless(Device.DEFAULT in {"CPU", "LLVM"} and AMX, "Only CPU with AMX upcasts float up to size 16")
@@ -777,10 +760,7 @@ class TestFloat4(unittest.TestCase):
      c = a + b

      s = c.schedule()[0]
-      k = Kernel(s.ast)
-      k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=4))
-      k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=shift))
-      return get_program(k.get_optimized_ast(), k.opts).uops
+      return get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]).uops

    sizes = [12, 8, 16]
    shifts = [3, 2, 4]
@@ -810,10 +790,7 @@ class TestFloat4(unittest.TestCase):
    c = a + b

    s = c.schedule()[0]
-    k = Kernel(s.ast)
-    k.apply_opt(Opt(op=OptOps.UPCAST, axis=1, arg=4))
-    k.apply_opt(Opt(op=OptOps.UPCAST, axis=1, arg=2))
-    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]).uops

    assert TestFloat4.count_float4(uops) == (0, 2)

@@ -825,10 +802,7 @@ class TestFloat4(unittest.TestCase):
      c = a + b

      s = c.schedule()[0]
-      k = Kernel(s.ast)
-      k.shift_to(1, 4, AxisType.UPCAST)  # manual trigger float4 dim
-      k.shift_to(1, shift, AxisType.UPCAST, insert_at=k.shape_len-1)
-      return get_program(k.get_optimized_ast(), k.opts).uops
+      return get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]).uops

    sizes = [13, 9, 17]
    shifts = [3, 2, 4]
@@ -846,9 +820,7 @@ class TestFloat4(unittest.TestCase):
    # float4 should be emitted (the reduce axis of size 4 is the float4 axis here)

    s = c.schedule()[0]
-    k = Kernel(s.ast)
-    k.apply_opt(Opt(op=OptOps.UNROLL, axis=0, arg=4))
-    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UNROLL, axis=0, arg=4)]).uops

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

@@ -862,10 +834,7 @@ class TestFloat4(unittest.TestCase):
    # UPDATE: now we do this fusion

    s = c.schedule()[0]
-    k = Kernel(s.ast)
-    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
-    k.apply_opt(Opt(op=OptOps.UNROLL, axis=0, arg=0))
-    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops

    assert TestFloat4.count_float4(uops) in {(0,1), (1,1)}

@@ -878,9 +847,7 @@ class TestFloat4(unittest.TestCase):
    # since the top axis is not contiguous.

    s = c.schedule()[0]
-    k = Kernel(s.ast)
-    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=4))
-    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops

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

@@ -892,9 +859,7 @@ class TestFloat4(unittest.TestCase):
    # should float4 b but not a

    s = c.schedule()[0]
-    k = Kernel(s.ast)
-    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=4))
-    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops

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

@@ -983,7 +948,7 @@ class TestHandCodedOpts(unittest.TestCase):
    layer_2 = Tensor.cat(layer_1.unsqueeze(0), Tensor.empty(6, 20))

    s = layer_2.schedule()[-1]
-    k = Kernel(s.ast)
+    k = Kernel(push_views(s.ast))
    k.apply_opts(hand_coded_optimizations(k))
    assert len(k.bufs) == 6  # make sure all ops are done in one kernel
    # masked upcast should upcast masked axis of size 7
@@ -996,7 +961,7 @@ class TestHandCodedOpts(unittest.TestCase):
    monster = Tensor.stack(*[Tensor.stack(*[Tensor.empty(16) for _ in range(6)]) for _ in range(6)])

    s = monster.schedule()[-1]
-    k = Kernel(s.ast)
+    k = Kernel(push_views(s.ast))
    k.apply_opts(hand_coded_optimizations(k))
    assert len(k.bufs) == 37  # make sure all ops are done in one kernel
    # should upcast the two Tensor.stacks
@@ -1012,7 +977,7 @@ class TestHandCodedOpts(unittest.TestCase):
      wino_schedule = out.schedule()
      # collect upcasts of tile transform kernels
      for i, si in enumerate(wino_schedule):
-        k = Kernel(si.ast)
+        k = Kernel(push_views(si.ast))
        k.apply_opts(hand_coded_optimizations(k))
        if k.reduceop is not None: continue  # not a tile transform kernel (there is a gemm reduce kernel)
        if len(k.bufs) < 22: continue  # not a tile transform kernel (there's a permute kernel at the end)
@@ -1024,7 +989,7 @@ class TestHandCodedOpts(unittest.TestCase):

      backward_schedule = Tensor.schedule(x.grad, w.grad)
      for si in backward_schedule:
-        k = Kernel(si.ast)
+        k = Kernel(push_views(si.ast))
        k.apply_opts(hand_coded_optimizations(k))
        if len(k.bufs) < 20: continue  # not a tile transform kernel
        # heuristic number to make sure that at least some upcasts but not too many upcasts are being done
@@ -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.opt.search import Opt, OptOps
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.search import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.engine.realize import get_program

 class TestLinearizerDumb(unittest.TestCase):
@@ -1,8 +1,8 @@
 # ruff: noqa: E501
 import unittest
 from tinygrad import dtypes
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.search import Opt, OptOps, bufs_from_lin
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.search import Opt, OptOps, bufs_from_lin
 from extra.optimization.helpers import time_linearizer

 # stuff needed to unpack a kernel
@@ -401,7 +401,7 @@ class TestNN(unittest.TestCase):
      torch_z = torch_layer(torch_x)
      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)

-  def test_embedding_one_kernel(self, ops=41410, kcount=3):
+  def test_embedding_one_kernel(self, ops=612000, kcount=2):
    GlobalCounters.reset()
    layer = Embedding(20, 30)
    layer.weight = Tensor.zeros_like(layer.weight).contiguous()
@@ -409,7 +409,7 @@ class TestNN(unittest.TestCase):
                [12, 19, 8, 1]])
    result = layer(a)
    schedule = result.schedule()
-    self.assertEqual(kcount, len([item for item in schedule if item.ast.op is Ops.SINK]), "first run realizes weight and embedding")
+    self.assertEqual(len([item for item in schedule if item.ast.op is Ops.SINK]), kcount, "first run realizes weight and embedding")
    run_schedule(schedule)

    b = Tensor([[1, 2, 3],
@@ -440,6 +440,13 @@ class TestNN(unittest.TestCase):
      result = layer(a)
      self.assertEqual(result.shape, shp + (embed_size,))

+  def test_embedding_regression(self):
+    # used to fail bounds check
+    with Context(FUSE_ARANGE=1):
+      embedding = Embedding(100, 1024)
+      input_ids = Tensor.empty(16, 16)
+      embedding(input_ids).realize()
+
  def test_load_state_dict(self):
    layer = Conv2d(3, 5, kernel_size=3)

@@ -699,6 +699,14 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, lambda x: x**-0.3, vals=[[0.0]])
    helper_test_op(None, lambda x: x**-1.0, vals=[[-1.0, 0.0, 1.0]])

+  def test_int_pow_const_int(self):
+    helper_test_op(None, lambda x: x**0, vals=[[-2,0,2]], forward_only=True, atol=0)
+    helper_test_op(None, lambda x: x**1, vals=[[-2,0,2]], forward_only=True, atol=0)
+    helper_test_op(None, lambda x: x**2, vals=[[-2,0,2]], forward_only=True, atol=0)
+    helper_test_op(None, lambda x: x**7, vals=[[11,12,13]], forward_only=True, atol=0)
+    helper_test_op(None, lambda x: x**29, vals=[[-2,0,2]], forward_only=True, atol=0)
+    self.helper_test_exception(None, lambda x: x**-2, vals=[[-2,0,2]], forward_only=True, expected=RuntimeError)
+
  @unittest.skip("not supported")
  def test_pow_int(self):
    def _test(base, exponent): helper_test_op(None, lambda x,y: x**y, vals=[base, exponent], forward_only=True)
@@ -957,8 +965,9 @@ class TestOps(unittest.TestCase):
    helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
    helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3), lambda t: Tensor.softplus(t, beta=3), grad_atol=1e-6)
    helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=1/3), lambda t: Tensor.softplus(t, beta=1/3), grad_atol=1e-6)
-    # # TODO: support threshold and enable this
-    # helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=300, high=400)
+    helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3, threshold=0.5),
+                              lambda t: Tensor.softplus(t, beta=3, threshold=0.5), grad_atol=1e-6)
+    helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=300, high=400)
    helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=-400, high=-300)
    helper_test_op([()], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)

@@ -1092,6 +1101,9 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, lambda x: x.type(torch.int32).argmin().type(torch.int32), lambda x: x.argmin(), forward_only=True, vals=[[True, False]])

  def test_sort(self):
+    for shape in [(0,), (0,5), (1,), (1,5)]:
+      helper_test_op([shape], lambda x: x.sort(0).values, lambda x: x.sort(0)[0], forward_only=True)
+      helper_test_op([shape], lambda x: x.sort(0).indices.type(torch.int32), lambda x: x.sort(0)[1], forward_only=True)
    for dim in [-1, 0, 1]:
      for descending in [True, False]:
        helper_test_op([(8,8,6)], lambda x: x.sort(dim, descending).values, lambda x: x.sort(dim, descending)[0], forward_only=True)
@@ -2738,10 +2750,10 @@ class TestOps(unittest.TestCase):
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[[0]]], lambda x: x[[[0]]])
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[0],b,c,d,:], lambda x: x[[0],j,k,o,:])
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[[[0]]],b,c,d,[[1]]], lambda x: x[[[[0]]],j,k,o,[[1]]])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[1,0],b,c,d,:], lambda x: x[[1,0],j,k,o,:])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[1,0,-1],b,c,d,:], lambda x: x[[1,0,-1],j,k,o,:])
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,b,c,[1,2,3],...], lambda x: x[i,j,k,[1,2,3],...])
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,b,c,[[1],[2],[3]],...], lambda x: x[i,j,k,[[1],[2],[3]],...])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,[2,1,0],c,[2,1,0],e], lambda x: x[i,[2,1,0],k,[2,1,0],p])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,[2,1,0],c,[-2,1,0],e], lambda x: x[i,[2,1,0],k,[-2,1,0],p])

  def test_slice_fancy_indexing_tuple_indices(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
@@ -2764,10 +2776,14 @@ class TestOps(unittest.TestCase):
    a = Tensor.ones(10,11,12)
    # tensors used as indices must be int tensors
    with self.assertRaises(IndexError): a[Tensor(1.1)]
-    with self.assertRaises(IndexError): a[Tensor([True, True])]
+    with self.assertRaises(IndexError): a[[1.1]]
+    with self.assertRaises(IndexError): a[Tensor([True, False])]
+    with self.assertRaises(IndexError): a[[True, False]]
    # shape mismatch, cannot broadcast. either exception is okay
    with self.assertRaises((IndexError, ValueError)): a[Tensor.randint(3,1,1,1), Tensor.randint(1,4,1,1), Tensor.randint(2,4,4,1)]
    with self.assertRaises((IndexError, ValueError)): a[Tensor.randint(3,1,1,1), Tensor.randint(1,4,1,1,1)]
+    # this is fine
+    helper_test_op([(5, 6)], lambda x: x[[True, False, 2]])

  def test_gather(self):
    # indices cannot have gradient
@@ -2,7 +2,7 @@ import numpy as np
 import unittest
 from tinygrad import Tensor
 from tinygrad.helpers import get_single_element
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program

 class TestOptGemm(unittest.TestCase):
@@ -17,9 +17,7 @@ class TestOptGemm(unittest.TestCase):
    t = self.a.T @ self.b.T
    # TODO: this should be a generic test helper
    si = get_single_element(t.schedule())
-    k = Kernel(si.ast)
-    k.apply_opts(opts)
-    run = CompiledRunner(get_program(k.get_optimized_ast(), k.opts))
+    run = CompiledRunner(get_program(si.ast, opts=opts))
    ExecItem(run, si.bufs).run()
    test = si.bufs[0].numpy().reshape(self.res.shape)
    np.testing.assert_allclose(self.res, test, atol=1e-4)
@@ -2,9 +2,10 @@ import numpy as np
 import torch
 import unittest
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.nn.optim import Adam, SGD, AdamW
+from tinygrad.nn.optim import Adam, SGD, AdamW, Muon
 from tinygrad.helpers import CI
 from tinygrad.device import is_dtype_supported
+from extra.torch_muon import SingleDeviceMuon as TorchMuon

 np.random.seed(1337)
 x_init = np.random.randn(1,4).astype(np.float32)
@@ -57,9 +58,12 @@ class TestOptim(unittest.TestCase):
  def _test_sgd(self, steps, opts, atol, rtol): self._test_optim(SGD, torch.optim.SGD, steps, opts, atol, rtol)
  def _test_adam(self, steps, opts, atol, rtol): self._test_optim(Adam, torch.optim.Adam, steps, opts, atol, rtol)
  def _test_adamw(self, steps, opts, atol, rtol): self._test_optim(AdamW, torch.optim.AdamW, steps, opts, atol, rtol)
+  #TODO: use torch.muon when it comes out
+  def _test_muon(self, steps, opts, atol, rtol): self._test_optim(Muon, TorchMuon, steps, opts, atol, rtol)

  def test_multistep_sgd_high_lr_teeny(self): self._test_sgd(2, {'lr': 1.1, 'teeny': True}, 1e-6, 1e-5)
  def test_multistep_adam_high_lr_teeny(self): self._test_adam(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)
+  def test_multistep_muon_high_lr_teeny(self): self._test_muon(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)

  def test_sgd(self): self._test_sgd(1, {'lr': 0.001}, 1e-6, 0)
  def test_sgd_high_lr(self): self._test_sgd(1, {'lr': 10}, 1e-6, 1e-5)
@@ -83,6 +87,28 @@ class TestOptim(unittest.TestCase):
  def test_multistep_sgd_high_lr_nesterov_momentum_wd(self):
    self._test_sgd(10, {'lr': 9, 'momentum': 0.9, 'nesterov': True, 'weight_decay': 0.1}, 1e-5, 3e-4)

+  def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-6, 0)
+  def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
+  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-6, 0)
+  def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 3e-4)
+
+  # NOTE: momentum set to 0.95 by default, nesterov set to True by default
+  def test_multistep_muon_momentum_wd(self): self._test_muon(10, {'lr': 0.001, 'weight_decay': 0.01}, 1e-5, 0)
+  # ns defaults are numerically unstable, but it is tolerable in real training (see nsteps/nparam tests)
+  def test_multistep_muon_high_lr_momentum_wd(self): self._test_muon(10, {'lr': 10, 'weight_decay': 0.01}, 1e-1, 3e-4)
+  def test_multistep_muon_no_nesterov_momentum(self): self._test_muon(10, {'lr': 0.001, 'nesterov': False}, 1e-5, 0)
+  def test_multistep_muon_high_lr_no_nesterov_momentum(self): self._test_muon(10, {'lr': 10, 'nesterov': False}, 0.5e-1, 1e-1)
+
+  def test_muon_ns_steps(self): self._test_muon(1, {'lr': 0.001, 'ns_steps': 3}, 1e-6, 0)
+  def test_muon_high_lr_ns_steps(self): self._test_muon(1, {'lr': 10, 'ns_steps': 3}, 1e-5, 3e-4)
+  def test_muon_ns_params(self): self._test_muon(1, {'lr': 0.001,'ns_params': (2.0,-1.5,0.5)}, 1e-6, 0)
+  def test_muon_high_lr_ns_params(self): self._test_muon(1, {'lr': 10,'ns_params': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+
+  def test_muon_momentum_wd_ns_steps_ns_params(self):
+    self._test_muon(10, {'lr': 0.001, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_params': (2.0,-1.5,0.5)}, 1e-5, 0)
+  def test_multistep_muon_high_lr_momentum_wd_ns_steps_ns_params(self):
+    self._test_muon(10, {'lr': 10, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_params': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+
  def test_adam(self): self._test_adam(1, {'lr': 0.001}, 1e-5, 0)
  def test_adam_high_lr(self): self._test_adam(1, {'lr': 10}, 1e-4, 1e-4)
  def test_adamw(self): self._test_adamw(1, {'lr': 0.001}, 1e-5, 0)
@@ -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.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item, get_program
-from tinygrad.opt.search import bufs_from_lin
+from tinygrad.codegen.opt.search import bufs_from_lin
 from tinygrad.shape.shapetracker import ShapeTracker, View # noqa: F401 # pylint: disable=unused-import

 N = 512
@@ -40,10 +40,7 @@ def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias=

 def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
  si = out.schedule()[-1]
-  k = Kernel(si.ast, opts=Device[Device.DEFAULT].renderer)
-  #opts = [Opt(op=OptOps.UPCAST, axis=0, arg=128)] #, Opt(op=OptOps.UNROLL, axis=0, arg=4)]
-  k.apply_opts(opts)
-  prg = get_program(k.get_optimized_ast(), k.opts)
+  prg = get_program(si.ast, opts=opts)
  if replace_src is not None:
    old_name = prg.src.split("__attribute__((noinline)) void ")[1].split("(")[0]
    prg = replace(prg, src=replace_src + "/* DSP boilerplate */" + prg.src.split("/* DSP boilerplate */")[1].replace(old_name, "fxn"))
@@ -297,10 +294,7 @@ class TestDSPCache(unittest.TestCase):
            x41,)),)),)),))""")
    opts = [Opt(op=OptOps.UNROLL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=1, arg=32), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
    with Context(DEVECTORIZE=0, QUANTIZE=1):
-      k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
-      k.apply_opts(opts)
-      prg = get_program(k.get_optimized_ast(), k.opts)
-      #print(prg.src)
+      prg = get_program(ast, opts=opts)

    new_src = """
 typedef int int32 __attribute__((aligned(128),vector_size(128)));
@@ -362,7 +356,7 @@ __attribute__((noinline)) void r_196_32_4_24_8(unsigned char* restrict __attribu
    prg = replace(prg, src=new_src+prg.src.split("/* DSP boilerplate */ ")[1])
    rt = CompiledRunner(prg)
    #Device.default.compiler.disassemble(rt.lib)
-    ei = ExecItem(rt, bufs_from_lin(k))
+    ei = ExecItem(rt, bufs_from_lin(Kernel(ast)))
    tm = ei.run(wait=True)
    print(f"final time {tm*1e6:.2f} us")

@@ -6,7 +6,7 @@ import unittest
 import numpy as np
 import functools
 from typing import cast
-from hypothesis import assume, given, strategies as strat
+from hypothesis import assume, given, settings, strategies as strat

 from tinygrad import nn, dtypes, Device, Tensor
 from tinygrad.device import is_dtype_supported
@@ -16,7 +16,7 @@ from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewr
 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.schedule.kernelize import merge_views, get_kernelize_map, Kernel
-from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars
+from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule

 class KernelCountException(Exception): pass
@@ -70,7 +70,7 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float, **kwargs):
 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):
+  def test_arange_avgpool2d(self, kcount=1):
    x = Tensor.arange(25).reshape(1,1,5,5).cast(dtypes.float32)
    t = x.avg_pool2d(padding=1)
    sched = t.schedule()
@@ -151,6 +151,7 @@ class TestSchedule(unittest.TestCase):
      self.assertEqual(root.item(), sum(range(N)))

  @given(strat.sampled_from(range(2,4)), strat.sampled_from(range(2,4)), strat.sampled_from(range(0,4)), strat.sampled_from(range(0,4)))
+  @settings(deadline=None)
  def test_indexing_scalars(self, x, y, a, b):
    assume(a<x and b<y)
    X = Tensor.randn(x, y).realize()
@@ -1027,14 +1028,14 @@ class TestSchedule(unittest.TestCase):
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 32).realize()
    out = x.argmin(-1)
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_equal(out.numpy(), x.numpy().argmin(axis=-1))

  def test_argmax_multireduce_fusion(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 32).realize()
    out = x.argmax(-1)
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_equal(out.numpy(), x.numpy().argmax(axis=-1))

  def test_scaled_dot_product_attention_multireduce_fusion(self):
@@ -1353,8 +1354,7 @@ class TestSchedule(unittest.TestCase):
    r = a.sum(0) + 6
    b = r.sum(0) * 4
    c = r.sum(1) * 2
-    schedule = check_schedule([b, c], 3)
-    self.assertIs(store_val(schedule[0]).op, Ops.ADD)
+    check_schedule([b, c], 3)

  def test_multireduce_simple_chase(self):
    Tensor.manual_seed(0)
@@ -1376,8 +1376,7 @@ class TestSchedule(unittest.TestCase):
    r = a.sum(2) + b
    d = r.T * 4
    e = r * d
-    schedule = check_schedule([d, e], 3)
-    self.assertIs(store_val(schedule[0]).op, Ops.ADD)
+    check_schedule([d, e], 3)

  def test_multireduce_push_permute_chase(self):
    Tensor.manual_seed(0)
@@ -1387,7 +1386,6 @@ class TestSchedule(unittest.TestCase):
    d = r.T * 4
    e = r * (d + a).sum(2)
    schedule = check_schedule([d, e], 3) # make sure it doesn't fuse
-    self.assertIs(store_val(schedule[0]).op, Ops.ADD)
    run_schedule(schedule)
    np.testing.assert_allclose(d.numpy(), (a.numpy().sum(2) + b.numpy()).T * 4, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(e.numpy(), (a.numpy().sum(2) + b.numpy()) * (d.numpy() + a.numpy()).sum(2), atol=1e-4, rtol=1e-4)
@@ -1398,8 +1396,7 @@ class TestSchedule(unittest.TestCase):
    c = Tensor.empty(16, )
    r = a.sum(1) + c
    d = r[:4] * b
-    schedule = check_schedule(d, 2)
-    self.assertIs(store_val(schedule[0]).op, Ops.ADD)
+    check_schedule(d, 2)

  def test_multireduce_push_shrink_chase(self):
    Tensor.manual_seed(0)
@@ -1411,15 +1408,13 @@ class TestSchedule(unittest.TestCase):
    out = r[:4] * b + d.sum(1)[:4]
    # schedule = check_schedule(out, 2)
    schedule = check_schedule(out, 3)
-    self.assertIs(store_val(schedule[0]).op, Ops.ADD)
    run_schedule(schedule)
    np.testing.assert_allclose(out.numpy(), (a.numpy().sum(1) + c.numpy())[:4] * b.numpy() + d.numpy().sum(1)[:4], atol=1e-4, rtol=1e-4)

  def test_midreduce_nochase(self):
    a = Tensor.empty(16, 16)
    b = (a.sum(0) + a.max(1)) + 2
-    schedule = check_schedule(b, 2)
-    self.assertIs(store_val(schedule[0]).op, Ops.REDUCE_AXIS)
+    check_schedule(b, 2)

  def test_multireduce_midreduce_nochase(self):
    Tensor.manual_seed(0)
@@ -1427,7 +1422,6 @@ class TestSchedule(unittest.TestCase):
    b = (a.sum(0)+a.max(0) + a.max(1)+a.sum(1)) + 2
    # schedule = check_schedule(b, 2)
    schedule = check_schedule(b, 4)
-    self.assertIs(store_val(schedule[0]).op, Ops.REDUCE_AXIS)
    run_schedule(schedule)
    np.testing.assert_allclose(b.numpy(), a.numpy().sum(0)+a.numpy().max(0) + a.numpy().max(1)+a.numpy().sum(1)+2, atol=1e-4, rtol=1e-4)

@@ -1619,7 +1613,7 @@ class TestSchedule(unittest.TestCase):
    Tensor.manual_seed(0)
    x = Tensor.randn(10, 20).realize()
    out = x.argmax(1)
-    run_schedule(check_schedule(out, 3)) # TODO: push a reduceop through a reshape
+    run_schedule(check_schedule(out, 2))

  def test_conv2d(self): _test_conv2d(7)
  def test_conv2d_fused(self): _test_conv2d(5, FUSE_CONV_BW=1)
@@ -1733,7 +1727,8 @@ class TestIndexing(unittest.TestCase):
      s = Tensor.schedule(*lst)
      lowered = [x[1] for x in lower_schedule(s.copy())]
      kernels = [ei for ei in list(lowered) if isinstance(ei.prg, CompiledRunner)]
-      if FUSE_ARANGE: self.assertEqual(len(kernels), cnt)
+      if FUSE_ARANGE and len(kernels) != cnt:
+        raise KernelCountException(f"{len(kernels)} != {cnt}")
      for ei in lowered: ei.run(do_update_stats=True)
    return s

@@ -1746,26 +1741,26 @@ class TestIndexing(unittest.TestCase):

  def test_simple_indexing_alt(self):
    X = Tensor.arange(16).reshape(4, 4)
-    xt = X[[1, 2], [1, 2]]
-    self.check_schedule(xt, 5)
-    np.testing.assert_equal(xt.numpy(), (np.arange(16).reshape(4, 4))[[1, 2], [1, 2]])
+    xt = X[[1, 2], [-1, 2]]
+    self.check_schedule(xt, 1)
+    np.testing.assert_equal(xt.numpy(), (np.arange(16).reshape(4, 4))[[1, 2], [-1, 2]])

  def test_advanced_indexing(self):
    X = Tensor.arange(10)+1
-    xt = X[[0]]
-    self.check_schedule(xt, 2)
-    np.testing.assert_equal(xt.numpy(), (np.arange(10)+1)[[0]])
+    xt = X[[0, -1]]
+    self.check_schedule(xt, 1)
+    np.testing.assert_equal(xt.numpy(), (np.arange(10)+1)[[0, -1]])

  def test_advanced_indexing_alt(self):
    X = Tensor.arange(6).reshape(3, 2)+1
    xt = X[[Tensor([2]), Tensor([1])]]
-    self.check_schedule(xt, 6)
+    self.check_schedule(xt, 3)
    np.testing.assert_equal(xt.numpy(), 6)

  def test_advanced_simple_indexing_combined(self):
    X = Tensor.arange(16).reshape(4, 4)
-    xt = X[1:2, [1, 2]]
-    self.check_schedule(xt, 4)
+    xt = X[1:2, [-1, 2]]
+    self.check_schedule(xt, 1)

  def test_push_through_reshape(self):
    Tensor.manual_seed(0)
@@ -1902,8 +1897,6 @@ class TestIndexing(unittest.TestCase):
    a = Tensor.arange(4).reshape(2, 2, 1).expand(2, 2, 2).contiguous().to("CPU")
    sched = self.check_schedule(a, 2) # NOTE: there is a contiguous between REDUCE_AXIS and COPY
    self.assertIs(sched[2].ast.op, Ops.COPY)
-    self.assertIs(store_val(sched[1]).op, Ops.LOAD)
-    self.assertIs(store_val(sched[0]).op, Ops.ADD)
    np.testing.assert_equal(a.numpy(), [[[0, 0], [1, 1]], [[2, 2], [3, 3]]])

  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@@ -1987,24 +1980,6 @@ class TestIndexing(unittest.TestCase):
    self.assertEqual(new_uop.st, ShapeTracker.from_shape((4,)).reshape((4, 1)))
    self.assertEqual(swizzle_cnt(new_uop), 0)

-  def test_no_rewrite_elementwise(self):
-    a = Tensor.empty(32, 32)
-    b = Tensor.empty(32, 32)
-    sink = (a+b).schedule()[0].ast
-    self.assertEqual(swizzle_cnt(sink), 0)
-
-  def test_simple_store_reshape(self):
-    a = Tensor.empty(32, 32).sum(axis=1)+Tensor.empty(1,32)
-    ast = a.schedule()[0].ast
-    self.assertEqual(ast.shape, (32, 1))
-    self.assertEqual(a.uop.shape, (1, 32))
-
-  def test_no_reshape_reduceop(self):
-    a = Tensor.empty(32, 32).sum(axis=(1,)).contiguous()
-    ast = a.schedule()[0].ast
-    self.assertEqual(ast.shape, (32, 1))
-    self.assertEqual(a.uop.shape, (32,))
-
 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, Ops.ASSIGN}])

@@ -2092,7 +2067,6 @@ class TestSwizzle(unittest.TestCase):
    np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)

  # kernels can only have 1 or n in each dim
-  @unittest.expectedFailure
  def test_dont_parallelize_different_n(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 2, 2).realize()
@@ -2108,7 +2082,6 @@ class TestSwizzle(unittest.TestCase):
    run_schedule(check_schedule(t, 3))
    np.testing.assert_equal(t.numpy(), [[0.5, 0.5], [0.5, 0.5], [0., 0.]])

-def store_val(si:ScheduleItem): return si.ast.src[0].src[1]
 zero_pm = UPat(Ops.CONST, arg=0)
 class TestView(unittest.TestCase):
  def test_all_masked_out(self):
@@ -2117,7 +2090,6 @@ class TestView(unittest.TestCase):
    # all masked out, degrades to const 0
    b = a.pad(((0, 10), None))[10:]
    sched = check_schedule(b.contiguous(), 1)
-    assert zero_pm.match(store_val(sched[-1]), {})
    run_schedule(sched)
    np.testing.assert_equal(b.numpy(), 0)

@@ -2128,7 +2100,6 @@ class TestView(unittest.TestCase):
    assert b.shape == (10, 10)
    sched = check_schedule(b.contiguous(), 1)
    self.assertEqual(sched[-1].ast.full_shape, (10, 10))
-    assert zero_pm.match(store_val(sched[-1]), {})
    run_schedule(sched)
    np.testing.assert_equal(b.numpy(), 0)

@@ -2143,8 +2114,6 @@ class TestView(unittest.TestCase):
    b = a.pad(((0, 5), None))[5:]
    assert b.shape == (10, 10)
    sched = check_schedule(b.contiguous(), 1)
-    self.assertEqual(store_val(sched[-1]).op, Ops.LOAD)
-    self.assertEqual(store_val(sched[-1]).st_arg, b.uop.st)
    run_schedule(sched)
    np.testing.assert_allclose(b.numpy(), np.pad(a.numpy(), ((0, 5), (0, 0)))[5:])

@@ -2260,24 +2229,6 @@ class TestConst(unittest.TestCase):
    sched = a.schedule()
    self.assertEqual(len(sched), 1)

-  def test_const_ast(self):
-    a = Tensor.ones((4,)).pad((1, 1)).contiguous()
-    sched = a.schedule()
-    print(sched[0].ast)
-    const_ast_pattern = UPat(Ops.SINK, src=(UPat.store(UPat(), UPat.where(UPat(Ops.VALID), UPat.cvar("x"), UPat(Ops.CONST, arg=0))),))
-    self.assertEqual(len(const_ast_pattern.match(sched[0].ast, {})), 1)
-    run_schedule(sched)
-    self.assertListEqual(a.tolist(), [0, 1, 1, 1, 1, 0])
-
-  def test_unmasked_const_ast(self):
-    a = Tensor.ones((4,)).contiguous()
-    sched = a.schedule()
-    print(sched[0].ast)
-    const_ast_pattern = UPat(Ops.SINK, src=(UPat.store(UPat(), UPat(Ops.CONST)),))
-    self.assertEqual(len(const_ast_pattern.match(sched[0].ast, {})), 1)
-    run_schedule(sched)
-    self.assertListEqual(a.tolist(), [1, 1, 1, 1])
-
  # ** part 2: scheduler behavior when const folding happens later

  def test_const_folding_no_realize(self):
@@ -1,8 +1,8 @@
 import unittest

-from tinygrad.opt.kernel import Opt, OptOps, Kernel
+from tinygrad.codegen.opt.kernel import Opt, OptOps, Kernel
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.opt.search import bufs_from_lin, actions, beam_search
+from tinygrad.codegen.opt.search import bufs_from_lin, actions, beam_search
 from tinygrad.device import Device
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
@@ -36,7 +36,7 @@ class TestBEAM(unittest.TestCase):

  def test_get_kernel_actions_dedup(self):
    from test.test_linearizer import helper_realized_ast
-    from tinygrad.opt.search import get_kernel_actions
+    from tinygrad.codegen.opt.search import get_kernel_actions
    a = Tensor.empty(4, 3)
    b = Tensor.empty(3)
    realized_ast, _ = helper_realized_ast(a @ b)
@@ -61,7 +61,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.opt.search import get_kernel_actions
+    from tinygrad.codegen.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]
@@ -78,7 +78,7 @@ class TestBEAM(unittest.TestCase):

  def test_get_kernel_actions_preserves_actions_state(self):
    from test.test_linearizer import helper_realized_ast
-    from tinygrad.opt.search import get_kernel_actions
+    from tinygrad.codegen.opt.search import get_kernel_actions
    a = Tensor.rand(16, 16)
    b = Tensor.rand(16, 16)
    realized_ast, _ = helper_realized_ast(a @ b)
@@ -131,10 +131,11 @@ class TestBEAM(unittest.TestCase):
    assert tm

  def test_beam_unnamed_kernels(self):
+    from test.test_linearizer import push_views
    a = Tensor.rand(100)
    b = Tensor.rand(100)
    si = (a+b).schedule()[-1]
-    lin = Kernel(si.ast)
+    lin = Kernel(push_views(si.ast))
    bufs = bufs_from_lin(lin)
    # TODO: beam should have better instrumentation so we don't have to check this indirect thing
    kcount = len(Kernel.kernel_cnt)
@@ -111,6 +111,12 @@ class TestFuse(unittest.TestCase):
    with Context(NOOPT=1):
      self._test_fuse(Tensor.scaled_dot_product_attention, q, k, v, atol=1e-5)

+  def test_mismatch_reduce(self):
+    a = Tensor.ones(16, 10).contiguous().realize()
+    b = Tensor.ones(16, 20).contiguous().realize()
+    c = (a.sum(axis=1) + b.sum(axis=1)).fuse()
+    self.assertListEqual(c.tolist(), [30]*16)
+
 class TestSoftmaxFusion(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
@@ -25,7 +25,7 @@ class TestStunning(unittest.TestCase):
    nv = a[12].cat(a[76]).tolist()

    vi = Variable('i', 0, a.shape[0]-1)
-    with self.assertRaisesRegex(AssertionError, "different values for the same key"):
+    with self.assertRaisesRegex(AssertionError, "bind mismatch on"):
      wv = a[vi.bind(12)].cat(a[vi.bind(76)]).tolist()
      self.assertListEqual(nv, wv)

@@ -14,7 +14,7 @@ from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps

 def to_uops_list(u:list[UOp], opts=None, skip_check=False) -> list[UOp]: return full_rewrite(UOp.sink(*u), opts)

@@ -415,6 +415,17 @@ class TestAssembly(unittest.TestCase):
    uops = program.uops
    self.assertEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)

+  def test_use_cmpeq(self):
+    g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
+    c = UOp(Ops.CONST, dtypes.uint, (), 7)
+    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
+    comp = l.ne(c).ne(True)
+    uops = to_uops_list([comp], opts=Device[Device.DEFAULT].renderer)
+    Device[Device.DEFAULT].renderer.render(uops)
+    ops = [x.op for x in uops]
+    self.assertIn(Ops.CMPEQ, ops)
+    self.assertNotIn(Ops.CMPNE, ops)
+
 class TestUOpMethod(unittest.TestCase):
  @unittest.skip("uops lt no longer ordered")
  def test_compare_alu_same_src_different_arg(self):
@@ -508,19 +519,6 @@ class TestShapeSpec(unittest.TestCase):
    a = Tensor.ones((4, 4)).uop
    self.assertEqual(a.st, ShapeTracker.from_shape(()).reshape((1,1)).expand((4,4)))

-  def test_padded_const(self):
-    a = Tensor.ones((1, 1)).pad(((1, 1), (1, 1)))
-    ast = a.contiguous().schedule()[0].ast
-    valid_pattern = UPat(Ops.WHERE, src=(UPat(Ops.VALID), UPat.cvar(), UPat.cvar()))
-    valid_ternary = [x for x in ast.toposort() if valid_pattern.match(x, {})][0]
-    # the WHERE outputs a contiguous (3, 3)
-    self.assertEqual(valid_ternary.st, ShapeTracker.from_shape((3, 3)))
-    valid, x, y = valid_ternary.src
-    # very notably, only the first source is padded
-    self.assertIsNotNone(valid.st.views[-1].mask)
-    assert x.st.views[-1].mask is y.st.views[-1].mask is None
-    assert all(s.shape == (3, 3) for s in valid_ternary.src)
-
  # NOTE: CONST ShapeTracker comes from its source
  def test_scalar_const(self):
    a = Tensor(0).uop
@@ -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.opt.kernel import Kernel, Opt, OptOps, KernelOptError
+from tinygrad.codegen.opt.kernel import Opt, OptOps, KernelOptError
 from tinygrad.device import Device

 def flops_mem(uops, ignore_indexing=False):
@@ -173,76 +173,60 @@ class TestStatsOptimized(unittest.TestCase):
    self.assertEqual(p.estimates.mem, 3*N*N*4) # 3 NxN mats with floats

  def test_gemm(self):
-    k = Kernel(self.ast_gemm)
-    p = get_program(k.get_optimized_ast(), k.opts)
+    p = get_program(self.ast_gemm, opts=[])
    self.check_gemm(p)
    self.assertEqual(p.estimates.lds, 2*N*N*N*4 + 4*N*N)

  def test_gemm_tc_unroll(self):
-    k = Kernel(self.ast_gemm)
-    if not k.apply_tensor_cores(): self.skipTest("no tensor cores")
-    k.apply_opt(Opt(OptOps.UNROLL, 0, 2))
-    p = get_program(k.get_optimized_ast(), k.opts)
+    try:
+      p = get_program(self.ast_gemm, opts=[Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)])
+    except KernelOptError:
+      raise unittest.SkipTest("no tensor cores")
    print(p.src)
    self.check_gemm(p)

  # this is a good lesson about why UPCASTing is a good idea

  def test_gemm_one_upcasted(self):
-    k = Kernel(self.ast_gemm)
-    k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
-    p = get_program(k.get_optimized_ast(), k.opts)
+    p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4)])
    self.check_gemm(p)
    self.assertEqual(p.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)

  def test_gemm_upcasted(self):
-    k = Kernel(self.ast_gemm)
-    k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
-    k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
-    k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
-    p = get_program(k.get_optimized_ast(), k.opts)
+    p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)])
    self.check_gemm(p)
    self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)

  def test_gemm_upcasted_locals(self):
-    k = Kernel(self.ast_gemm)
-    k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
-    k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
    try:
-      k.apply_opt(Opt(OptOps.LOCAL, 0, 5))
-      k.apply_opt(Opt(OptOps.LOCAL, 1, 5))
+      p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4),
+                                           Opt(OptOps.LOCAL, 0, 4),  Opt(OptOps.LOCAL, 1, 4)])
    except KernelOptError:
      raise unittest.SkipTest("no locals")
-    p = get_program(k.get_optimized_ast(), k.opts)
    self.check_gemm(p)
    self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)

  def test_gemm_group(self):
-    k = Kernel(self.ast_gemm)
    try:
-      k.apply_opt(Opt(OptOps.GROUP, 0, 4))
+      p = get_program(self.ast_gemm, opts=[Opt(OptOps.GROUP, 0, 4)])
    except KernelOptError:
      raise unittest.SkipTest("no locals")
    SZ = N*N*4
-    p = get_program(k.get_optimized_ast(), k.opts)
    # NOTE: these are sort of wrong. they aren't honoring the IF statement
    self.check_gemm(p, extra_flops=SZ*4)
    self.assertEqual(p.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)

  def test_reduce(self):
-    k = Kernel(self.ast_reduce)
-    p = get_program(k.get_optimized_ast(), k.opts)
+    p = get_program(self.ast_reduce, opts=[])
    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
    self.assertEqual(p.estimates.ops, N*N)
    self.assertEqual(p.estimates.mem, N*N*4 + 4)

  def test_reduce_group(self):
-    k = Kernel(self.ast_reduce)
    try:
-      k.apply_opt(Opt(OptOps.GROUP, 0, 50))
+      p = get_program(self.ast_reduce, opts=[Opt(OptOps.GROUP, 0, 50)])
    except KernelOptError:
      raise unittest.SkipTest("no locals")
-    p = get_program(k.get_optimized_ast(), k.opts)
    # NOTE: these are wrong, they don't respect the if statement
    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)

@@ -1,10 +1,7 @@
 import unittest
 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.opt.kernel import Kernel
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.helpers import CI, Profiling, WINO, getenv

 class TestWinogradClose(unittest.TestCase):
  def test_close(self):
@@ -28,30 +25,6 @@ class TestWinograd(unittest.TestCase):
  def tearDown(self):
    WINO.value = self.old

-  def test_speed(self):
-    x = Tensor.empty(1,4,9,9)
-    w = Tensor.empty(4,4,3,3)
-
-    with Timing("running conv: "):
-      out = Tensor.conv2d(x, w)
-
-    with Timing("scheduling: "):
-      sched = out.schedule()
-
-    for i,s in enumerate(sched):
-      if s.ast.op is not Ops.SINK: continue
-      ops = s.ast.toposort()
-      with Timing(f"linearize {i} with {len(ops):4d} ops: "):
-        l = Kernel(s.ast)
-        l.apply_opts(hand_coded_optimizations(l))
-      assert len(l.sts) <= 256  # just the current value to prevent regression
-      if DEBUG >= 2: print(f"{len(l.sts):4d} shapetrackers with max {max(len(x.views) for x in l.sts)} views")
-      for st in l.sts:
-        assert len(st.views) <= 2, "too many views in winograd"
-        if DEBUG >= 3:
-          print(f"{len(st.views):3d} views")
-          for v in st.views: print(v)
-
  def test_profile(self):
    x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
    with Profiling(enabled=not CI, sort='time'):
@@ -1,55 +0,0 @@
-#!/usr/bin/env python
-import unittest
-from tinygrad.uop.ops import Ops
-from tinygrad.tensor import Tensor
-from tinygrad.nn import Conv2d
-from tinygrad.shape.shapetracker import ShapeTracker, View
-from tinygrad.helpers import prod
-from test.unit.test_shapetracker import shapetracker_getitem
-
-class TestConvShapetracker(unittest.TestCase):
-  def test_conv_3x3_one_view(self):
-    conv = Conv2d(16, 32, (3, 3))
-    # first run to init the weights, they are scheduled.
-    conv(Tensor.empty(1, 16, 10, 10)).schedule()
-    # run it again to get the kernels
-    sched = [si for si in conv(Tensor.empty(1, 16, 10, 10)).schedule() if si.ast.op is Ops.SINK]
-    assert len(sched) == 1, f"conv should only have one kernel, getting {len(sched)}"
-    for st in [x.st_arg for x in sched[0].ast.toposort() if x.op is Ops.LOAD]:
-      assert len(st.views) == 1
-
-  def test_conv_2x2_backward_one_view(self):
-    X = Tensor.rand(1, 1, 3, 3, requires_grad=True)
-    conv = Conv2d(1, 1, (2, 2), bias=False)
-    conv(X).mean().backward()
-    si = X.grad.schedule()[-1]
-    print(si)
-    ldb = [x for x in si.ast.toposort() if x.op is Ops.LOAD][0]
-    st: ShapeTracker = ldb.st_arg.simplify()
-    print(si.bufs[1].size)
-    self.assertEqual(si.bufs[1].size, st.real_size())
-    for v in st.views: print(v)
-
-    # same st
-    test_st = ShapeTracker((
-      View(shape=(1, 1, 2, 4, 2, 4), strides=(0, 0, 2, 8, 1, 4), offset=0, mask=((0, 1), (0, 1), (0, 2), (0, 2), (0, 2), (0, 2)), contiguous=False),
-      View(shape=(1, 1, 1, 1, 3, 3, 3, 3), strides=(0, 0, 0, 0, 24, 8, 3, 1), offset=0,
-           mask=((0, 1), (0, 1), (0, 1), (0, 1), (0, 2), (0, 3), (0, 2), (0, 3)), contiguous=False)))
-    #test_st = ShapeTracker((
-    #  View(shape=(2,4), strides=(1,4), offset=0, mask=None, contiguous=False),
-    #)).simplify()
-      #View(shape=(1, 1, 2, 4, 2, 4), strides=(0, 0, 2, 8, 1, 4), offset=0, mask=((0, 1), (0, 1), (0, 2), (0, 2), (0, 2), (0, 2)), contiguous=False),
-      #View(shape=(1, 1, 1, 1, 3, 3, 3, 3), strides=(0, 0, 0, 0, 24, 8, 3, 1), offset=0,
-      #     mask=((0, 1), (0, 1), (0, 1), (0, 1), (0, 2), (0, 3), (0, 2), (0, 3)), contiguous=False))).simplify()
-    print("*** new ***")
-    for v in test_st.views: print(v)
-    for i in range(prod(st.shape)):
-      i1, i2 = shapetracker_getitem(st, i), shapetracker_getitem(test_st, i)
-      print(i, i1, i2, si.bufs[1].size, i1==i2)
-      #self.assertEqual(i1, i2)
-
-    with self.assertRaises(AssertionError):
-      assert len(st.views) <= 2
-
-if __name__ == '__main__':
-  unittest.main()
@@ -303,8 +303,8 @@ class TestRecurse(unittest.TestCase):
  def test_inf_loop(self):
    a = UOp.variable('a', 0, 10)
    pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.DEFINE_REG)),
-      (UPat(Ops.DEFINE_REG, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
+      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
    ])
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm)
@@ -312,8 +312,8 @@ class TestRecurse(unittest.TestCase):
  def test_inf_loop_bottom_up(self):
    a = UOp.variable('a', 0, 10)
    pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.DEFINE_REG)),
-      (UPat(Ops.DEFINE_REG, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
+      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
    ])
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm, bottom_up=True)
@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Context, Device
 from tinygrad.engine.realize import get_program
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps
 from tinygrad.uop.ops import KernelInfo

 class TestLinearizerRewrite(unittest.TestCase):
@@ -1,7 +1,8 @@
-import unittest, base64, functools
+import unittest, base64, functools, sys
 from tinygrad.apps.llm import SimpleTokenizer, get_llama_re
 from tinygrad.helpers import fetch

+@unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
 class TestLLMTokenizer(unittest.TestCase):
  @functools.cached_property
  def basic_tok(self): return SimpleTokenizer(".*", { b"a": 0, b"b": 1, b"c": 2, b"ab": 3, b"bc": 4 }, { "<x>": 5, "<y>": 6, "<z>": 7 })
@@ -0,0 +1,110 @@
+import unittest
+from dataclasses import dataclass, field
+from tinygrad.uop.ops import PatternMatcher, UOp, graph_rewrite, Ops, UPat, GroupOp, RewriteNotReady
+
+# we could insert CHILDREN node
+
+@dataclass
+class ChildrenContext:
+  children: dict[UOp, list[UOp]]|None = None
+
+# this is a generic child labeller
+def extract_children(ctx:ChildrenContext, x:UOp):
+  if ctx.children is not None: return
+  ctx.children = {k:list(v.keys()) for k,v in x.get_children_map().items() if len(v) > 1}
+
+def mark_children(ctx:ChildrenContext, x:UOp):
+  new_srcs = [(UOp(Ops.CHILD, s.dtype, src=(s,), arg=(ctx.children[s].index(x), len(ctx.children[s]))) if s in ctx.children else s) for s in x.src]
+  return x.replace(src=tuple(new_srcs))
+
+pm_children = PatternMatcher([
+  (UPat(Ops.SINK, name="x"), extract_children),
+  (UPat(GroupOp.All-{Ops.CHILD}, name="x"), mark_children),
+])
+
+@dataclass
+class TestContext:
+  seen_children: dict[UOp, set[int]] = field(default_factory=dict)
+  ready_children: dict[UOp, set[int]] = field(default_factory=dict)
+  seen_consts:int = 0
+  saved_seen_consts:int = 0
+  exp2_visit_count:int = 0
+
+# this is a generic pattern
+def visit_child(ctx:ChildrenContext, x:UOp):
+  if x.src[0] not in ctx.seen_children:
+    ctx.seen_children[x.src[0]] = set()
+    ctx.ready_children[x.src[0]] = set()
+  ctx.seen_children[x.src[0]].add(x.arg[0])
+  if len(ctx.seen_children[x.src[0]]) != x.arg[1]:
+    print(f"visit CHILD {x.arg} bottom up -- not ready {ctx.seen_children[x.src[0]]}")
+    raise RewriteNotReady
+  print(f"visit CHILD {x.arg} bottom up -- READY {ctx.seen_children[x.src[0]]}")
+  ctx.ready_children[x.src[0]].add(x.arg[0])
+
+pm_child_visitor = PatternMatcher([
+  (UPat(Ops.CHILD, name="x"), visit_child),
+])
+
+# this is for the test
+def see_const(ctx:ChildrenContext, c:UOp): ctx.seen_consts += c.arg
+def see_exp2(ctx:ChildrenContext): ctx.exp2_visit_count += 1
+def save_seen_consts(ctx:ChildrenContext, x:UOp): ctx.saved_seen_consts = ctx.seen_consts
+pm_consts = PatternMatcher([
+  (UPat(Ops.DEFINE_VAR, name="x"), save_seen_consts),
+  (UPat()+UPat.cvar("c"), see_const),
+  (UPat(Ops.EXP2), see_exp2),
+])
+
+class TestChildrenRewrite(unittest.TestCase):
+  def test_not_ready_double_simple(self):
+    global_a = UOp.variable("a", 0, 10).exp2()
+    inter = (global_a+global_a).exp2()
+    global_sink = (inter+inter).sink()
+
+    sink = graph_rewrite(global_sink, pm_children, ctx=ChildrenContext(), bottom_up=True)
+    ctx = TestContext()
+    graph_rewrite(sink, pm_consts, ctx=ctx, bottom_up=True)
+    self.assertEqual(ctx.exp2_visit_count, 2)
+
+  def test_not_ready_double(self):
+    global_a = UOp.variable("a", 0, 10).exp2()
+    inter = ((global_a+1000)+(global_a+100)).exp2()
+    global_sink = ((inter+10)+(inter+1)).sink()
+
+    sink = graph_rewrite(global_sink, pm_children, ctx=ChildrenContext(), bottom_up=True)
+    print("test_not_ready_double")
+    ctx = TestContext()
+    graph_rewrite(sink, pm_child_visitor+pm_consts, ctx=ctx, bottom_up=True)
+    self.assertEqual(ctx.exp2_visit_count, 2)
+    self.assertEqual(ctx.seen_consts, ctx.saved_seen_consts)
+    self.assertEqual(ctx.seen_consts, 1111)
+
+  def test_in_srcs_twice(self):
+    global_a = UOp.variable("a", 0, 10).exp2()
+    global_sink = (global_a+global_a).sink()
+
+    ctx = TestContext()
+    graph_rewrite(global_sink, pm_consts, ctx=ctx, bottom_up=True)
+    self.assertEqual(ctx.exp2_visit_count, 1)
+
+  def test_not_ready(self):
+    global_a = UOp.variable("a", 0, 10).exp2()
+    global_sink = ((global_a+2)+(global_a+3)).sink()
+
+    # without children and not ready, we don't see both adds before the DEFINE_VAR
+    ctx = TestContext()
+    graph_rewrite(global_sink, pm_consts, ctx=ctx, bottom_up=True)
+    self.assertNotEqual(ctx.seen_consts, ctx.saved_seen_consts)
+    self.assertEqual(ctx.exp2_visit_count, 1)
+
+    # with children and not ready we do
+    sink = graph_rewrite(global_sink, pm_children, ctx=ChildrenContext(), bottom_up=True)
+    ctx = TestContext()
+    graph_rewrite(sink, pm_child_visitor+pm_consts, ctx=ctx, bottom_up=True)
+    self.assertEqual(ctx.seen_consts, ctx.saved_seen_consts)
+    self.assertEqual(ctx.exp2_visit_count, 1)
+    self.assertSetEqual(list(ctx.ready_children.values())[0], {0,1})
+
+if __name__ == '__main__':
+  unittest.main()
@@ -1,10 +1,11 @@
 import unittest
 from tinygrad import Tensor, Device
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.device import Buffer
-from tinygrad.opt.search import get_test_global_size, bufs_from_lin
+from tinygrad.codegen.opt.search import get_test_global_size, bufs_from_lin
 from tinygrad.helpers import GlobalCounters
 from extra.optimization.helpers import time_linearizer
+from test.test_linearizer import push_views

 class TestSearchUtil(unittest.TestCase):
  def test_get_test_global_size(self):
@@ -25,7 +26,7 @@ class TestSearchUtil(unittest.TestCase):
    a = Tensor.randn(4, 4).realize()
    b = a+a[0]
    si = b.schedule()[0]
-    rawbufs = bufs_from_lin(Kernel(si.ast))
+    rawbufs = bufs_from_lin(Kernel(push_views(si.ast)))
    assert len(rawbufs) == 2
    assert all(r is not None for r in rawbufs)
    assert all(isinstance(r, Buffer) for r in rawbufs)
@@ -38,13 +39,13 @@ class TestTimeLinearizer(unittest.TestCase):
    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)
+    tm = time_linearizer(Kernel(push_views(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)
+    lin = Kernel(push_views(ast))
    bufs = bufs_from_lin(lin)

    kernel_count = GlobalCounters.kernel_count
@@ -2,8 +2,8 @@ import unittest, math
 import numpy as np
 from tinygrad import dtypes
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.uop.transcendental import TRANSCENDENTAL_SUPPORTED_DTYPES, payne_hanek_reduction, cody_waite_reduction
-from tinygrad.uop.transcendental import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if
+from tinygrad.uop.decompositions import TRANSCENDENTAL_SUPPORTED_DTYPES, payne_hanek_reduction, cody_waite_reduction
+from tinygrad.uop.decompositions import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if
 from test.helpers import eval_uop

 class TestTranscendentalFunctions(unittest.TestCase):
@@ -34,17 +34,6 @@ class TestUOpSpec(unittest.TestCase):
    store = UOp(Ops.STORE, dtypes.void, (buf_0.view(ShapeTracker.from_shape((32, 1))), a+b))
    helper_test_verify_ast(store)

-  def test_exactly_one_full_shape(self):
-    dtype = dtypes.int
-    bufs = [UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), i) for i in range(6)]
-    a = UOp(Ops.LOAD, dtype, (bufs[2].view(ShapeTracker.from_shape((32, 1))),))
-    b = UOp(Ops.LOAD, dtype, (bufs[3].view(ShapeTracker.from_shape((32, 1))),))
-    st0 = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 1))), a+b)
-    a = UOp(Ops.LOAD, dtype, (bufs[4].view(ShapeTracker.from_shape((32, 32))),))
-    b = UOp(Ops.LOAD, dtype, (bufs[5].view(ShapeTracker.from_shape((32, 32))),))
-    st1 = UOp.store(bufs[1].view(ShapeTracker.from_shape((32, 32))), a+b)
-    with self.assertRaises(InvalidASTException): helper_test_verify_ast(st0, st1)
-
  def test_no_implicit_broadcasting(self):
    bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)]
    a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker.from_shape((4, 32))),))
@@ -73,14 +62,6 @@ class TestUOpSpec(unittest.TestCase):
    st = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 1))), r+a)
    with self.assertRaises(InvalidASTException): helper_test_verify_ast(st)

-  def test_buffer_uops_st(self):
-    a = Tensor.randn(4, 4)+2
-    helper_test_verify_ast(ast:=a.schedule()[-1].ast)
-    store_st = [u.st for u in ast.toposort() if u.op is Ops.STORE][0]
-    self.assertEqual(store_st, ShapeTracker.from_shape((4, 4)))
-    const_st = [u.st for u in ast.toposort() if u.op is Ops.CONST][0]
-    self.assertEqual(const_st, ShapeTracker.from_shape((1, 1)).expand((4, 4)))
-
  def test_assert_swizzle(self):
    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
    a = UOp(Ops.LOAD, dtypes.float, (buf.view(ShapeTracker.from_shape((32, 1))),))
@@ -203,9 +203,9 @@ class TestSymbolic(unittest.TestCase):
  def test_mod_min_max(self):
    self.helper_test_variable(Variable("x", 0, 10)%Variable("y", 1, 10), 0, 9, "(x%y)")
    self.helper_test_variable(Variable("x", -10, 0)%Variable("y", 1, 10), -9, 0, "(((x*-1)%y)*-1)")
-    self.helper_test_variable(Variable("x", 0, 10)%Variable("y", -10, -1), 0, 9, "(x%y)")
-    self.helper_test_variable(Variable("x", -10, 0)%Variable("y", -10, -1), -9, 0, "(((x*-1)%y)*-1)")
-    self.helper_test_variable(Variable("x", -10, 10)%Variable("y", -10, -1), -9, 9, "(x%y)")
+    self.helper_test_variable(Variable("x", 0, 10)%Variable("y", -10, -1), 0, 9, "(x%(y*-1))")
+    self.helper_test_variable(Variable("x", -10, 0)%Variable("y", -10, -1), -9, 0, "(((x*-1)%(y*-1))*-1)")
+    self.helper_test_variable(Variable("x", -10, 10)%Variable("y", -10, -1), -9, 9, "(x%(y*-1))")

    # test _min_max directly without the rewrite taking out the sign
    self.assertEqual((Variable("x", -10, 0)%Variable("y", -10, -1))._min_max, (-9, 0))
@@ -296,7 +296,7 @@ class TestSymbolic(unittest.TestCase):
  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)")
+    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))")
@@ -97,9 +97,9 @@ class TestViz(BaseTestViz):

  # name can also come from a function that returns a string
  def test_dyn_name_fxn(self):
-    @track_rewrites(name=lambda a,ret: a.render())
-    def name_from_fxn(s:UOp): return graph_rewrite(s, PatternMatcher([]))
-    name_from_fxn(UOp.variable("a", 1, 10)+1)
+    @track_rewrites(name=lambda *args,ret,**kwargs: ret.render())
+    def name_from_fxn(s:UOp, arg:list|None=None): return graph_rewrite(s, PatternMatcher([]))
+    name_from_fxn(UOp.variable("a", 1, 10)+1, arg=["test"])
    lst = get_viz_list()
    # name gets deduped by the function call counter
    self.assertEqual(lst[0]["name"], "(a+1) n1")
@@ -124,10 +124,10 @@ class TestViz(BaseTestViz):

  def test_inf_loop(self):
    a = UOp.variable('a', 0, 10)
-    b = a.replace(op=Ops.DEFINE_REG)
+    b = a.replace(op=Ops.CONST)
    pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.DEFINE_REG)),
-      (UPat(Ops.DEFINE_REG, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
+      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
    ])
    with self.assertRaises(RuntimeError): exec_rewrite(a, [pm])
    graphs = flatten(x["graph"].values() for x in get_details(tracked_ctxs[0][0]))
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	260da2017c	fix unit test dtypes	2025-08-13 12:15:12 -07:00
geohot	65dcd6dd45	render ranges in viz, name gbufs with sizes. changes from rangeify	2025-08-13 12:06:34 -07:00
kevvz GitHub George Hotz	e2873a3a41	[bounty] Muon optim (#11414 ) * newton schulz * add muon + move newton schulz to tensor * compact newton schulz * better tests * cleanup * add comments for muon * cleanup * add export with tests * match muon optim with test optim * cleanup * unsed import * correct comment * whitespace * move export * muon test fix * match reference impl + tests * remove export by moving muon device * add credit * cleanup * remove print * spacing * spacing * comma * cleanup * removal * fix tests + optim momentum * consistent is not/ not * more consistency * fix test * cleanup * fix the nones * remove comment * cast * comment * comment * muon teeny test * muon flag beautiful mnist * set steps * steps as hyperparam * match default test steps * name * large cleanup * dont care about steps * nesterov false default * match each other impl * steps * switch nest * swap defaults * update docstring * add no nesterov test * ban fuse_optim * prints * classical momentum * alternative condition * recon * pre + post wd * false default * detach * signature changes * context * swap order * big cleanup * 0 step instead * parity * remove fuse * remove fused * better paper * assert message * correct shape check + eps * multidim * add eps * cleanup * correct assert message * lint * better tests * naming * ns_steps,ns_params * update docstring * docstring * match sgd and muon together * sandwich * add back fused * parity --------- Co-authored-by: George Hotz <[email protected]>	2025-08-13 14:27:55 -04:00
chenyuandGitHub	94e6d84e32	rewrite Tensor.round to not use cast int (#11654 )	2025-08-13 13:51:08 -04:00
George HotzandGitHub	d2521d828a	transcendental+idiv+threefry are uop decompositions (#11636 ) * transcendental+idiv+threefry are uop decompositions [pr] * threefry decomp * fix randomness tests * fix webgpu * unneeded now * fix * move prematcher * all cast should probably be cast_vec	2025-08-13 09:37:12 -07:00
geohotstanandGitHub	cf7224ce3e	fully lint onnx.py (#11647 ) * mypy * ruff ruff ruff	2025-08-13 08:22:06 -07:00
geohotstanandGitHub	925555b62a	Fix onnx Domain bug (#11650 )	2025-08-13 08:20:50 -07:00
Sieds LyklesandGitHub	67df617fe1	add launch bounds to ptx (#11646 )	2025-08-13 13:05:39 +02:00
qazalandGitHub	88f95e9f59	viz: minor fixups for firefox (#11645 ) * fix circle attr * set fill color	2025-08-13 12:59:28 +03:00
qazalandGitHub	6f88eac0fc	viz: refactor node and edge tagging (#11644 )	2025-08-13 12:41:01 +03:00
qazalandGitHub	8140bf9778	viz: create layout once (#11643 ) * start * work * works * diff cleanup	2025-08-13 09:24:58 +03:00
chenyuandGitHub	3fb79bb43a	minor onnx cleanups (#11642 )	2025-08-13 01:05:19 -04:00
chenyuandGitHub	e9e5a08a04	simplify onnx cubic (#11641 ) we can drop the double where and abs since we know which ranges the inputs map into	2025-08-12 19:57:31 -04:00
George HotzandGitHub	18cdbec447	split decompositions pass (#11638 ) * split decompositions pass * fix ptx * pack load store early * restore that	2025-08-12 12:56:05 -07:00
chenyuandGitHub	0d8a0d7a96	update test_multi_const_folding_tensor to include pow (#11635 ) pow folds now	2025-08-12 13:35:37 -04:00
Sieds LyklesandGitHub	4d6e407eb0	Extend fast_idiv to negative ints (#11632 ) * fast idiv for signed ints * Add rule and test * fix tests * redo fuzz_fast_idiv to do negative ints as well * adjust comments * remove unused imports	2025-08-12 19:34:49 +02:00
qazalandGitHub	17adbe86d8	hotfix: do not default to capturing args in track_rewrites (#11634 )	2025-08-12 20:01:24 +03:00
geohotstan GitHub chenyu	ad9dec25b3	combine onnx parser and onnx (#11485 ) * start * more * fix onnx_runner test * pass * patch for disk and add domains from huggingface * simpler docs * revert domain changes * rerun ci * revert onnx ops test change * add fix from strenum stuff * correct way * revert correct way to leave the fix for another PR * test segfault * Revert "test segfault" This reverts commit `4e1aaf41e7`. * remove some unnecessary documentation * test segfault again * Revert "test segfault again" This reverts commit `56fc5f03e7`. * try gemini suggested patch for sys._getframe * keep trying with gemini * revert not working gemini suggestions and try faulthandler * remove pythonfaulthandler * trigger CI a few times * minimize diff --------- Co-authored-by: chenyu <[email protected]>	2025-08-12 12:56:39 -04:00
Sieds LyklesandGitHub	4c3982c44e	Take sign out of mod (#11631 ) * Add rule and test * fix tests	2025-08-12 18:44:36 +02:00
qazalandGitHub	e28605e324	rename profile point event fields [pr] (#11633 )	2025-08-12 19:11:21 +03:00
nimlgenandGitHub	8a7be0a747	metal: workaround for transfers sync issue (#11622 ) * metal: workaround for transfers sync issue * metal tracsfer sync is broken * hm * rm it? * keep it	2025-08-12 16:16:34 +03:00
qazalandGitHub	efe8b5611d	move ProfilePointEvent out of device.py [pr] (#11630 ) Generic profiling events exist in helpers so they can be imported from everywhere in tinygrad.	2025-08-12 09:58:32 +03:00
chenyuandGitHub	0d7075f2de	assign should broadcast input tensor (#11629 ) fixed test_assign_broadcast	2025-08-11 23:36:35 -04:00
Joshua KissoonandGitHub	c44760c89d	torch backend: fix arange, add linalg.cross, add tests (#11628 )	2025-08-11 23:34:41 -04:00
George HotzandGitHub	ca41b5e38b	skip_0 in graph rewrite [pr] (#11627 ) * skip_0 in graph rewrite [pr] * no track_rewrites on test * use dict instead of set	2025-08-11 18:29:04 -07:00
SardorandGitHub	ca7a641442	fix bugs at examples/yolov3.py (#11614 ) * Update load_weight. Give valid model url * Fix bug in iou function	2025-08-11 21:14:47 -04:00
chenyuandGitHub	0c97d6de1b	don't round pow output for int pow int (#11625 ) also added atol=0 and big pows for the tests	2025-08-11 20:57:47 -04:00
chenyuandGitHub	d623f6d850	support int Tensor pow to const non-negative int (#11624 ) matches torch	2025-08-11 19:50:19 -04:00
chenyuandGitHub	857a830dcc	fix test_arange_float_step (#11623 )	2025-08-11 16:58:42 -04:00
chenyuandGitHub	0806677b51	rewrite sort idx (#11613 )	2025-08-11 16:20:56 -04:00
George HotzandGitHub	700c11597b	switch contextvars.ContextVar to _ContextVar (#11621 )	2025-08-11 12:20:09 -07:00
ttomsa GitHub wozeparrot	ae0c3cfff6	change clang -march flag to -mcpu on arm (#10970 ) Co-authored-by: wozeparrot <[email protected]>	2025-08-11 13:38:48 -04:00
geohotstanandGitHub	27bcb9fd1c	Support cubic mode for ONNX Resize OP (#11612 ) * start * add reference * this is so much slower * this makes sense but differs from official impl, but results are still correct..? * add a comment * Just keep it simple for now since I don't fully get it yet * address comments * correct * teeny clean up * another small comment improvement lol	2025-08-11 11:49:30 -04:00
nimlgenandGitHub	d2bb1bcb97	cloud: a bit better err handling (#11616 ) * cloud: err propagation to client * fix * print exc * linter * excs * fix * hm * flaky	2025-08-11 15:51:22 +03:00
qazalandGitHub	6a232ccdac	viz: add tiny range drawing helper (#11620 ) * viz: add tiny range drawing helper * less	2025-08-11 15:15:43 +03:00
qazalandGitHub	e768773e13	viz: use colors helper (#11618 )	2025-08-11 13:10:15 +03:00
qazalandGitHub	7d6c0a8cc7	viz: refactor progress msg (#11617 )	2025-08-11 13:01:36 +03:00
chenyuandGitHub	630edcffd8	remove .float calls in olmoe (#11610 ) still matches torch	2025-08-10 20:33:22 -04:00
chenyuandGitHub	a67e0917c3	list indexing can normalize in python (#11609 ) * list indexing can normalize in python list index does not need to be normalized in tensor * update those	2025-08-10 20:02:38 -04:00
chenyuandGitHub	1181ec0cd2	few more tensor indexing test cases (#11608 )	2025-08-10 18:56:42 -04:00
George HotzandGitHub	996c907c0b	rewrite not ready + children machinery (#11607 ) * rewrite not ready + children machinery * it doesn't like track rewrites	2025-08-10 15:28:30 -07:00
Sieds LyklesandGitHub	1875bc69f9	Late rewrite rules for CMPLT (#11591 ) * add rules * more rules * fix comment spelling * remove two rules	2025-08-10 22:18:13 +02:00
nimlgenandGitHub	5403a4aeaf	null dev: support offset on buffers (#11606 ) * null dev: support offset on buffers * nolimit	2025-08-10 21:58:37 +03:00
geohotstanandGitHub	b0dab6a4cd	onnx Resize OP clean up (#11603 ) * start * slight clean up	2025-08-10 14:10:39 -04:00
Sieds LyklesandGitHub	10540414cd	Add Ops.CMPEQ (#10431 ) * Add op * add to Groupop.ALU * fix spec * fix ptx * temporary pickle by name to see process replay * add Ops.EQ to binary ops * Actuall rename properly * add test to assert CMPEQ is being used * Ops.CMPEQ is automatic cast to bool * add Ops.CMPEQ to llvm * add Ops.CMPEQ to llvm	2025-08-10 13:13:16 +02:00
chenyuandGitHub	f7aa1b85fe	minor sort cleanups (#11602 )	2025-08-10 01:51:23 -04:00
chenyuandGitHub	dfb702ef33	fix sort for small dim (#11601 ) * fix sort for small dim * fixed test_sort_empty	2025-08-10 01:17:41 -04:00
chenyuandGitHub	ef17af85c6	remove .float call in llama logit (#11598 ) * remove .float call in llama logit * bfloat item	2025-08-10 00:02:18 -04:00
chenyuandGitHub	dd3d2eb36c	add training llama3 test in ci (#11599 )	2025-08-09 22:35:39 -04:00
chenyuandGitHub	3e64467322	remove freqs_cis contiguous in llama (#11597 )	2025-08-09 21:11:12 -04:00
chenyuandGitHub	7338ffead0	small beautiful_mnist update (#11596 ) gather is fast now. there's a conv/bw kernel that only gets fast with BEAM, but whole thing runs < 5 seconds now regardless	2025-08-09 19:51:14 -04:00
chenyuandGitHub	45baec1aab	model parallel llama (#11588 ) MP=8 GRADIENT_ACC_STEPS=3 BS=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=70B SEQLEN=512 PYTHONPATH=. MODEL=llama3 python3 examples/mlperf/model_train.py	2025-08-09 16:54:27 -04:00
nimlgenandGitHub	09bc377da3	search: print runtime failures on debug (#11593 )	2025-08-09 23:01:19 +03:00
nimlgenandGitHub	14f99ff1a1	amd: doorbell_cpu_addr is not used (#11592 ) * amd: doorbell_cpu_addr is not used * hm	2025-08-09 20:03:21 +03:00
Sieds LyklesandGitHub	01c770c77b	Fix z3 float cast in indexing (#11590 ) * adjust dtype of z3_renderer and add rule for cast * dtypes.bool is also cast noop * add regression test * make embedding smaller * even smaller test	2025-08-09 17:59:23 +02:00
Sieds LyklesandGitHub	10d388499d	Refactor optional.py (#11578 ) * move fast_idiv to transcendental * move optional.py * adjust comment * change import * mypy needs this?	2025-08-09 17:35:05 +02:00
nimlgenandGitHub	20e46a175c	do not use disk with usb (#11119 ) * not use disk with usb * better name	2025-08-09 11:58:02 +03:00
qazalandGitHub	53179953fc	viz: factor out memory graph render (#11586 )	2025-08-08 20:18:11 +03:00
qazalandGitHub	8ce72d3fad	simpler disassembly table spec (#11583 ) * simpler disassembly table spec * update ui * move to scalar/vec render	2025-08-08 17:59:26 +03:00
qazalandGitHub	44a222a9b2	viz: move resource usage summary to server (#11582 )	2025-08-08 17:08:28 +03:00
qazalandGitHub	793ace530e	update amd_uop_matmul.py import (#11581 ) Using this for testing SQTT	2025-08-08 17:07:35 +03:00
chenyuandGitHub	b232c60def	benchmark openpilot 0.9.9 (#11575 ) * benchmark openpilot 0.9.9 not sure what to do with the 0.9.7 ones with IMAGE=2 and validate * name	2025-08-08 01:26:14 -04:00
qazalandGitHub	16f0edbe90	pass opts arg in get_program process replay [pr] (#11571 ) * fix ptx process replay * keyword arg * renderer is also optional [pr] * test_linearizer fixup * name function order is args,ret,kwargs * can use opts_to_apply * pass through p.applied_opts * sink_arg * now it opens devices too	2025-08-08 03:05:09 +03:00
qazalandGitHub	960cc6533a	pass through name function args in track_rewrites (#11572 )	2025-08-08 02:28:52 +03:00
wozeparrotandGitHub	1826004ef9	feat: add tinyos builder link (#11570 )	2025-08-07 17:42:18 -04:00
George HotzandGitHub	82be8abfd2	move opt under codegen (#11569 )	2025-08-07 14:19:17 -07:00
chenyuandGitHub	702e38dc19	remove FUSE_ARANGE_UINT (#11567 ) also add IGNORE_OOB=1 to bert runs. lowered BS on tinybox to 90 since 96 oom during eval without reset	2025-08-07 16:49:06 -04:00
George HotzandGitHub	6ed2dfd187	delete the arange dim mismatch restriction (#11568 ) * delete the arange dim mismatch restriction * skip that test race	2025-08-07 13:46:17 -07:00
wozeparrotandGitHub	7ae4335127	feat: generate blend index (#11566 )	2025-08-07 14:20:28 -04:00
chenyuandGitHub	594cbdc66f	skip AM ResNet50 benchmark (#11565 ) hanging with FUSE_ARANGE?	2025-08-07 14:07:01 -04:00
chenyuandGitHub	aa1a6f2132	support threshold in Tensor.softplus (#11564 ) fix gradient for large input	2025-08-07 13:43:18 -04:00
chenyu GitHub George Hotz	7ee3770961	FUSE_ARANGE=1 (#11427 ) * FUSE_ARANGE=1 * fix test --------- Co-authored-by: George Hotz <[email protected]>	2025-08-07 13:32:34 -04:00
geohot	4dfcfb1ae5	Revert "Revert "viz: align-center checkbox (#11555 )"" This reverts commit `c52facfd29`.	2025-08-07 08:15:57 -07:00
geohot	7e42427a7b	Revert "Revert "viz: remove color for unbind step (#11554 )"" This reverts commit `5650c7b86c`.	2025-08-07 08:15:51 -07:00
geohot	dc765fbeb7	Revert "viz: timeline perf (#11533 )" This reverts commit `031f26632b`.	2025-08-07 08:08:51 -07:00
geohot	5650c7b86c	Revert "viz: remove color for unbind step (#11554 )" This reverts commit `1e205775bd`.	2025-08-07 08:08:50 -07:00
geohot	c52facfd29	Revert "viz: align-center checkbox (#11555 )" This reverts commit `91ec093464`.	2025-08-07 08:08:50 -07:00
geohot	974cfbe76d	Revert "viz: add support for colored tooltip text (#11556 )" This reverts commit `b3f7ea6f93`.	2025-08-07 08:08:49 -07:00
geohot	3bf0db80ef	Revert "viz: pick the largest rect for proxy fillColor (#11558 )" This reverts commit `76079bc7f2`.	2025-08-07 08:08:48 -07:00
George HotzandGitHub	9764c6cdee	fix mismatch reduce, try 2 (#11560 ) * fix mismatch reduce, try 2 * fix heuristic * delete that test * don't start allowing ones	2025-08-07 07:57:58 -07:00
qazalandGitHub	76079bc7f2	viz: pick the largest rect for proxy fillColor (#11558 )	2025-08-07 16:40:17 +03:00
nimlgenandGitHub	4f29a2c441	fix flaky test on macos (#11557 )	2025-08-07 15:55:35 +03:00
qazalandGitHub	b3f7ea6f93	viz: add support for colored tooltip text (#11556 )	2025-08-07 15:04:43 +03:00
qazalandGitHub	91ec093464	viz: align-center checkbox (#11555 )	2025-08-07 14:22:02 +03:00
qazalandGitHub	1e205775bd	viz: remove color for unbind step (#11554 )	2025-08-07 14:16:21 +03:00
nimlgenandGitHub	031f26632b	viz: timeline perf (#11533 ) * viz: timeline perf * progress * fast * less lines * less lines * less lines * fix chrome	2025-08-07 13:16:17 +03:00
George HotzandGitHub	a1aa5670aa	Revert "fix mismatch reduce (#11547 )" (#11549 ) This reverts commit `49d21a9055`.	2025-08-06 22:43:15 -07:00
George HotzandGitHub	49d21a9055	fix mismatch reduce (#11547 ) * fix mismatch reduce * cleanups * fix shape * fix mypy * resolve	2025-08-06 21:12:51 -07:00
George HotzandGitHub	21570545d3	move view pushing to codegen, try 2 (#11534 ) * move view pushing to codegen, try 2 * fix up some linearizer tests * fix test search * fix test schedule * delete that test * fix test arange * fix a few tests * update tests * push views * ebs cleanup * fix local/reg * test and lint * fix more tests * test cleanups * skipped that one	2025-08-06 15:58:38 -07:00
wozeparrotandGitHub	2d5bdc939d	faster llama3 dataloader (#11540 )	2025-08-06 18:25:57 -04:00
George HotzandGitHub	80d9cced07	more test cleanups (#11544 ) * more test cleanups * revert that	2025-08-06 15:05:21 -07:00
George HotzandGitHub	6fd1332763	update some tests for less Kernel (#11543 ) * update some tests for less Kernel * get_program update	2025-08-06 14:19:59 -07:00
George HotzandGitHub	09dc7af8e9	move bind to big graph (#11539 ) * move bind to big graph * fix tests * unbind inside kernel only * merge views * fix multitensor * failure text change	2025-08-06 13:27:51 -07:00
George HotzandGitHub	7c5e115747	test_mismatch_reduce (#11538 )	2025-08-06 10:02:14 -07:00
George HotzandGitHub	4fe11725c6	pass through sink arg, update linearizer test (#11536 ) * pass through sink arg, update linearizer test * get_program help * bump line count * use new api	2025-08-06 09:48:48 -07:00