fix wmma

96 trains fine now

.github/workflows/test.yml

@@ -272,6 +272,8 @@ jobs:
     #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
     - name: Run Clip tests for SD MLPerf on NULL backend
       run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
+    - name: Run AMD emulated BERT training on NULL backend
+      run: EMULATE=AMD_RDNA4 NULL=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
     # TODO: support fake weights
     #- name: Run LLaMA 7B on 4 fake devices
     #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_beam.sh

@@ -2,7 +2,7 @@
 
 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_run.sh

@@ -2,7 +2,7 @@
 
 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh

@@ -5,7 +5,7 @@ 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=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_beam.sh

@@ -2,7 +2,7 @@
 
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_run.sh

@@ -2,7 +2,7 @@
 
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh

@@ -5,7 +5,7 @@ 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=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000

test/external/external_test_dev_var.py

@@ -0,0 +1,39 @@
+import subprocess, unittest, os, sys
+from tinygrad.device import Device
+
+class TestTinygradSlow(unittest.TestCase):
+  def test_env_overwrite_default_device(self):
+    subprocess.run([f'{Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                    shell=True, check=True)
+    subprocess.run([f'DISK=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                    shell=True, check=True)
+    subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                    shell=True, check=True)
+
+    if Device.DEFAULT != "CPU":
+      # setting multiple devices fail
+      with self.assertRaises(subprocess.CalledProcessError):
+        subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                        shell=True, check=True)
+
+      # setting device via DEV
+      subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+      subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+      subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+
+      with self.assertRaises(subprocess.CalledProcessError):
+        subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                        shell=True, check=True)
+
+class TestRunAsModule(unittest.TestCase):
+  def test_module_runs(self):
+    p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+      env={**os.environ, "DEBUG": "1"}, timeout=40,)
+    out = (p.stdout + p.stderr).decode()
+    self.assertEqual(p.returncode, 0, msg=out)
+
+if __name__ == '__main__':
+  unittest.main()

test/test_multitensor.py

@@ -658,7 +658,7 @@ class TestMultiTensor(unittest.TestCase):
 
   # it doesn't work like this anymore
   # NOTE: this never failed in assign_multi, it failed tensor spec because MULTI was never pushed in the graph
-  @unittest.expectedFailure
+  @unittest.skip("this test is broken")
   def test_mlb_assign_change_axis(self):
     t_none = Tensor.zeros((16, 16)).shard(devices_2).contiguous().realize()
     t_zero = Tensor.ones((16, 16)).shard(devices_2, axis=0)

test/test_tensor.py

@@ -1,4 +1,3 @@
-import subprocess
 import numpy as np
 import torch
 import unittest, copy, mmap, random, math, array
@@ -515,32 +514,6 @@ class TestTinygrad(unittest.TestCase):
     print(a)
     print(c)
 
-  def test_env_overwrite_default_device(self):
-    subprocess.run([f'{Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                    shell=True, check=True)
-    subprocess.run([f'DISK=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                    shell=True, check=True)
-    subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                    shell=True, check=True)
-
-    if Device.DEFAULT != "CPU":
-      # setting multiple devices fail
-      with self.assertRaises(subprocess.CalledProcessError):
-        subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                        shell=True, check=True)
-
-      # setting device via DEV
-      subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                      shell=True, check=True)
-      subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                      shell=True, check=True)
-      subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                      shell=True, check=True)
-
-      with self.assertRaises(subprocess.CalledProcessError):
-        subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                        shell=True, check=True)
-
   def test_no_attributeerror_after_apply_uop_exception(self):
     try:
       Tensor.arange(4).reshape(3,2)

test/unit/test_device.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
-import unittest, os, subprocess, sys
+import unittest, os, subprocess
 from tinygrad import Tensor
-from tinygrad.device import Device, Compiler
+from tinygrad.device import Device, Compiler, enumerate_devices_str
 from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, WIN, CI
 
 class TestDevice(unittest.TestCase):
@@ -100,10 +100,7 @@ class TestCompiler(unittest.TestCase):
 
 class TestRunAsModule(unittest.TestCase):
   def test_module_runs(self):
-    p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE,
-      env={**os.environ, "DEBUG": "1"}, timeout=40,)
-    out = (p.stdout + p.stderr).decode()
-    self.assertEqual(p.returncode, 0, msg=out)
+    out = '\n'.join(enumerate_devices_str())
     self.assertIn("CPU", out) # for sanity check
 
 if __name__ == "__main__":

test/unit/test_indexing.py

@@ -180,474 +180,6 @@ class TestIndexing(unittest.TestCase):
     # def delitem(): del reference[0]
     # self.assertRaises(TypeError, delitem)
 
-  # TODO: LLVM is quite fast, why are other compiled backends slow?
-  @unittest.skipIf(CI and Device.DEFAULT in ["CPU", "CL", "METAL", "NV", "AMD"], "slow")
-  def test_advancedindex(self):
-    # integer array indexing
-
-    # pick a random valid indexer type
-    def ri(indices):
-      choice = random.randint(0, 2)
-      if choice == 0: return Tensor(indices)
-      if choice == 1: return list(indices)
-      return tuple(indices)
-
-    def validate_indexing(x):
-      numpy_testing_assert_equal_helper(x[[0]], consec((1,)))
-      numpy_testing_assert_equal_helper(x[ri([0]),], consec((1,)))
-      numpy_testing_assert_equal_helper(x[ri([3]),], consec((1,), 4))
-      numpy_testing_assert_equal_helper(x[[2, 3, 4]], consec((3,), 3))
-      numpy_testing_assert_equal_helper(x[ri([2, 3, 4]),], consec((3,), 3))
-      numpy_testing_assert_equal_helper(x[ri([0, 2, 4]),], np.array([1, 3, 5]))
-
-    def validate_setting(x):
-      x[[0]] = -2
-      numpy_testing_assert_equal_helper(x[[0]], np.array([-2]))
-      x[[0]] = -1
-      numpy_testing_assert_equal_helper(x[ri([0]), ], np.array([-1]))
-      x[[2, 3, 4]] = 4
-      numpy_testing_assert_equal_helper(x[[2, 3, 4]], np.array([4, 4, 4]))
-      x[ri([2, 3, 4]), ] = 3
-      numpy_testing_assert_equal_helper(x[ri([2, 3, 4]), ], np.array([3, 3, 3]))
-      x[ri([0, 2, 4]), ] = Tensor([5, 4, 3])
-      numpy_testing_assert_equal_helper(x[ri([0, 2, 4]), ], np.array([5, 4, 3]))
-
-    # Case 1: Purely Integer Array Indexing
-    reference = consec((10,))
-    validate_indexing(reference)
-    # setting values
-    validate_setting(reference)
-
-    # Tensor with stride != 1
-    # strided is [1, 3, 5, 7]
-
-    # # TODO: set stride
-    # reference = consec((10,))
-    # strided = set_(reference, (4,), (2,), 0)
-
-    # numpy_testing_assert_equal_helper(strided[[0]], np.array([1]))
-    # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([1]))
-    # numpy_testing_assert_equal_helper(strided[ri([3]), ], np.array([7]))
-    # numpy_testing_assert_equal_helper(strided[[1, 2]], np.array([3, 5]))
-    # numpy_testing_assert_equal_helper(strided[ri([1, 2]), ], np.array([3, 5]))
-    # numpy_testing_assert_equal_helper(strided[ri([[2, 1], [0, 3]]), ],
-    #                   np.array([[5, 3], [1, 7]]))
-
-    # stride is [4, 8]
-
-    # strided = set_(reference, (2,), (4,), offset=4)
-
-    # numpy_testing_assert_equal_helper(strided[[0]], np.array([5]))
-    # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([5]))
-    # numpy_testing_assert_equal_helper(strided[ri([1]), ], np.array([9]))
-    # numpy_testing_assert_equal_helper(strided[[0, 1]], np.array([5, 9]))
-    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ], np.array([5, 9]))
-    # numpy_testing_assert_equal_helper(strided[ri([[0, 1], [1, 0]]), ],
-    #                   np.array([[5, 9], [9, 5]]))
-
-    # reference is 1 2
-    #              3 4
-    #              5 6
-    reference = consec((3, 2))
-    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], np.array([1, 3, 5]))
-    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([1])], np.array([2, 4, 6]))
-    numpy_testing_assert_equal_helper(reference[ri([0]), ri([0])], consec((1,)))
-    numpy_testing_assert_equal_helper(reference[ri([2]), ri([1])], consec((1,), 6))
-    numpy_testing_assert_equal_helper(reference[[ri([0, 0]), ri([0, 1])]], np.array([1, 2]))
-    numpy_testing_assert_equal_helper(reference[[ri([0, 1, 1, 0, 2]), ri([1])]], np.array([2, 4, 4, 2, 6]))
-    numpy_testing_assert_equal_helper(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 2, 3, 3]))
-
-    rows = ri([[0, 0],
-               [1, 2]])
-    columns = [0],
-    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[1, 1],
-                                                                          [3, 5]]))
-
-    rows = ri([[0, 0],
-               [1, 2]])
-    columns = ri([1, 0])
-    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[2, 1],
-                                                                          [4, 5]]))
-    rows = ri([[0, 0],
-               [1, 2]])
-    columns = ri([[0, 1],
-                  [1, 0]])
-    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[1, 2],
-                                                                          [4, 5]]))
-
-    # setting values
-    reference[ri([0]), ri([1])] = -1
-    numpy_testing_assert_equal_helper(reference[ri([0]), ri([1])], np.array([-1]))
-    reference[ri([0, 1, 2]), ri([0])] = Tensor([-1, 2, -4])
-    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])],
-                      np.array([-1, 2, -4]))
-    reference[rows, columns] = Tensor([[4, 6], [2, 3]])
-    numpy_testing_assert_equal_helper(reference[rows, columns],
-                      np.array([[4, 6], [2, 3]]))
-
-    # Verify still works with Transposed (i.e. non-contiguous) Tensors
-    reference = Tensor([[0, 1, 2, 3],
-                        [4, 5, 6, 7],
-                        [8, 9, 10, 11]]).T
-
-    # Transposed: [[0, 4, 8],
-    #              [1, 5, 9],
-    #              [2, 6, 10],
-    #              [3, 7, 11]]
-
-    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], np.array([0, 1, 2]))
-    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([1])], np.array([4, 5, 6]))
-    numpy_testing_assert_equal_helper(reference[ri([0]), ri([0])], np.array([0]))
-    numpy_testing_assert_equal_helper(reference[ri([2]), ri([1])], np.array([6]))
-    numpy_testing_assert_equal_helper(reference[[ri([0, 0]), ri([0, 1])]], np.array([0, 4]))
-    numpy_testing_assert_equal_helper(reference[[ri([0, 1, 1, 0, 3]), ri([1])]], np.array([4, 5, 5, 4, 7]))
-    numpy_testing_assert_equal_helper(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([0, 4, 1, 1]))
-
-    rows = ri([[0, 0],
-               [1, 2]])
-    columns = [0],
-    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[0, 0], [1, 2]]))
-
-    rows = ri([[0, 0],
-               [1, 2]])
-    columns = ri([1, 0])
-    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[4, 0], [5, 2]]))
-    rows = ri([[0, 0],
-               [1, 3]])
-    columns = ri([[0, 1],
-                  [1, 2]])
-    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[0, 4], [5, 11]]))
-
-    # TODO: non contiguous setitem
-    '''
-    # setting values
-    reference[ri([0]), ri([1])] = -1
-    numpy_testing_assert_equal_helper(reference[ri([0]), ri([1])],
-                      np.array([-1]))
-    reference[ri([0, 1, 2]), ri([0])] = np.array([-1, 2, -4])
-    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])],
-                      np.array([-1, 2, -4]))
-    reference[rows, columns] = np.array([[4, 6], [2, 3]])
-    numpy_testing_assert_equal_helper(reference[rows, columns],
-                      np.array([[4, 6], [2, 3]]))
-    '''
-
-    # stride != 1
-
-    # strided is [[1 3 5 7],
-    #             [9 11 13 15]]
-
-    # # TODO: set stride
-    # reference = Tensor.arange(0., 24).reshape(3, 8)
-    # strided = set_(reference, (2,4), (8,2), 1)
-
-    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([0])], np.array([1, 9]))
-    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1])], np.array([3, 11]))
-    # numpy_testing_assert_equal_helper(strided[ri([0]), ri([0])], np.array([1]))
-    # numpy_testing_assert_equal_helper(strided[ri([1]), ri([3])], np.array([15]))
-    # numpy_testing_assert_equal_helper(strided[[ri([0, 0]), ri([0, 3])]], np.array([1, 7]))
-    # numpy_testing_assert_equal_helper(strided[[ri([1]), ri([0, 1, 1, 0, 3])]], np.array([9, 11, 11, 9, 15]))
-    # numpy_testing_assert_equal_helper(strided[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 3, 9, 9]))
-
-    # rows = ri([[0, 0],
-    #            [1, 1]])
-    # columns = [0],
-    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 1], [9, 9]]))
-
-    # rows = ri([[0, 1],
-    #            [1, 0]])
-    # columns = ri([1, 2])
-    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[3, 13], [11, 5]]))
-    # rows = ri([[0, 0],
-    #            [1, 1]])
-    # columns = ri([[0, 1],
-    #               [1, 2]])
-    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 3], [11, 13]]))
-
-    # setting values
-
-    # strided is [[10, 11],
-    #             [17, 18]]
-
-    # # TODO: set stride
-    # reference = Tensor.arange(0., 24).reshape(3, 8)
-    # strided = set_(reference, (2,2), (7,1), 10)
-
-    # numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])], np.array([11]))
-
-    # TODO non contiguous setitem
-    '''
-    strided[ri([0]), ri([1])] = -1
-    numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])],
-                      Tensor([-1]))
-    '''
-    # # TODO: set stride
-    # reference = Tensor.arange(0., 24).reshape(3, 8)
-    # strided = set_(reference, (2,2), (7,1), 10)
-
-    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])], np.array([11, 17]))
-
-    # TODO non contiguous setitem
-    '''
-    strided[ri([0, 1]), ri([1, 0])] = Tensor([-1, 2])
-    numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])],
-                      Tensor([-1, 2]))
-    '''
-
-    # # TODO: set stride
-    # reference = Tensor.arange(0., 24).realize().reshape(3, 8)
-    # strided = set_(reference, (2,2), (7,1), 10)
-
-    # rows = ri([[0],
-    #             [1]])
-    # columns = ri([[0, 1],
-    #               [0, 1]])
-    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[10, 11], [17, 18]]))
-
-    # TODO non contiguous setitem
-    '''
-    strided[rows, columns] = Tensor([[4, 6], [2, 3]])
-    numpy_testing_assert_equal_helper(strided[rows, columns],
-                      Tensor([[4, 6], [2, 3]]))
-    '''
-
-    # Tests using less than the number of dims, and ellipsis
-
-    # reference is 1 2
-    #              3 4
-    #              5 6
-    reference = consec((3, 2))
-    numpy_testing_assert_equal_helper(reference[ri([0, 2]),], np.array([[1, 2], [5, 6]]))
-    numpy_testing_assert_equal_helper(reference[ri([1]), ...], np.array([[3, 4]]))
-    numpy_testing_assert_equal_helper(reference[..., ri([1])], np.array([[2], [4], [6]]))
-
-    # verify too many indices fails
-    with self.assertRaises(IndexError): reference[ri([1]), ri([0, 2]), ri([3])]
-
-    # test invalid index fails
-    reference = Tensor.empty(10)
-    for err_idx in (10, -11):
-      with self.assertRaises(IndexError):
-        reference[err_idx]
-      # NOTE cannot check for out of bounds with Tensor indexing
-      # see tensor.py: __getitem__ (Tiny Things)
-      '''
-      with self.assertRaises(IndexError):
-        reference[Tensor([err_idx], dtype=dtypes.int64)]
-      with self.assertRaises(IndexError):
-        reference[[err_idx]]
-      '''
-
-    def tensor_indices_to_np(tensor: Tensor, indices):
-      npt = tensor.numpy()
-      idxs = tuple(i.numpy().tolist() if isinstance(i, Tensor) and i.dtype == dtypes.int64 else
-                  i for i in indices)
-      return npt, idxs
-
-    def get_numpy(tensor, indices):
-      npt, idxs = tensor_indices_to_np(tensor, indices)
-      return Tensor(npt[idxs])
-
-    def set_numpy(tensor:Tensor, indices, value):
-      if not isinstance(value, int):
-        value = value.numpy()
-      npt, idxs = tensor_indices_to_np(tensor, indices)
-      npt[idxs] = value
-      return npt
-
-    def assert_get_eq(tensor, indexer):
-      numpy_testing_assert_equal_helper(tensor[indexer], get_numpy(tensor, indexer))
-
-    def assert_set_eq(tensor: Tensor, indexer, val):
-      pyt = clone(tensor)
-      numt = clone(tensor)
-      pyt[indexer] = val
-      numt = set_numpy(numt, indexer, val)
-      numpy_testing_assert_equal_helper(pyt, numt)
-
-    # NOTE: torch initiates the gradients using g0cpu (rand as gradients)
-    def assert_backward_eq(tensor: Tensor, indexer):
-      cpu = clone(tensor.float())
-      cpu.requires_grad = True
-      outcpu = cpu[indexer].sum()
-      outcpu.backward()
-      dev = cpu.detach()
-      dev.requires_grad = True
-      outdev = dev[indexer].sum()
-      outdev.backward()
-      numpy_testing_assert_equal_helper(cpu.grad, dev.grad)
-
-    def get_set_tensor(indexed: Tensor, indexer):
-      set_size = indexed[indexer].shape
-      set_count = indexed[indexer].numel()
-      set_tensor = Tensor.randint(set_count, high=set_count).reshape(set_size) #.cast(dtypes.float64)
-      return set_tensor
-
-    # Tensor is  0  1  2  3  4
-    #            5  6  7  8  9
-    #           10 11 12 13 14
-    #           15 16 17 18 19
-    reference = Tensor.arange(0., 20).reshape(4, 5)
-
-    indices_to_test = [
-      # grab the second, fourth columns
-      [slice(None), [1, 3]],
-
-      # first, third rows,
-      [[0, 2], slice(None)],
-
-      # weird shape
-      [slice(None), [[0, 1],
-                      [2, 3]]],
-      # negatives
-      [[-1], [0]],
-      [[0, 2], [-1]],
-      [slice(None), [-1]],
-    ]
-
-    # only test dupes on gets
-    get_indices_to_test = indices_to_test + [[slice(None), [0, 1, 1, 2, 2]]]
-
-    for indexer in get_indices_to_test:
-      assert_get_eq(reference, indexer)
-      assert_backward_eq(reference, indexer)
-
-    for indexer in indices_to_test:
-      assert_set_eq(reference, indexer, 44)
-      assert_set_eq(reference, indexer, get_set_tensor(reference, indexer))
-
-    reference = Tensor.arange(0., 160).reshape(4, 8, 5)
-
-    indices_to_test = [
-      [slice(None), slice(None), [0, 3, 4]],
-      [slice(None), [2, 4, 5, 7], slice(None)],
-      [[2, 3], slice(None), slice(None)],
-      [slice(None), [0, 2, 3], [1, 3, 4]],
-      [slice(None), [0], [1, 2, 4]],
-      [slice(None), [0, 1, 3], [4]],
-      [slice(None), [[0, 1], [1, 0]], [[2, 3]]],
-      [slice(None), [[0, 1], [2, 3]], [[0]]],
-      [slice(None), [[5, 6]], [[0, 3], [4, 4]]],
-      [[0, 2, 3], [1, 3, 4], slice(None)],
-      [[0], [1, 2, 4], slice(None)],
-      [[0, 1, 3], [4], slice(None)],
-      [[[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None)],
-      [[[0, 1], [1, 0]], [[2, 3]], slice(None)],
-      [[[0, 1], [2, 3]], [[0]], slice(None)],
-      [[[2, 1]], [[0, 3], [4, 4]], slice(None)],
-      [[[2]], [[0, 3], [4, 1]], slice(None)],
-      # non-contiguous indexing subspace
-      [[0, 2, 3], slice(None), [1, 3, 4]],
-
-      # less dim, ellipsis
-      [[0, 2], ],
-      [[0, 2], slice(None)],
-      [[0, 2], Ellipsis],
-      [[0, 2], slice(None), Ellipsis],
-      [[0, 2], Ellipsis, slice(None)],
-      [[0, 2], [1, 3]],
-      [[0, 2], [1, 3], Ellipsis],
-      [Ellipsis, [1, 3], [2, 3]],
-      [Ellipsis, [2, 3, 4]],
-      [Ellipsis, slice(None), [2, 3, 4]],
-      [slice(None), Ellipsis, [2, 3, 4]],
-
-      # ellipsis counts for nothing
-      [Ellipsis, slice(None), slice(None), [0, 3, 4]],
-      [slice(None), Ellipsis, slice(None), [0, 3, 4]],
-      [slice(None), slice(None), Ellipsis, [0, 3, 4]],
-      [slice(None), slice(None), [0, 3, 4], Ellipsis],
-      [Ellipsis, [[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None)],
-      [[[0, 1], [1, 0]], [[2, 1], [3, 5]], Ellipsis, slice(None)],
-      [[[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None), Ellipsis],
-    ]
-
-    for indexer in indices_to_test:
-      assert_get_eq(reference, indexer)
-
-      assert_set_eq(reference, indexer, 212)
-      assert_set_eq(reference, indexer, get_set_tensor(reference, indexer))
-      assert_backward_eq(reference, indexer)
-
-    reference = Tensor.arange(0., 1296).reshape(3, 9, 8, 6)
-
-    indices_to_test = [
-      [slice(None), slice(None), slice(None), [0, 3, 4]],
-      [slice(None), slice(None), [2, 4, 5, 7], slice(None)],
-      [slice(None), [2, 3], slice(None), slice(None)],
-      [[1, 2], slice(None), slice(None), slice(None)],
-      [slice(None), slice(None), [0, 2, 3], [1, 3, 4]],
-      [slice(None), slice(None), [0], [1, 2, 4]],
-      [slice(None), slice(None), [0, 1, 3], [4]],
-      [slice(None), slice(None), [[0, 1], [1, 0]], [[2, 3]]],
-      [slice(None), slice(None), [[0, 1], [2, 3]], [[0]]],
-      [slice(None), slice(None), [[5, 6]], [[0, 3], [4, 4]]],
-      [slice(None), [0, 2, 3], [1, 3, 4], slice(None)],
-      [slice(None), [0], [1, 2, 4], slice(None)],
-      [slice(None), [0, 1, 3], [4], slice(None)],
-      [slice(None), [[0, 1], [3, 4]], [[2, 3], [0, 1]], slice(None)],
-      [slice(None), [[0, 1], [3, 4]], [[2, 3]], slice(None)],
-      [slice(None), [[0, 1], [3, 2]], [[0]], slice(None)],
-      [slice(None), [[2, 1]], [[0, 3], [6, 4]], slice(None)],
-      [slice(None), [[2]], [[0, 3], [4, 2]], slice(None)],
-      [[0, 1, 2], [1, 3, 4], slice(None), slice(None)],
-      [[0], [1, 2, 4], slice(None), slice(None)],
-      [[0, 1, 2], [4], slice(None), slice(None)],
-      [[[0, 1], [0, 2]], [[2, 4], [1, 5]], slice(None), slice(None)],
-      [[[0, 1], [1, 2]], [[2, 0]], slice(None), slice(None)],
-      [[[2, 2]], [[0, 3], [4, 5]], slice(None), slice(None)],
-      [[[2]], [[0, 3], [4, 5]], slice(None), slice(None)],
-      [slice(None), [3, 4, 6], [0, 2, 3], [1, 3, 4]],
-      [slice(None), [2, 3, 4], [1, 3, 4], [4]],
-      [slice(None), [0, 1, 3], [4], [1, 3, 4]],
-      [slice(None), [6], [0, 2, 3], [1, 3, 4]],
-      [slice(None), [2, 3, 5], [3], [4]],
-      [slice(None), [0], [4], [1, 3, 4]],
-      [slice(None), [6], [0, 2, 3], [1]],
-      [slice(None), [[0, 3], [3, 6]], [[0, 1], [1, 3]], [[5, 3], [1, 2]]],
-      [[2, 2, 1], [0, 2, 3], [1, 3, 4], slice(None)],
-      [[2, 0, 1], [1, 2, 3], [4], slice(None)],
-      [[0, 1, 2], [4], [1, 3, 4], slice(None)],
-      [[0], [0, 2, 3], [1, 3, 4], slice(None)],
-      [[0, 2, 1], [3], [4], slice(None)],
-      [[0], [4], [1, 3, 4], slice(None)],
-      [[1], [0, 2, 3], [1], slice(None)],
-      [[[1, 2], [1, 2]], [[0, 1], [2, 3]], [[2, 3], [3, 5]], slice(None)],
-
-      # less dim, ellipsis
-      [Ellipsis, [0, 3, 4]],
-      [Ellipsis, slice(None), [0, 3, 4]],
-      [Ellipsis, slice(None), slice(None), [0, 3, 4]],
-      [slice(None), Ellipsis, [0, 3, 4]],
-      [slice(None), slice(None), Ellipsis, [0, 3, 4]],
-      [slice(None), [0, 2, 3], [1, 3, 4]],
-      [slice(None), [0, 2, 3], [1, 3, 4], Ellipsis],
-      [Ellipsis, [0, 2, 3], [1, 3, 4], slice(None)],
-      [[0], [1, 2, 4]],
-      [[0], [1, 2, 4], slice(None)],
-      [[0], [1, 2, 4], Ellipsis],
-      [[0], [1, 2, 4], Ellipsis, slice(None)],
-      [[1], ],
-      [[0, 2, 1], [3], [4]],
-      [[0, 2, 1], [3], [4], slice(None)],
-      [[0, 2, 1], [3], [4], Ellipsis],
-      [Ellipsis, [0, 2, 1], [3], [4]],
-    ]
-
-    for indexer in indices_to_test:
-      assert_get_eq(reference, indexer)
-      assert_set_eq(reference, indexer, 1333)
-      assert_set_eq(reference, indexer, get_set_tensor(reference, indexer))
-
-    indices_to_test += [
-      [slice(None), slice(None), [[0, 1], [1, 0]], [[2, 3], [3, 0]]],
-      [slice(None), slice(None), [[2]], [[0, 3], [4, 4]]],
-    ]
-    for indexer in indices_to_test:
-      assert_get_eq(reference, indexer)
-      assert_set_eq(reference, indexer, 1333)
-      assert_backward_eq(reference, indexer)
-
   # TODO setitem backward
   '''
   def test_set_item_to_scalar_tensor(self):
@@ -1568,5 +1100,474 @@ class TestNumpy(unittest.TestCase):
     numpy_testing_assert_equal_helper(kernel, kernel2)
   '''
 
+def tensor_indices_to_np(tensor: Tensor, indices):
+  npt = tensor.numpy()
+  idxs = tuple(i.numpy().tolist() if isinstance(i, Tensor) and i.dtype == dtypes.int64 else
+              i for i in indices)
+  return npt, idxs
+
+def get_numpy(tensor, indices):
+  npt, idxs = tensor_indices_to_np(tensor, indices)
+  return Tensor(npt[idxs])
+
+def set_numpy(tensor:Tensor, indices, value):
+  if not isinstance(value, int):
+    value = value.numpy()
+  npt, idxs = tensor_indices_to_np(tensor, indices)
+  npt[idxs] = value
+  return npt
+
+def assert_get_eq(tensor, indexer):
+  numpy_testing_assert_equal_helper(tensor[indexer], get_numpy(tensor, indexer))
+
+def assert_set_eq(tensor: Tensor, indexer, val):
+  pyt = clone(tensor)
+  numt = clone(tensor)
+  pyt[indexer] = val
+  numt = set_numpy(numt, indexer, val)
+  numpy_testing_assert_equal_helper(pyt, numt)
+
+# NOTE: torch initiates the gradients using g0cpu (rand as gradients)
+def assert_backward_eq(tensor: Tensor, indexer):
+  cpu = clone(tensor.float())
+  cpu.requires_grad = True
+  outcpu = cpu[indexer].sum()
+  outcpu.backward()
+  dev = cpu.detach()
+  dev.requires_grad = True
+  outdev = dev[indexer].sum()
+  outdev.backward()
+  numpy_testing_assert_equal_helper(cpu.grad, dev.grad)
+
+def get_set_tensor(indexed: Tensor, indexer):
+  set_size = indexed[indexer].shape
+  set_count = indexed[indexer].numel()
+  set_tensor = Tensor.randint(set_count, high=set_count).reshape(set_size) #.cast(dtypes.float64)
+  return set_tensor
+
+@unittest.skipIf(CI and Device.DEFAULT in ["CPU", "CL", "METAL", "NV", "AMD"], "slow")
+class TestAdvancedIndexing(unittest.TestCase):
+  def test_integer_array_indexing(self):
+    # pick a random valid indexer type
+    def ri(indices):
+      choice = random.randint(0, 2)
+      if choice == 0: return Tensor(indices)
+      if choice == 1: return list(indices)
+      return tuple(indices)
+
+    def validate_indexing(x):
+      numpy_testing_assert_equal_helper(x[[0]], consec((1,)))
+      numpy_testing_assert_equal_helper(x[ri([0]),], consec((1,)))
+      numpy_testing_assert_equal_helper(x[ri([3]),], consec((1,), 4))
+      numpy_testing_assert_equal_helper(x[[2, 3, 4]], consec((3,), 3))
+      numpy_testing_assert_equal_helper(x[ri([2, 3, 4]),], consec((3,), 3))
+      numpy_testing_assert_equal_helper(x[ri([0, 2, 4]),], np.array([1, 3, 5]))
+
+    def validate_setting(x):
+      x[[0]] = -2
+      numpy_testing_assert_equal_helper(x[[0]], np.array([-2]))
+      x[[0]] = -1
+      numpy_testing_assert_equal_helper(x[ri([0]), ], np.array([-1]))
+      x[[2, 3, 4]] = 4
+      numpy_testing_assert_equal_helper(x[[2, 3, 4]], np.array([4, 4, 4]))
+      x[ri([2, 3, 4]), ] = 3
+      numpy_testing_assert_equal_helper(x[ri([2, 3, 4]), ], np.array([3, 3, 3]))
+      x[ri([0, 2, 4]), ] = Tensor([5, 4, 3])
+      numpy_testing_assert_equal_helper(x[ri([0, 2, 4]), ], np.array([5, 4, 3]))
+
+    # Case 1: Purely Integer Array Indexing
+    reference = consec((10,))
+    validate_indexing(reference)
+    # setting values
+    validate_setting(reference)
+
+    # Tensor with stride != 1
+    # strided is [1, 3, 5, 7]
+
+    # # TODO: set stride
+    # reference = consec((10,))
+    # strided = set_(reference, (4,), (2,), 0)
+
+    # numpy_testing_assert_equal_helper(strided[[0]], np.array([1]))
+    # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([1]))
+    # numpy_testing_assert_equal_helper(strided[ri([3]), ], np.array([7]))
+    # numpy_testing_assert_equal_helper(strided[[1, 2]], np.array([3, 5]))
+    # numpy_testing_assert_equal_helper(strided[ri([1, 2]), ], np.array([3, 5]))
+    # numpy_testing_assert_equal_helper(strided[ri([[2, 1], [0, 3]]), ],
+    #                   np.array([[5, 3], [1, 7]]))
+
+    # stride is [4, 8]
+
+    # strided = set_(reference, (2,), (4,), offset=4)
+
+    # numpy_testing_assert_equal_helper(strided[[0]], np.array([5]))
+    # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([5]))
+    # numpy_testing_assert_equal_helper(strided[ri([1]), ], np.array([9]))
+    # numpy_testing_assert_equal_helper(strided[[0, 1]], np.array([5, 9]))
+    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ], np.array([5, 9]))
+    # numpy_testing_assert_equal_helper(strided[ri([[0, 1], [1, 0]]), ],
+    #                   np.array([[5, 9], [9, 5]]))
+
+    # reference is 1 2
+    #              3 4
+    #              5 6
+    reference = consec((3, 2))
+    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], np.array([1, 3, 5]))
+    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([1])], np.array([2, 4, 6]))
+    numpy_testing_assert_equal_helper(reference[ri([0]), ri([0])], consec((1,)))
+    numpy_testing_assert_equal_helper(reference[ri([2]), ri([1])], consec((1,), 6))
+    numpy_testing_assert_equal_helper(reference[[ri([0, 0]), ri([0, 1])]], np.array([1, 2]))
+    numpy_testing_assert_equal_helper(reference[[ri([0, 1, 1, 0, 2]), ri([1])]], np.array([2, 4, 4, 2, 6]))
+    numpy_testing_assert_equal_helper(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 2, 3, 3]))
+
+    rows = ri([[0, 0],
+               [1, 2]])
+    columns = [0],
+    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[1, 1],
+                                                                          [3, 5]]))
+
+    rows = ri([[0, 0],
+               [1, 2]])
+    columns = ri([1, 0])
+    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[2, 1],
+                                                                          [4, 5]]))
+    rows = ri([[0, 0],
+               [1, 2]])
+    columns = ri([[0, 1],
+                  [1, 0]])
+    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[1, 2],
+                                                                          [4, 5]]))
+
+    # setting values
+    reference[ri([0]), ri([1])] = -1
+    numpy_testing_assert_equal_helper(reference[ri([0]), ri([1])], np.array([-1]))
+    reference[ri([0, 1, 2]), ri([0])] = Tensor([-1, 2, -4])
+    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])],
+                      np.array([-1, 2, -4]))
+    reference[rows, columns] = Tensor([[4, 6], [2, 3]])
+    numpy_testing_assert_equal_helper(reference[rows, columns],
+                      np.array([[4, 6], [2, 3]]))
+
+    # Verify still works with Transposed (i.e. non-contiguous) Tensors
+    reference = Tensor([[0, 1, 2, 3],
+                        [4, 5, 6, 7],
+                        [8, 9, 10, 11]]).T
+
+    # Transposed: [[0, 4, 8],
+    #              [1, 5, 9],
+    #              [2, 6, 10],
+    #              [3, 7, 11]]
+
+    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], np.array([0, 1, 2]))
+    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([1])], np.array([4, 5, 6]))
+    numpy_testing_assert_equal_helper(reference[ri([0]), ri([0])], np.array([0]))
+    numpy_testing_assert_equal_helper(reference[ri([2]), ri([1])], np.array([6]))
+    numpy_testing_assert_equal_helper(reference[[ri([0, 0]), ri([0, 1])]], np.array([0, 4]))
+    numpy_testing_assert_equal_helper(reference[[ri([0, 1, 1, 0, 3]), ri([1])]], np.array([4, 5, 5, 4, 7]))
+    numpy_testing_assert_equal_helper(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([0, 4, 1, 1]))
+
+    rows = ri([[0, 0],
+               [1, 2]])
+    columns = [0],
+    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[0, 0], [1, 2]]))
+
+    rows = ri([[0, 0],
+               [1, 2]])
+    columns = ri([1, 0])
+    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[4, 0], [5, 2]]))
+    rows = ri([[0, 0],
+               [1, 3]])
+    columns = ri([[0, 1],
+                  [1, 2]])
+    numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[0, 4], [5, 11]]))
+
+    # TODO: non contiguous setitem
+    '''
+    # setting values
+    reference[ri([0]), ri([1])] = -1
+    numpy_testing_assert_equal_helper(reference[ri([0]), ri([1])],
+                      np.array([-1]))
+    reference[ri([0, 1, 2]), ri([0])] = np.array([-1, 2, -4])
+    numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])],
+                      np.array([-1, 2, -4]))
+    reference[rows, columns] = np.array([[4, 6], [2, 3]])
+    numpy_testing_assert_equal_helper(reference[rows, columns],
+                      np.array([[4, 6], [2, 3]]))
+    '''
+
+    # stride != 1
+
+    # strided is [[1 3 5 7],
+    #             [9 11 13 15]]
+
+    # # TODO: set stride
+    # reference = Tensor.arange(0., 24).reshape(3, 8)
+    # strided = set_(reference, (2,4), (8,2), 1)
+
+    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([0])], np.array([1, 9]))
+    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1])], np.array([3, 11]))
+    # numpy_testing_assert_equal_helper(strided[ri([0]), ri([0])], np.array([1]))
+    # numpy_testing_assert_equal_helper(strided[ri([1]), ri([3])], np.array([15]))
+    # numpy_testing_assert_equal_helper(strided[[ri([0, 0]), ri([0, 3])]], np.array([1, 7]))
+    # numpy_testing_assert_equal_helper(strided[[ri([1]), ri([0, 1, 1, 0, 3])]], np.array([9, 11, 11, 9, 15]))
+    # numpy_testing_assert_equal_helper(strided[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 3, 9, 9]))
+
+    # rows = ri([[0, 0],
+    #            [1, 1]])
+    # columns = [0],
+    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 1], [9, 9]]))
+
+    # rows = ri([[0, 1],
+    #            [1, 0]])
+    # columns = ri([1, 2])
+    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[3, 13], [11, 5]]))
+    # rows = ri([[0, 0],
+    #            [1, 1]])
+    # columns = ri([[0, 1],
+    #               [1, 2]])
+    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 3], [11, 13]]))
+
+    # setting values
+
+    # strided is [[10, 11],
+    #             [17, 18]]
+
+    # # TODO: set stride
+    # reference = Tensor.arange(0., 24).reshape(3, 8)
+    # strided = set_(reference, (2,2), (7,1), 10)
+
+    # numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])], np.array([11]))
+
+    # TODO non contiguous setitem
+    '''
+    strided[ri([0]), ri([1])] = -1
+    numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])],
+                      Tensor([-1]))
+    '''
+    # # TODO: set stride
+    # reference = Tensor.arange(0., 24).reshape(3, 8)
+    # strided = set_(reference, (2,2), (7,1), 10)
+
+    # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])], np.array([11, 17]))
+
+    # TODO non contiguous setitem
+    '''
+    strided[ri([0, 1]), ri([1, 0])] = Tensor([-1, 2])
+    numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])],
+                      Tensor([-1, 2]))
+    '''
+
+    # # TODO: set stride
+    # reference = Tensor.arange(0., 24).realize().reshape(3, 8)
+    # strided = set_(reference, (2,2), (7,1), 10)
+
+    # rows = ri([[0],
+    #             [1]])
+    # columns = ri([[0, 1],
+    #               [0, 1]])
+    # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[10, 11], [17, 18]]))
+
+    # TODO non contiguous setitem
+    '''
+    strided[rows, columns] = Tensor([[4, 6], [2, 3]])
+    numpy_testing_assert_equal_helper(strided[rows, columns],
+                      Tensor([[4, 6], [2, 3]]))
+    '''
+
+    # Tests using less than the number of dims, and ellipsis
+
+    # reference is 1 2
+    #              3 4
+    #              5 6
+    reference = consec((3, 2))
+    numpy_testing_assert_equal_helper(reference[ri([0, 2]),], np.array([[1, 2], [5, 6]]))
+    numpy_testing_assert_equal_helper(reference[ri([1]), ...], np.array([[3, 4]]))
+    numpy_testing_assert_equal_helper(reference[..., ri([1])], np.array([[2], [4], [6]]))
+
+    # verify too many indices fails
+    with self.assertRaises(IndexError): reference[ri([1]), ri([0, 2]), ri([3])]
+
+    # test invalid index fails
+    reference = Tensor.empty(10)
+    for err_idx in (10, -11):
+      with self.assertRaises(IndexError):
+        reference[err_idx]
+      # NOTE cannot check for out of bounds with Tensor indexing
+      # see tensor.py: __getitem__ (Tiny Things)
+      '''
+      with self.assertRaises(IndexError):
+        reference[Tensor([err_idx], dtype=dtypes.int64)]
+      with self.assertRaises(IndexError):
+        reference[[err_idx]]
+      '''
+
+  def test_numpy_parity_and_backward_2d(self):
+    # Tensor is  0  1  2  3  4
+    #            5  6  7  8  9
+    #           10 11 12 13 14
+    #           15 16 17 18 19
+    reference = Tensor.arange(0., 20).reshape(4, 5)
+
+    indices_to_test = [
+      # grab the second, fourth columns
+      [slice(None), [1, 3]],
+
+      # first, third rows,
+      [[0, 2], slice(None)],
+
+      # weird shape
+      [slice(None), [[0, 1],
+                      [2, 3]]],
+      # negatives
+      [[-1], [0]],
+      [[0, 2], [-1]],
+      [slice(None), [-1]],
+    ]
+
+    # only test dupes on gets
+    get_indices_to_test = indices_to_test + [[slice(None), [0, 1, 1, 2, 2]]]
+
+    for indexer in get_indices_to_test:
+      assert_get_eq(reference, indexer)
+      assert_backward_eq(reference, indexer)
+
+    for indexer in indices_to_test:
+      assert_set_eq(reference, indexer, 44)
+      assert_set_eq(reference, indexer, get_set_tensor(reference, indexer))
+
+  def test_numpy_parity_and_backward_3d(self):
+    reference = Tensor.arange(0., 160).reshape(4, 8, 5)
+
+    indices_to_test = [
+      [slice(None), slice(None), [0, 3, 4]],
+      [slice(None), [2, 4, 5, 7], slice(None)],
+      [[2, 3], slice(None), slice(None)],
+      [slice(None), [0, 2, 3], [1, 3, 4]],
+      [slice(None), [0], [1, 2, 4]],
+      [slice(None), [0, 1, 3], [4]],
+      [slice(None), [[0, 1], [1, 0]], [[2, 3]]],
+      [slice(None), [[0, 1], [2, 3]], [[0]]],
+      [slice(None), [[5, 6]], [[0, 3], [4, 4]]],
+      [[0, 2, 3], [1, 3, 4], slice(None)],
+      [[0], [1, 2, 4], slice(None)],
+      [[0, 1, 3], [4], slice(None)],
+      [[[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None)],
+      [[[0, 1], [1, 0]], [[2, 3]], slice(None)],
+      [[[0, 1], [2, 3]], [[0]], slice(None)],
+      [[[2, 1]], [[0, 3], [4, 4]], slice(None)],
+      [[[2]], [[0, 3], [4, 1]], slice(None)],
+      # non-contiguous indexing subspace
+      [[0, 2, 3], slice(None), [1, 3, 4]],
+
+      # less dim, ellipsis
+      [[0, 2], ],
+      [[0, 2], slice(None)],
+      [[0, 2], Ellipsis],
+      [[0, 2], slice(None), Ellipsis],
+      [[0, 2], Ellipsis, slice(None)],
+      [[0, 2], [1, 3]],
+      [[0, 2], [1, 3], Ellipsis],
+      [Ellipsis, [1, 3], [2, 3]],
+      [Ellipsis, [2, 3, 4]],
+      [Ellipsis, slice(None), [2, 3, 4]],
+      [slice(None), Ellipsis, [2, 3, 4]],
+
+      # ellipsis counts for nothing
+      [Ellipsis, slice(None), slice(None), [0, 3, 4]],
+      [slice(None), Ellipsis, slice(None), [0, 3, 4]],
+      [slice(None), slice(None), Ellipsis, [0, 3, 4]],
+      [slice(None), slice(None), [0, 3, 4], Ellipsis],
+      [Ellipsis, [[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None)],
+      [[[0, 1], [1, 0]], [[2, 1], [3, 5]], Ellipsis, slice(None)],
+      [[[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None), Ellipsis],
+    ]
+
+    for indexer in indices_to_test:
+      assert_get_eq(reference, indexer)
+
+      assert_set_eq(reference, indexer, 212)
+      assert_set_eq(reference, indexer, get_set_tensor(reference, indexer))
+      assert_backward_eq(reference, indexer)
+
+  def test_numpy_parity_and_backward_4d(self):
+    reference = Tensor.arange(0., 1296).reshape(3, 9, 8, 6)
+
+    indices_to_test = [
+      [slice(None), slice(None), slice(None), [0, 3, 4]],
+      [slice(None), slice(None), [2, 4, 5, 7], slice(None)],
+      [slice(None), [2, 3], slice(None), slice(None)],
+      [[1, 2], slice(None), slice(None), slice(None)],
+      [slice(None), slice(None), [0, 2, 3], [1, 3, 4]],
+      [slice(None), slice(None), [0], [1, 2, 4]],
+      [slice(None), slice(None), [0, 1, 3], [4]],
+      [slice(None), slice(None), [[0, 1], [1, 0]], [[2, 3]]],
+      [slice(None), slice(None), [[0, 1], [2, 3]], [[0]]],
+      [slice(None), slice(None), [[5, 6]], [[0, 3], [4, 4]]],
+      [slice(None), [0, 2, 3], [1, 3, 4], slice(None)],
+      [slice(None), [0], [1, 2, 4], slice(None)],
+      [slice(None), [0, 1, 3], [4], slice(None)],
+      [slice(None), [[0, 1], [3, 4]], [[2, 3], [0, 1]], slice(None)],
+      [slice(None), [[0, 1], [3, 4]], [[2, 3]], slice(None)],
+      [slice(None), [[0, 1], [3, 2]], [[0]], slice(None)],
+      [slice(None), [[2, 1]], [[0, 3], [6, 4]], slice(None)],
+      [slice(None), [[2]], [[0, 3], [4, 2]], slice(None)],
+      [[0, 1, 2], [1, 3, 4], slice(None), slice(None)],
+      [[0], [1, 2, 4], slice(None), slice(None)],
+      [[0, 1, 2], [4], slice(None), slice(None)],
+      [[[0, 1], [0, 2]], [[2, 4], [1, 5]], slice(None), slice(None)],
+      [[[0, 1], [1, 2]], [[2, 0]], slice(None), slice(None)],
+      [[[2, 2]], [[0, 3], [4, 5]], slice(None), slice(None)],
+      [[[2]], [[0, 3], [4, 5]], slice(None), slice(None)],
+      [slice(None), [3, 4, 6], [0, 2, 3], [1, 3, 4]],
+      [slice(None), [2, 3, 4], [1, 3, 4], [4]],
+      [slice(None), [0, 1, 3], [4], [1, 3, 4]],
+      [slice(None), [6], [0, 2, 3], [1, 3, 4]],
+      [slice(None), [2, 3, 5], [3], [4]],
+      [slice(None), [0], [4], [1, 3, 4]],
+      [slice(None), [6], [0, 2, 3], [1]],
+      [slice(None), [[0, 3], [3, 6]], [[0, 1], [1, 3]], [[5, 3], [1, 2]]],
+      [[2, 2, 1], [0, 2, 3], [1, 3, 4], slice(None)],
+      [[2, 0, 1], [1, 2, 3], [4], slice(None)],
+      [[0, 1, 2], [4], [1, 3, 4], slice(None)],
+      [[0], [0, 2, 3], [1, 3, 4], slice(None)],
+      [[0, 2, 1], [3], [4], slice(None)],
+      [[0], [4], [1, 3, 4], slice(None)],
+      [[1], [0, 2, 3], [1], slice(None)],
+      [[[1, 2], [1, 2]], [[0, 1], [2, 3]], [[2, 3], [3, 5]], slice(None)],
+
+      # less dim, ellipsis
+      [Ellipsis, [0, 3, 4]],
+      [Ellipsis, slice(None), [0, 3, 4]],
+      [Ellipsis, slice(None), slice(None), [0, 3, 4]],
+      [slice(None), Ellipsis, [0, 3, 4]],
+      [slice(None), slice(None), Ellipsis, [0, 3, 4]],
+      [slice(None), [0, 2, 3], [1, 3, 4]],
+      [slice(None), [0, 2, 3], [1, 3, 4], Ellipsis],
+      [Ellipsis, [0, 2, 3], [1, 3, 4], slice(None)],
+      [[0], [1, 2, 4]],
+      [[0], [1, 2, 4], slice(None)],
+      [[0], [1, 2, 4], Ellipsis],
+      [[0], [1, 2, 4], Ellipsis, slice(None)],
+      [[1], ],
+      [[0, 2, 1], [3], [4]],
+      [[0, 2, 1], [3], [4], slice(None)],
+      [[0, 2, 1], [3], [4], Ellipsis],
+      [Ellipsis, [0, 2, 1], [3], [4]],
+    ]
+
+    for indexer in indices_to_test:
+      assert_get_eq(reference, indexer)
+      assert_set_eq(reference, indexer, 1333)
+      assert_set_eq(reference, indexer, get_set_tensor(reference, indexer))
+
+    indices_to_test += [
+      [slice(None), slice(None), [[0, 1], [1, 0]], [[2, 3], [3, 0]]],
+      [slice(None), slice(None), [[2]], [[0, 3], [4, 4]]],
+    ]
+    for indexer in indices_to_test:
+      assert_get_eq(reference, indexer)
+      assert_set_eq(reference, indexer, 1333)
+      assert_backward_eq(reference, indexer)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_linalg.py

@@ -26,18 +26,22 @@ class TestLinAlg(unittest.TestCase):
       orthogonality_helper(V)
       reconstruction_helper([U,s_diag,V],a)
 
-  def test_svd_nonfull(self):
-    sizes = [(2,2),(5,3),(3,5),(2,2,2,2,3)]
-    for size in sizes:
-      a = Tensor.randn(size).realize()
-      U,S,V = a.svd(full_matrices=False)
-      b_shape,m,n = size[0:-2],size[-2],size[-1]
-      k = min(m,n)
-      s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)).expand(b_shape + (k,k)))
-      #reduced U,V is only orthogonal along smaller dim
-      if (m < n): orthogonality_helper(U),orthogonality_helper(V)
-      else: orthogonality_helper(U.transpose(-2,-1)),orthogonality_helper(V.transpose(-2,-1))
-      reconstruction_helper([U,s_diag,V],a)
+  def _test_svd_nonfull(self, size):
+    a = Tensor.randn(size).realize()
+    U,S,V = a.svd(full_matrices=False)
+    b_shape,m,n = size[0:-2],size[-2],size[-1]
+    k = min(m,n)
+    s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)).expand(b_shape + (k,k)))
+    #reduced U,V is only orthogonal along smaller dim
+    if (m < n): orthogonality_helper(U),orthogonality_helper(V)
+    else: orthogonality_helper(U.transpose(-2,-1)),orthogonality_helper(V.transpose(-2,-1))
+    reconstruction_helper([U,s_diag,V],a)
+
+  # faster for parallel pytest
+  def test_svd_nonfull_2_2(self): self._test_svd_nonfull((2,2))
+  def test_svd_nonfull_5_3(self): self._test_svd_nonfull((5,3))
+  def test_svd_nonfull_3_5(self): self._test_svd_nonfull((3,5))
+  def test_svd_nonfull_2_2_2_2_3(self): self._test_svd_nonfull((2,2,2,2,3))
 
   @unittest.skip("very big. recommend wrapping with TinyJit around inner function")
   def test_svd_large(self):

tinygrad/codegen/simplify.py

@@ -99,14 +99,10 @@ pm_reduce_collapse = PatternMatcher([
   # MUL casted bool
   ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast().or_broadcasted(name="b")),
    lambda x,gate,b=None: gate.broadcast(x.dtype.count).where(x, 0) if b is not None else gate.where(x, 0)),
-  # WHERE on LOAD (works on max too)
-  (UPat.var("gate").where(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).load(), 0).reduce(arg=Ops.ADD, allow_any_len=True),
-   lambda buf,idx,gate: buf.index(idx.valid(gate)).load()),
-  (UPat.var("gate").where(0, UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).load()).reduce(arg=Ops.ADD, allow_any_len=True),
-   lambda buf,idx,gate: buf.index(idx.valid(gate.logical_not())).load()),
-  # INDEX on RANGE / gated RANGE
-  (UPat.var("buf").index(UPat.var("idx").eq(UPat(Ops.RANGE, name="r").or_casted()).where(UPat.var("expr"), invalid_pat)),
-   lambda buf,r,idx,expr,i: buf.index(expr.substitute({r:idx.cast(r.dtype)}).valid((idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])))),
+  # reduce on gated load becomes can substitute the range and remove the reduce
+  (UPat.var("buf").index(UPat.var("idx").eq(UPat(Ops.RANGE, name="r").or_casted()).where(UPat.var("expr"), invalid_pat)).load()
+    .reduce(arg=Ops.ADD, allow_any_len=True), lambda buf,r,idx,expr,i:
+    buf.index(expr.substitute({r:idx.cast(r.dtype)}).valid((idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0]))).load()),
   # AND on WHERE
   ((UPat.any(UPat(Ops.DEFINE_VAR, name="x"), UPat(Ops.DEFINE_VAR).gep(name="x")) & UPat.var("y")) \
    .where(UPat.cvar("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),

tinygrad/device.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
-from typing import Any, Generic, TypeVar, Iterator, Sequence, cast
+from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM
 from tinygrad.helpers import Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
@@ -357,7 +357,7 @@ if PROFILE:
     from tinygrad.uop.ops import launch_viz
     launch_viz("PROFILE", fn)
 
-if __name__ == "__main__":
+def enumerate_devices_str() -> Generator[str, None, None]:
   from tinygrad import Tensor, Device
 
   for device in ALL_DEVICES:
@@ -376,4 +376,7 @@ if __name__ == "__main__":
       result = (colored('PASS', 'green') if any_works else f"{colored('FAIL', 'yellow')}") + ''.join([f'\n{" "*16} {x}' for x in compilers_results])
     except Exception as e:
       result = f"{colored('FAIL', 'red')} {e}"
-    print(f"{'*' if device == Device.DEFAULT else ' '} {device:10s}: {result}")
+    yield f"{'*' if device == Device.DEFAULT else ' '} {device:10s}: {result}"
+
+if __name__ == "__main__":
+  for s in enumerate_devices_str(): print(s)

tinygrad/renderer/llvmir.py

@@ -248,7 +248,7 @@ class AMDLLVMRenderer(LLVMRenderer):
         (UPat(Ops.WMMA, name="x"), lambda x: UOp(Ops.WMMA, x.dtype, (x.src[0].bitcast(dtypes.uint16.vec(16)), x.src[1].bitcast(dtypes.uint16.vec(16)),
           x.src[2]), x.arg) if x.src[0].dtype == dtypes.bfloat16.vec(16) else None),
       ])
-    if self.arch.split(":")[0] == "gfx1201":
+    if self.arch.split(":")[0] in {"gfx1200", "gfx1201"}:
       self.extra_matcher += PatternMatcher([
         (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16.vec(8)), lambda x: UOp(Ops.WMMA, dtypes.uint16.vec(8),
           (x.src[0].bitcast(dtypes.uint16.vec(8)), x.src[1].bitcast(dtypes.uint16.vec(8)), x.src[2].bitcast(dtypes.uint16.vec(8))), (*x.arg,))

tinygrad/runtime/ops_amd.py

@@ -458,7 +458,7 @@ class AMDProgram(HCQProgram):
       if typ == 5: image[apply_image_offset:apply_image_offset+8] = struct.pack('<q', rel_sym_offset - apply_image_offset + addent) # R_AMDGPU_REL64
       else: raise RuntimeError(f"unknown AMD reloc {typ}")
 
-    self.lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000), buf_spec:=BufferSpec(cpu_access=True, nolru=True))
+    self.lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000), buf_spec:=BufferSpec(nolru=True))
     self.dev.allocator._copyin(self.lib_gpu, image)
     self.dev.synchronize()
 

tinygrad/runtime/ops_null.py

@@ -1,9 +1,11 @@
 import functools
+from typing import cast
 from tinygrad.device import Compiled, Compiler, Allocator
 from tinygrad.engine.jit import MultiGraphRunner
-from tinygrad.renderer.cstyle import CStyleLanguage
+from tinygrad.renderer.cstyle import Renderer, CStyleLanguage
+from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.uop.ops import Ops
-from tinygrad.helpers import cpu_profile
+from tinygrad.helpers import cpu_profile, EMULATE
 
 class NullRenderer(CStyleLanguage):
   device = "NULL"
@@ -29,5 +31,11 @@ class NullGraph(MultiGraphRunner):
   def __call__(self, input_rawbuffers, var_vals, wait=False) -> float|None: return 1e-3
 
 class NullDevice(Compiled):
-  def __init__(self, device:str): super().__init__(device, NullAllocator(self), [(NullRenderer, Compiler)], functools.partial(NullProgram, device),
-                                                   NullGraph)
+  def __init__(self, device:str):
+    renderer:functools.partial|type[Renderer]
+    match cast(str, EMULATE.value):
+      case "AMD": renderer = functools.partial(AMDLLVMRenderer, "gfx1100")
+      case "AMD_RDNA4": renderer = functools.partial(AMDLLVMRenderer, "gfx1201")
+      case "": renderer = NullRenderer
+      case _: raise RuntimeError(f"can't EMULATE device: {EMULATE.value}")
+    super().__init__(device, NullAllocator(self), [(renderer, Compiler)], functools.partial(NullProgram, device), NullGraph)

tinygrad/runtime/support/hcq.py

@@ -310,7 +310,7 @@ class HCQProgram(Generic[HCQDeviceType]):
     Returns:
       Arguments state with the given buffers and values set for the program.
     """
-    argsbuf = kernargs or self.dev.kernargs_buf.offset(offset=self.dev.kernargs_offset_allocator.alloc(self.kernargs_alloc_size),
+    argsbuf = kernargs or self.dev.kernargs_buf.offset(offset=self.dev.kernargs_offset_allocator.alloc(self.kernargs_alloc_size, 8),
                                                        size=self.kernargs_alloc_size)
     return self.args_state_t(argsbuf, self, bufs, vals=vals)
 

tinygrad/schedule/rangeify.py

@@ -60,7 +60,7 @@ earliest_rewrites = PatternMatcher([
    lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if x.size == 0 and reduce.size != 0 else None),
 
   # handle size 0
-  (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x.st is not None and x.size == 0 else None),
+  (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and x.size == 0 else None),
 
   # remove contiguous on movement ops before a copy on disk
   (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, allow_any_len=True, name="copy"),

tinygrad/uop/ops.py

@@ -175,6 +175,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   # *** uop shape stuff ***
 
+  # TODO: remove this. it's used by the jit and split_reduceop
   @recursive_property
   def st(self) -> ShapeTracker|None:
     if self.op is Ops.INDEX and self.src[0].op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.MSTACK,
@@ -223,12 +224,98 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
         shape = tuple(1 if i in axis_arg else s for i,s in enumerate(shape))
     return ShapeTracker.from_shape(shape)
 
+  @recursive_property
+  def _shape(self) -> tuple[sint, ...]|None:
+    match self.op:
+      # late ops don't have shape
+      case Ops.UNIQUE | Ops.DEVICE | Ops.RANGE | Ops.INDEX | Ops.LOAD | Ops.IF | Ops.BARRIER | \
+           Ops.VECTORIZE | Ops.VCONST | Ops.SUBSTITUTE | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.PRECAST:
+        return None
+
+      # some ops init the shape
+      case Ops.CONST | Ops.DEFINE_VAR | Ops.BIND: return () if self._device is not None else None
+      case Ops.BUFFER: return (self.arg,)
+      case Ops.BUFFER_VIEW: return (self.arg[0],)
+      case Ops.BUFFERIZE: return tuple([int(r.vmax+1) for r in self.src[1:]])
+      case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
+
+      # passthrough ops
+      case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.FUSE: return self.src[0]._shape
+
+      # ops with custom handling
+      case Ops.KERNEL: return self.arg.ast._shape
+      case Ops.STORE:
+        if isinstance(self.dtype, PtrDType): return (self.ptrdtype.size,)
+        if self.dtype is not dtypes.void: return self.src[0].src[0].shape
+        return None
+
+      # TODO: disallow shape changing bitcast
+      case Ops.BITCAST:
+        ps = self.src[0]._shape
+        if ps is None: return None
+        if (output_sz:=self.dtype.itemsize) != (input_sz:=self.src[0].dtype.itemsize): return ps[:-1]+(ssimplify((ps[-1]*input_sz) // output_sz),)
+        return ps
+
+      # TODO: disallow reshape from nothing. tested by TestOpenClip.test_multigpu_clip_score
+      case Ops.RESHAPE:
+        if self.src[0]._shape is None: return tuple(ssimplify(s) for s in self.arg)
+
+    # movement ops change the shape. this is the logic from the old ShapeTracker
+    # NOTE: ssimplify is required because the shape needs to be canonical for broadcasting and same shape checking
+    if self.op in GroupOp.Movement.union({Ops.MULTI, Ops.REDUCE_AXIS, Ops.WMMA}):
+      ps = self.src[0]._shape
+      # TODO: WMMA is used for both axis WMMA and op WMMA. fix this and remove this hack. tested by BERT on AMD LLVM
+      if ps is None and self.op is Ops.WMMA: return None
+      if ps is None: raise RuntimeError(f"movement op {self.op} requires shape")
+      match self.op:
+        case Ops.RESHAPE:
+          if not all(x >= 0 for x in self.arg): raise ValueError(f"shape can't contain negative numbers {self.arg}")
+          if prod(ps) != prod(self.arg): raise ValueError(f"bad reshape: {ps} -> {self.arg}")
+          return tuple(ssimplify(s) for s in self.arg)
+        case Ops.EXPAND:
+          if len(ps) != len(self.arg) or not all(s==ns or (s==1 and ns>=0) for s,ns in zip(ps, self.arg)):
+            raise ValueError(f"bad expand: {ps} -> {self.arg}")
+          return tuple(ssimplify(s) for s in self.arg)
+        case Ops.PERMUTE:
+          if sorted(self.arg) != list(range(len(ps))): raise ValueError(f"invalid permutation {self.arg} of len {len(ps)}")
+          return tuple(ps[i] for i in self.arg)
+        case Ops.PAD:
+          # TODO: why do i need resolve here?
+          if len(ps) != len(self.arg) or not all(resolve(b>=0) and resolve(e>=0) for b,e in self.arg): raise ValueError(f"invalid pad {self.arg}")
+          return tuple(ssimplify(s+b+e) for s,(b,e) in zip(ps, self.arg))
+        case Ops.SHRINK:
+          # TODO: why do i need resolve here?
+          if len(ps) != len(self.arg) or not all(resolve(0<=b) and resolve(b<=e) and resolve(e<=s) for s,(b,e) in zip(ps, self.arg)):
+            raise ValueError(f"invalid shrink {self.arg} for {ps}")
+          return tuple(ssimplify(e-s) for s,e in self.arg)
+        case Ops.FLIP:
+          if len(ps) != len(self.arg) or not all(isinstance(x, bool) for x in self.arg): raise ValueError(f"bad flip on {ps}, {self.arg}")
+          return ps
+        case Ops.MULTI: return tuple(s*len(self.device) if a == self.axis else s for a,s in enumerate(ps))
+        case Ops.REDUCE_AXIS | Ops.WMMA:
+          axis_arg = self.arg[1] if self.op is Ops.REDUCE_AXIS else self.arg[7]
+          if not isinstance(axis_arg, tuple) or not all(isinstance(x, int) and x>=0 and x<len(ps) for x in axis_arg):
+            raise ValueError(f"invalid type for axis: {axis_arg}")
+          return tuple(1 if i in axis_arg else s for i,s in enumerate(ps))
+
+    # elementwise ops keep the shape the same. all inputs with shape must match
+    if self.op in (GroupOp.Elementwise-{Ops.BITCAST}).union({Ops.COPY, Ops.ASSIGN, Ops.NOOP, Ops.SINK, Ops.ALLREDUCE}):
+      # TODO: remove this hack for 3 op assign
+      input_shapes = [x._shape for x in (self.src[:2] if self.op is Ops.ASSIGN else self.src) if x._shape is not None]
+      if len(input_shapes) == 0: return None
+      if not all_same(input_shapes): raise RuntimeError(f"shape mismatch at {self.op}: {input_shapes}")
+      return input_shapes[0]
+
+    # all Ops must be explicitly handled
+    raise NotImplementedError(f"no shape handling for {self.op} with {self.dtype}")
+
   @property
   def shape(self) -> tuple[sint, ...]:
-    assert self.st is not None, f"{self.op} doesn't have a shape"
-    return unwrap(self.st).shape
+    if (ret:=self._shape) is None: raise RuntimeError(f"shape requested, but {self.op} doesn't have a shape")
+    return ret
+
   @property
-  def size(self) -> int: return self.arg[0] if self.op is Ops.BUFFER_VIEW else self.arg if self.op is Ops.BUFFER else unwrap(self.st).size
+  def size(self) -> int: return prod([int(x.vmax) if isinstance(x, UOp) else x for x in self.shape])
 
   # determine what ranges this is in
   @recursive_property
@@ -290,7 +377,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def __getitem__(self, idx): return self.index(idx)
   def const_like(self, b:ConstLike):
     # constants can optionally have a DEVICE source
-    return UOp.const(self.dtype, b, device=self._device, shape=self.shape if self.st is not None else None)
+    return UOp.const(self.dtype, b, device=self._device, shape=self._shape)
   def broadcast(self, count:int):
     assert self.dtype.count == 1
     if count == 1: return self
@@ -344,7 +431,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     return ret.reshape(tuple([x if i not in axis else 1 for i,x in enumerate(self.shape)]))
   @staticmethod
   def invalid(count=1): return UOp(Ops.CONST, dtypes.index.vec(count), src=(), arg=Invalid)
-  def valid(self, cond): return cond.where(self, UOp.invalid(self.dtype.count))
+  def valid(self, cond): return self if cond.op is Ops.WHERE and cond.arg else cond.where(self, UOp.invalid(self.dtype.count))
   def get_idx(self) -> UOp:
     assert self.dtype.scalar() is dtypes.index, "Can only call get_idx on index dtype"
     return self.src[1] if self.op is Ops.WHERE and self.src[2].arg is Invalid else self
@@ -428,19 +515,22 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if self.op is Ops.MULTI: return self.src[0].base  # MULTI is really a VIEW
     return self
 
-  def _mop(self, op:Ops, arg) -> UOp:
+  def _mop(self, op:Ops, arg, no_reshape_is_no_op:bool=False) -> UOp:
     ret = UOp(op, self.dtype, (self,), arg)
-    if self.st == ret.st: return self  # ignore NOOPs, also check ret.st
+    # for all movement ops, we check shape property
+    if ret.shape == self.shape and no_reshape_is_no_op: return self
     return ret
-
   def forced_reshape(self, arg:tuple[sint, ...], **kwargs): return UOp(Ops.RESHAPE, kwargs.pop("dtype", self.dtype), src=(self,), arg=arg)
 
-  def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg)
-  def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg)
-  def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg)
-  def pad(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.PAD, arg)
-  def permute(self, arg:tuple[int, ...]): return self._mop(Ops.PERMUTE, arg)
-  def flip(self, arg:tuple[bool, ...]): return self._mop(Ops.FLIP, arg)
+  # in these four, if the shape doesn't change we can return self
+  def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, no_reshape_is_no_op=True)
+  def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg, no_reshape_is_no_op=True)
+  def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg, no_reshape_is_no_op=True)
+  def pad(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.PAD, arg, no_reshape_is_no_op=True)
+
+  # in these two, we have custom logic to check if they are a no-op
+  def permute(self, arg:tuple[int, ...]): return self._mop(Ops.PERMUTE, arg) if arg != tuple(range(len(self.shape))) else self
+  def flip(self, arg:tuple[bool, ...]): return self._mop(Ops.FLIP, arg) if any(arg) and len(arg) == len(self.shape) else self
 
   # *** uop UNIQUE ***
 

tinygrad/uop/symbolic.py

@@ -473,6 +473,17 @@ def drop_and_clauses(cond:UOp, x:UOp, i:UOp) -> UOp|None:
   if not (dropped_clauses:=[c for c in cond.split_uop(Ops.AND) if not any(r in x.ranges for r in c.ranges)]): return None
   return functools.reduce(operator.and_, [c for c in cond.split_uop(Ops.AND) if c not in dropped_clauses], UOp.const(dtypes.bool, True)).where(x, i)
 pm_drop_and_clauses = PatternMatcher([(UPat.var("cond").where(UPat.var("x", dtype=dtypes.index), invalid_pat), drop_and_clauses)])
+def where_on_load(l, c1, buf, x):
+  c2 = x.get_valid()
+  duplicate_clauses = [c for c in c1.split_uop(Ops.AND) if c in c2.split_uop(Ops.AND)]
+  # we move the condition from the where to the load _as long as_ the condtition doesn't have some range that would place it inside of a new range
+  # also no data dependent loads!
+  moved_clauses = [c for c in c1.split_uop(Ops.AND) if c not in duplicate_clauses and all(r in x.ranges for r in c.ranges)
+    and not c.op_in_backward_slice_with_self(Ops.LOAD)]
+  if not (removed:=moved_clauses+duplicate_clauses): return None
+  # aditionally we can drop the clause on the where if it already exists in the load
+  remaining_clause = functools.reduce(operator.and_, [c for c in c1.split_uop(Ops.AND) if c not in removed], UOp.const(dtypes.bool, True))
+  return remaining_clause.where(UOp.load(buf.index(x.get_idx().valid(functools.reduce(operator.and_, moved_clauses, c2)), *l.src[1:])), 0)
 
 pm_simplify_valid = PatternMatcher([
   # simplify valid
@@ -518,8 +529,9 @@ sym = symbolic_flat+pm_simplify_valid+PatternMatcher([
   (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.index, Invalid)).or_casted(),), allow_any_len=True, name="x"),
     lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0
   # # Where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer
-  (UPat.var("c1").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("c2").where(UPat(), invalid_pat)).or_casted(),), name="l"), 0),
-    lambda c1,c2,l,i: l.replace(src=(l.src[0],)+l.src[1:]) if all(c in list(c2.split_uop(Ops.AND)) for c in c1.split_uop(Ops.AND)) else None),
+  (UPat.var("c1").where(UPat(Ops.LOAD, src=(UPat.var("buf").index(UPat.var("x")),), name="l"), 0), where_on_load),
+  (UPat.var("c1").where(0, UPat(Ops.LOAD, src=(UPat.var("buf").index(UPat.var("x")),), name="l")),
+    lambda l,c1,buf,x: where_on_load(l,c1.logical_not(),buf,x)),
   # remove VECTORIZE from SINK/BARRIER. TODO: SINK/BARRIER are really the same thing at GLOBAL/LOCAL levels
   (UPat(Ops.BARRIER, name="root"),
     lambda root: UOp(Ops.BARRIER, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_BARRIER else (x,) for x in root.src)), root.arg)