split test_advancedindex

* disable cpu_access in the sqtt buffer allocation

not sure if this is required, it results in a very slow call to
pcie_mem_write over USB GPU, removing it worked fine.

* fix itrace_se_mask on gfx12

on gfx11 it gave 6 se, on gfx11 this value is 2 so no instructions were
traced.

* Revert "fix itrace_se_mask on gfx12"

This reverts commit 0644adbcd1.

.github/workflows/test.yml

@@ -238,8 +238,6 @@ jobs:
         pip3 install --upgrade --force-reinstall ruff==0.11.0
         python3 -m ruff check .
         python3 -m ruff check examples/mlperf/ --ignore E501
-    - name: Lint tinygrad with pylint
-      run: python -m pylint tinygrad/
     - name: Run mypy
       run: |
         python -m mypy --strict-equality --lineprecision-report .

.pre-commit-config.yaml

@@ -20,21 +20,15 @@ repos:
         language: system
         always_run: true
         pass_filenames: false
-      - id: tests
-        name: subset of tests
-        entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
-        language: system
-        always_run: true
-        pass_filenames: false
       - id: example
-        name: multi device tests
+        name: test all devices
         entry: python3 test/external/external_test_example.py
         language: system
         always_run: true
         pass_filenames: false
-      - id: pylint
-        name: pylint
-        entry: python3 -m pylint tinygrad/
+      - id: tests
+        name: subset of tests
+        entry: env PYTHONPATH="." python3 -m pytest -n=8 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
         language: system
         always_run: true
         pass_filenames: false

extra/torch_backend/backend.py

@@ -155,16 +155,14 @@ def index_tensor(x, y):
 def zero_(x):
   if TORCH_DEBUG: print(f"zero_ {x.shape}")
   tt = unwrap(x)
-  # NOTE: unconditional contiguous covers if x is contiguous (match it) or if x is view (realize for inplace)
-  # TODO: consolidate
-  tt.assign(tt.zeros_like().contiguous())
+  tt.assign(tt.zeros_like())
 
 @torch.library.impl("aten::fill_.Scalar", "privateuseone")
 @inplace_fn("x")
 def fill_scalar(x, y):
   if TORCH_DEBUG: print(f"fill_.Scalar {x.shape} {y}")
   tt = unwrap(x)
-  tt.assign(tt.full_like(y).contiguous())
+  tt.assign(tt.full_like(y))
 
 @torch.library.impl("aten::_local_scalar_dense", "privateuseone")
 def _local_scalar_dense(tensor): return unwrap(tensor).item()

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/external/external_test_example.py

@@ -58,8 +58,8 @@ class TestExample(unittest.TestCase):
       print(f"WARNING: {device} test isn't running")
       return
 
-    x = Tensor.eye(64, device=device, requires_grad=True)
-    y = Tensor.eye(64, device=device, requires_grad=True)
+    x = Tensor.eye(8, device=device, requires_grad=True)
+    y = Tensor.eye(8, device=device, requires_grad=True)
     z = y.matmul(x).sum()
     z.backward()
 

test/test_assign.py

@@ -129,6 +129,7 @@ class TestAssign(unittest.TestCase):
   @unittest.expectedFailure
   def test_assign_changes_realized_alt(self): return self.test_assign_changes_alt(realize=True)
 
+  @unittest.skip("assign to contiguous shouldn't change the base buffer")
   def test_assign_changes_buffer_alt(self):
     a, b = [Tensor(Tensor(0).contiguous().realize().uop.as_buf()) for _ in range(2)]
     Tensor.realize(a.contiguous().assign(1), b.contiguous().assign(2))

test/test_ops.py

@@ -3177,6 +3177,7 @@ class TestOps(unittest.TestCase):
   def test_bitcast(self):
     helper_test_op([(3, 3)], lambda x: x.view(torch.int32), lambda x: x.bitcast(dtypes.int32), forward_only=True)
 
+  @unittest.skip("we have test_linalg, no need to test here. TODO: should be in torch backend tests")
   def test_svd(self):
     # test for tiny backend. real svd tests are in test_linalg
     A = torch.randn(5, 5)

test/test_outerworld.py

@@ -1,63 +1,7 @@
 import unittest
-from tinygrad import Tensor, UOp, Variable, nn
+from tinygrad import Tensor, UOp
 from tinygrad.uop.ops import AxisType, Ops
 
-class TestOuterworldTrain(unittest.TestCase):
-  @Tensor.train()
-  def test_train(self):
-    # same example over and over
-    X = Tensor.rand(1, 32).expand(16,32).contiguous()
-    Y = Tensor.rand(1, 1).expand(16,1).contiguous()
-
-    layer = nn.Linear(32, 1, bias=False)
-    opt = nn.optim.SGD(nn.state.get_parameters(layer))
-    Tensor.realize(X, Y, *nn.state.get_parameters(layer))
-
-    print("train")
-
-    # if everything is correct, this should be a 16 step training loop
-    steps = UOp.range(16, -1)
-    opt.zero_grad()
-    loss = (layer(X[steps]) - Y[steps]).square().mean().backward()
-    sched = opt.schedule_step() # TODO: does this need to know anything about steps?
-    # NOTE: this can't work. the inputs to layer are not the assign, need to run twice for the fixed point?
-    all_losses = Tensor.realize(loss.reshape(1).expand(steps).contiguous(), *sched)
-    print(all_losses.numpy())
-
-#@unittest.skip("TODO: understand assign")
-class TestOuterworldAssign(unittest.TestCase):
-  def test_triple_add_inner(self):
-    t = Tensor.zeros(5).contiguous().realize()
-    t2 = Tensor.ones(3).contiguous().realize()
-    a = UOp.range(3, -1)
-    t = t.reshape(1,5).expand(a+1,5)[a].assign(t+t2[a])
-    self.assertListEqual(t.tolist(), [3,3,3,3,3])
-
-  def test_triple_add_outer(self):
-    t = Tensor.zeros(5).contiguous().realize()
-    t2 = Tensor.ones(3).contiguous().realize()
-
-    # OUTER is a loop at the schedule level
-    a = UOp.range(3, -1, AxisType.OUTER)
-    va = Variable("loop", 0, 2).bind(a)
-    t = t.assign(t+t2[va])
-    t = Tensor(UOp(Ops.ENDRANGE, dtype=t.uop.dtype, src=(a, t.uop)))
-
-    self.assertListEqual(t.tolist(), [3,3,3,3,3])
-
-  def test_triple_gemm(self):
-    x = Tensor.rand(1, 16).realize()
-    W = Tensor.rand(3, 16, 16).realize()
-
-    #manual = (x @ W[0] @ W[1] @ W[2]).contiguous().realize()
-
-    a = UOp.range(3, -1)
-
-    out = (x @ W[a]).contiguous()
-    t = Tensor(UOp(Ops.ASSIGN, dtype=out.uop.dtype, src=(x.uop, out.uop, a)))
-    #t = Tensor(UOp(Ops.REDUCE, dtype=out.uop.dtype, src=(out.uop, x.uop, a), arg=Ops.NOOP))
-    t.realize()
-
 class TestOuterworldReduce(unittest.TestCase):
   def test_reduce(self):
     x = Tensor.ones(5, 5).contiguous()
@@ -96,7 +40,6 @@ class TestOuterworld(unittest.TestCase):
     # passthrough ranges
     a = UOp.range(10, -1)
     sel = t[9-a]
-    assert sel.shape == (10,)
     cpy = sel.reshape(1, 10).expand(a, 10).contiguous().realize()
 
     self.assertTrue((t.flip(0)==cpy).all().item())

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/test_tiny.py

@@ -134,8 +134,8 @@ class TestTiny(unittest.TestCase):
   def test_mnist_backward(self):
     # NOTE: we don't have the whole model here for speed
     layers = [
-      nn.Conv2d(1, 32, 5), Tensor.relu,
-      nn.Conv2d(32, 32, 5), Tensor.relu]
+      nn.Conv2d(1, 8, 5), Tensor.relu,
+      nn.Conv2d(8, 8, 5), Tensor.relu]
 
     # replace random weights with ones
     # TODO: there's a bug here where it's tying two of the biases together. we need UNIQUE const
@@ -144,7 +144,7 @@ class TestTiny(unittest.TestCase):
 
     # realize gradients
     for x in nn.state.get_parameters(layers): x.requires_grad_()
-    Tensor.empty(4, 1, 28, 28).sequential(layers).sum().backward()
+    Tensor.empty(4, 1, 14, 14).sequential(layers).sum().backward()
     Tensor.realize(*[x.grad for x in nn.state.get_parameters(layers) if x.grad is not None])
 
   # *** image ***

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):

test/unit/test_winograd.py

@@ -42,7 +42,7 @@ class TestWinograd(unittest.TestCase):
     out = Tensor.conv2d(x,w, padding=1)
     out.mean().backward()
     backward_schedule = Tensor.schedule(x.grad, w.grad)
-    self.assertEqual(len(backward_schedule), 5)
+    self.assertEqual(len(backward_schedule), 4)
 
   def test_counters(self):
     IC, OC, X, Y = 4,4,9,9

tinygrad/codegen/late/devectorizer.py

@@ -20,8 +20,8 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
   # can drop valid if idx is out of bound when valid is False
   drop_stmt = []
   for stmt in valid.split_uop(Ops.AND):
-    try: X, is_upper_bound, c = parse_valid(stmt)
-    except ValueError: return None
+    if (res:=parse_valid(stmt)) is None: continue
+    X, is_upper_bound, c = res
 
     # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i
     if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in X.split_uop(Ops.ADD)):

tinygrad/codegen/opt/__init__.py

@@ -18,8 +18,8 @@ class Opt:
 
 axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
                 AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
-axis_colors = {AxisType.OUTER: "GREEN", AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN",
-               AxisType.LOOP: "WHITE", AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
+axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
+               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
 
 class KernelOptError(Exception): pass
 def check(cond:bool, msg:str=""):

tinygrad/codegen/opt/postrange.py

@@ -13,8 +13,8 @@ from tinygrad.renderer import Renderer
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
 # NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
-axis_to_pos = {AxisType.OUTER: -2, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2,
-               AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
+axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
+               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
 
 class Scheduler:
   def __init__(self, ast:UOp, opts:Renderer):

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/__init__.py

@@ -7,7 +7,7 @@ from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, Gro
 from tinygrad.dtype import AddrSpace, PtrDType
 if TYPE_CHECKING:
   from tinygrad.codegen.opt.tc import TensorCore
-  from tinygrad.codegen.opt.kernel import Opt
+  from tinygrad.codegen.opt import Opt
 
 @dataclass(frozen=True)
 class Estimates:

tinygrad/runtime/ops_amd.py

@@ -819,7 +819,7 @@ class AMDDevice(HCQCompiled):
                            f"ppfeaturemask={(ppfeaturemask&~0x8000):#x} (current {ppfeaturemask=:#x} & ~PP_GFXOFF_MASK) to amdgpu module parameters\n"
                            "For more information read https://github.com/tinygrad/tinygrad/blob/master/extra/sqtt/README.md")
       SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine
-      self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE*1024*1024, BufferSpec(cpu_access=True, nolru=True)) for _ in range(self.se_cnt)]
+      self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE*1024*1024, BufferSpec(nolru=True)) for _ in range(self.se_cnt)]
       self.sqtt_itrace_se_mask = getenv("SQTT_ITRACE_SE_MASK", 2) # -1 enable all, 0 disable all, >0 bitmask for where to enable instruction tracing
       self.sqtt_next_cmd_id = itertools.count(0)
       cast(AMDComputeQueue, self.hw_compute_queue_t()).sqtt_start(self.sqtt_buffers, self.sqtt_itrace_se_mask).submit(self)

tinygrad/schedule/indexing.py

@@ -17,7 +17,7 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
     if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
 
 def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
-  #if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
+  if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
   # if it's a kernel, we don't realize it
   if a.src[1].op is not Ops.KERNEL: ctx[a] = None
 
@@ -25,7 +25,7 @@ pm_generate_realize_map = PatternMatcher([
   # always realize SINK src
   (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
   # always realize COPY/BUFFER_VIEW/CONTIGUOUS
-  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.ENDRANGE}, name="tr"), realize),
+  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS}, name="tr"), realize),
   # realize srcs of COPY, MSELECT, MSTACK
   (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
   # realize ASSIGN and input to assign (might be optimized out)
@@ -128,8 +128,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
         axes_out.append(combined_axes % s)
         combined_axes //= s
       # this simplify is doing a lot of heavy lifting. this is the replacement for the reshape view merging code
-      rngs = graph_rewrite(graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic+pm_simplify_valid, name="reshape"),
-        pm_drop_and_clauses, name="reshape drop ands").src
+      rngs = graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic+pm_simplify_valid+pm_drop_and_clauses, name="reshape").src
     case _: raise RuntimeError(f"{op} is not a MovementOp")
   return rngs
 

tinygrad/schedule/rangeify.py

@@ -92,9 +92,6 @@ earliest_rewrites = PatternMatcher([
 
    # realize before assign if input permutes the target buffer
    (UPat(Ops.ASSIGN, src=(UPat.var("a"), UPat.var("b")), name="assign"), find_permutes),
-
-  # contiguous buffer is buffer, this is for *correctness* of assign, not just speed
-  (UPat(Ops.CONTIGUOUS, name="root", src=(UPat(Ops.BUFFER),)), lambda root: root.src[0].forced_reshape(root.shape).rtag(root.tag)),
 ])
 
 # *****************
@@ -110,7 +107,7 @@ pm_mops = PatternMatcher([
 # 3.5 cleanups
 
 # Ops.NOOP happens when we have a COPY to the device the Tensor is already on. We treat it like COPY here for MSTACK.
-ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.NOOP, Ops.ENDRANGE}
+ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.NOOP}
 
 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
 def cleanup_dead_axes(b:UOp):
@@ -338,7 +335,6 @@ def handle_assign(ctx:LocalAddBufferContext, assign:UOp):
 
 def renumber_range(ctx:LocalAddBufferContext, r:UOp):
   if r.tag is not None: return None
-  if r.arg[-1] is AxisType.OUTER: return None
   ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=())
   ctx.range += 1
   return ret
@@ -413,7 +409,7 @@ class Kernel:
     return f"<Kernel {len(list(self.ast.toposort()))} {ast_rep} {self.metadata}>"
 
 def split_store(ctx:list[UOp], x:UOp):
-  if len([r for r in x.ranges if r.arg[-1] != AxisType.OUTER]): return None
+  if len(x.ranges): return None
   if x.src[0].ptrdtype.addrspace is AddrSpace.LOCAL: return None
 
   # local kernel rewrite
@@ -425,7 +421,7 @@ def split_store(ctx:list[UOp], x:UOp):
 
   # NOTE: the hack for COPY is here
   ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts) if lctx.opts is not None else None) \
-    if ret.src[1].op not in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENDRANGE} else ret.src[1]
+    if ret.src[1].op not in {Ops.COPY, Ops.BUFFER_VIEW} else ret.src[1]
   kernel_arg = Kernel(ret,tuple(dedup(flatten([x for x in metadatas if x is not None])))[::-1])
   kernel = UOp(Ops.KERNEL, src=tuple(lctx.map.values())+tuple(lctx.vars.keys()), arg=kernel_arg)
   if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src if x.op is not Ops.BIND]):
@@ -482,11 +478,6 @@ def do_sub_recurse(s:UOp):
   return x.replace(src=tuple([UOp(Ops.SUBSTITUTE, dtype=y.dtype, src=(y,uop_keys,uop_values)) for y in x.src]))
 pm_substitute_recurse = PatternMatcher([(UPat(Ops.SUBSTITUTE, src=(UPat(), UPat(Ops.NOOP), UPat(Ops.NOOP)), name="s"), do_sub_recurse)])
 
-pm_localize_bufs = PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), lambda x:
-   x.replace(arg=BufferizeOpts(device=None, addrspace=AddrSpace.LOCAL), tag=None) if len(x.ranges) > 0 else None),
-])
-
 @track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len([u for u in UOp.sink(*ret.values()).toposort() if u.op is Ops.KERNEL]))}", True)
 def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   uop_list: list[UOp] = []
@@ -502,7 +493,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers")
   # TODO: can you substitute and remove costly buffers at the same time?
   tsink = graph_rewrite(tsink, pm_substitute_recurse, bottom_up=True, name="run substitutes")
-  tsink = graph_rewrite(tsink, pm_localize_bufs+pm_limit_bufs, ctx=rctx, name="localize/limit buffers")
+  tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
 
   # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph
   # MSTACK stacks multiple BUFFERIZEs in one tagged tensor

tinygrad/tensor.py

@@ -19,6 +19,9 @@ from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars
 from tinygrad.schedule.rangeify import get_rangeify_map
 from tinygrad.schedule.multi import get_multi_map
 
+# TODO: this should be the only usage of Device
+def canonicalize_device(device:str|None) -> str: return Device.canonicalize(device)
+
 # *** all in scope Tensors are here. this gets relevant UOps ***
 
 all_tensors: dict[weakref.ref[Tensor], None] = {}
@@ -113,9 +116,10 @@ class Tensor(MathTrait):
 
   def __init__(self, data:ConstType|bytes|list|tuple|UOp|'np.ndarray'|pathlib.Path|None,  # type: ignore [name-defined] # noqa: F821
                device:str|tuple|list|None=None, dtype:DTypeLike|None=None, requires_grad:bool|None=None):
-    if dtype is not None: dtype = to_dtype(dtype)
     if device is None and isinstance(data, pathlib.Path): device = f"DISK:{data.resolve()}"  # keep it on the disk if device is None
-    device = tuple(Device.canonicalize(x) for x in device) if isinstance(device, (tuple, list)) else Device.canonicalize(device)
+    _dtype:DType|None = to_dtype(dtype) if dtype is not None else None
+    _device:str|tuple[str, ...] = tuple(canonicalize_device(x) for x in device) if isinstance(device, (tuple, list)) else canonicalize_device(device)
+    del device, dtype
 
     # tensors can have gradients if you have called .backward
     self.grad:Tensor|None = None
@@ -126,41 +130,41 @@ class Tensor(MathTrait):
 
     # create a UOp from the different types of inputs
     if isinstance(data, UOp):
-      assert dtype is None or dtype==data.dtype, "dtype doesn't match, and casting isn't supported"
+      assert _dtype is None or _dtype==data.dtype, "dtype doesn't match, and casting isn't supported"
       # if data is dtype.index that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of
       if data.dtype==dtypes.index: data = _index_to_concrete_int(data)
       if data.op is Ops.BIND:  # type: ignore  # mypy type narrowing is bugged here
         var, val = data.unbind()  # type: ignore
         # give the bound constant a device
-        const = UOp.const(var.dtype, val, device, ())
+        const = UOp.const(var.dtype, val, _device, ())
         data = data.replace(src=(var.replace(src=const.src), const))  # type: ignore
-    elif data is None: data = UOp.const(dtype or dtypes.default_float, 0, device, ())
-    elif isinstance(data, get_args(ConstType)): data = UOp.const(dtype or dtypes.from_py(data), data, device, ())
-    elif isinstance(data, bytes): data = _frompy(data, dtypes.uint8 if dtype is None else dtype)
+    elif data is None: data = UOp.const(_dtype or dtypes.default_float, 0, _device, ())
+    elif isinstance(data, get_args(ConstType)): data = UOp.const(_dtype or dtypes.from_py(data), data, _device, ())
+    elif isinstance(data, bytes): data = _frompy(data, dtypes.uint8 if _dtype is None else _dtype)
     elif isinstance(data, (list, tuple)):
-      if dtype is None:
-        if (d := fully_flatten(data)) and all(isinstance(s, bool) for s in d): dtype = dtypes.bool
-        else: dtype = dtypes.default_int if d and all_int(d) else dtypes.default_float  # NOTE: this works because all_int([True, False]) is True
-      if dtype in [dtypes.bfloat16, *dtypes.fp8s]: data = Tensor(_frompy(data, dtypes.float32), device=device).cast(dtype).uop
-      else: data = _frompy(data, dtype)
+      if _dtype is None:
+        if (d := fully_flatten(data)) and all(isinstance(s, bool) for s in d): _dtype = dtypes.bool
+        else: _dtype = dtypes.default_int if d and all_int(d) else dtypes.default_float  # NOTE: this works because all_int([True, False]) is True
+      if _dtype in [dtypes.bfloat16, *dtypes.fp8s]: data = Tensor(_frompy(data, dtypes.float32), device=_device).cast(_dtype).uop
+      else: data = _frompy(data, _dtype)
     elif is_numpy_ndarray(data):
       import numpy as np
       assert isinstance(data, np.ndarray), f"expected np.ndarray, got {data}"
-      if data.shape == (): data = UOp.const(dtype or _from_np_dtype(data.dtype), data.item(), device, ())
-      else: data = _fromnp(data.astype(npdtype) if dtype is not None and (npdtype:=_to_np_dtype(dtype)) is not None else data)  # type: ignore [name-defined]
+      if data.shape == (): data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item(), _device, ())
+      else: data = _fromnp(data.astype(npdtype) if _dtype is not None and (npdtype:=_to_np_dtype(_dtype)) is not None else data)  # type: ignore [name-defined]
     elif isinstance(data, pathlib.Path):
-      dtype = dtype or dtypes.uint8
-      data = UOp.new_buffer(f"DISK:{data.resolve()}", data.stat().st_size // dtype.itemsize, dtype)
+      _dtype = _dtype or dtypes.uint8
+      data = UOp.new_buffer(f"DISK:{data.resolve()}", data.stat().st_size // _dtype.itemsize, _dtype)
 
     # by this point, it has to be a UOp
     if not isinstance(data, UOp): raise RuntimeError(f"can't create Tensor from {data!r} with type {type(data)}")
 
     # data might be on a different device
-    if isinstance(device, str): self.uop:UOp = data if data.device == device else data.copy_to_device(device)
+    if isinstance(_device, str): self.uop:UOp = data if data.device == _device else data.copy_to_device(_device)
     # if device is a tuple, we should have/construct a MultiLazyBuffer
-    elif isinstance(data.device, str): self.uop = Tensor(data).shard(device).uop
+    elif isinstance(data.device, str): self.uop = Tensor(data).shard(_device).uop
     else:
-      assert data.device == device, f"MultiLazyBuffer device mismatch, {data.device} != {device}"
+      assert data.device == _device, f"MultiLazyBuffer device mismatch, {data.device} != {_device}"
       self.uop = data
 
     # add to all_tensors after construction succeeds
@@ -376,7 +380,7 @@ class Tensor(MathTrait):
     """
     Moves the tensor to the given device.
     """
-    device = tuple(Device.canonicalize(x) for x in device) if isinstance(device, (tuple, list)) else Device.canonicalize(device)
+    device = tuple(canonicalize_device(x) for x in device) if isinstance(device, (tuple, list)) else canonicalize_device(device)
     if device == self.device: return self
     if not isinstance(device, str): return self.shard(device)
     ret = Tensor(self.uop, device, requires_grad=self.requires_grad)
@@ -401,7 +405,7 @@ class Tensor(MathTrait):
     ```
     """
     assert isinstance(self.device, str), "can't shard a MultiLazyBuffer"
-    devices = tuple(Device.canonicalize(x) for x in devices)
+    devices = tuple(canonicalize_device(x) for x in devices)
     mlb = self.uop.shard(devices, self._resolve_dim(axis)) if axis is not None else self.uop.copy_to_device(devices)
     return Tensor(mlb, device=devices, requires_grad=self.requires_grad)
 
@@ -490,7 +494,7 @@ class Tensor(MathTrait):
     dtype, shape = to_dtype(dtype) if dtype is not None else dtypes.default_float, argfix(*shape)
     if not isinstance(size:=prod([x.vmax if isinstance(x, UOp) else x for x in shape]), int): raise ValueError(f"size must be int {size}")
     # TODO: add test for multidevice tensor
-    device = tuple(Device.canonicalize(d) for d in device) if isinstance(device, tuple) else Device.canonicalize(device)
+    device = tuple(canonicalize_device(d) for d in device) if isinstance(device, tuple) else canonicalize_device(device)
     return Tensor(UOp.new_buffer(device, size, dtype), device, dtype, **kwargs).shrink(((0,prod(shape)),)).reshape(shape)
 
   def empty_like(self, **kwargs) -> Tensor:
@@ -572,7 +576,7 @@ class Tensor(MathTrait):
     if not dtypes.is_float(dtype := to_dtype(dtype or dtypes.default_float)): raise ValueError(f"rand only supports float dtypes, got {dtype}")
     if not all_int(shape:=argfix(*shape)) or not all(s >= 0 for s in shape): raise ValueError(f"invalid input {shape=}")
     if device is not None and not isinstance(device, str): raise ValueError(f"rand only supports single device, got {device=}")
-    device = Device.canonicalize(device)
+    device = canonicalize_device(device)
 
     # if shape has 0, return zero tensor
     if (numel := prod(shape)) == 0: return Tensor.zeros(shape, device=device, dtype=dtype, **kwargs)
@@ -1218,7 +1222,7 @@ class Tensor(MathTrait):
           if not dtypes.is_int((ti:=Tensor(index)).dtype): raise IndexError(f"{index=} contains non-int element")
           index = Tensor([i+size if i<0 else i for i in fully_flatten(index)], self.device, requires_grad=False).reshape(ti.shape)
         case int() | UOp(): # sint
-          #if index >= size or index < -size: raise IndexError(f"{index=} is out of bounds with {size=}")
+          if index >= size or index < -size: raise IndexError(f"{index=} is out of bounds with {size=}")
           # TODO: is this right for (negative) symbolic?
           boundary = [index, index+1] if index >= 0 else [index+size, index+size+1]
         case slice():

tinygrad/uop/ops.py

@@ -15,7 +15,6 @@ if TYPE_CHECKING:
 
 class AxisType(Enum):
   def __repr__(self): return str(self)
-  OUTER = auto()
   GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
   THREAD = auto()
 
@@ -241,7 +240,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
         if s in ret: del ret[s]
     else:
       for s in self.src: ret.update(s.ranges)
-    if self.op is Ops.ENDRANGE: del ret[self.src[0]]
     return ret
 
   @property
@@ -847,17 +845,15 @@ class PatternMatcher:
       if (ret:=match(uop, ctx)) is not None and ret is not uop: return ret
     return None
 
-# *** non-blocking UOp tracker ***
-
-ucount = itertools.count()
-uop_fields:dict[int, tuple] = {}
-def track_uop(u:UOp): return u.trace_num
-
 # *** tracking pattern matcher ***
 
 TRACK_MATCH_STATS = ContextVar("TRACK_MATCH_STATS", 2 if VIZ else 0)
 match_stats:dict[UPat, list[int|float]] = dict()
 
+# TRACK_MATCH_STATS>=2 or VIZ=1 saves all matches
+ucount = itertools.count()
+uop_fields:dict[int, tuple] = {}
+
 @dataclass(frozen=True)
 class TrackedGraphRewrite:
   loc:tuple[str, int]                           # location that called graph_rewrite
@@ -914,7 +910,7 @@ def track_matches(func):
       loc = ((frm:=sys._getframe(1)).f_code.co_filename, frm.f_lineno)
       depth = len(active_rewrites)
       if not tracked_ctxs: add_trace_group(TracingKey(f"default {func.__name__}"))
-      tracked_ctxs[-1].append(ctx:=TrackedGraphRewrite(loc, track_uop(args[0]), [], kwargs.get("name", None), depth, kwargs.get("bottom_up", False)))
+      tracked_ctxs[-1].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], kwargs.get("name", None), depth, kwargs.get("bottom_up", False)))
       active_rewrites.append(ctx)
     with cpu_profile(kwargs.get("name", "<unnamed>"), "TINY", display=tracking):
       ret = func(*args, **kwargs)
@@ -936,14 +932,14 @@ class TrackedPatternMatcher(PatternMatcher):
       try: ret = match(uop, ctx)
       except Exception:
         if TRACK_MATCH_STATS >= 2 and active_rewrites:
-          active_rewrites[-1].matches.append((track_uop(uop), track_uop(UOp(Ops.REWRITE_ERROR,src=uop.src,arg=str(sys.exc_info()[1]))),p.location,0))
+          active_rewrites[-1].matches.append((uop.trace_num, UOp(Ops.REWRITE_ERROR,src=uop.src,arg=str(sys.exc_info()[1])).trace_num,p.location,0))
         raise
       if ret is not None and ret is not uop:
         match_stats[p][0] += 1
         match_stats[p][3] += (et:=time.perf_counter()-st)
         if TRACK_MATCH_STATS >= 3: print(f"{et*1e6:7.2f} us -- ", printable(p.location))
         if TRACK_MATCH_STATS >= 2 and isinstance(ret, UOp) and active_rewrites:
-          active_rewrites[-1].matches.append((track_uop(uop), track_uop(ret), p.location, et))
+          active_rewrites[-1].matches.append((uop.trace_num, ret.trace_num, p.location, et))
         return ret
       match_stats[p][2] += time.perf_counter()-st
     return None
@@ -1100,8 +1096,6 @@ pm_lower_index_dtype = PatternMatcher([
   (UPat((Ops.STORE, Ops.LOAD), src=(UPat(), UPat(), UPat().cast(dtypes.index)), allow_any_len=True, name="s"),
     lambda s: s.replace(src=s.src[:2]+tuple(u.src[0] for u in s.src[2:]))),
   (UPat((Ops.SINK, Ops.NOOP), src=UPat().cast(dtypes.index), name="n"), lambda n: n.replace(src=tuple(s.src[0] for s in n.src))),
-  # hack for ENDRANGE
-  (UPat(Ops.ENDRANGE, src=(UPat(Ops.RANGE, name="r").cast(dtypes.index),), allow_any_len=True, name="x"), lambda x,r: x.replace(src=(r,)+x.src[1:])),
 ])
 def _index_to_concrete_int(u:UOp): return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]
 

tinygrad/uop/spec.py

@@ -109,9 +109,6 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([
   (UPat(Ops.ALLREDUCE, name="red", src=(UPat.var("x"), UPat(Ops.DEVICE))), lambda red,x: red.dtype == x.dtype and isinstance(red.arg, Ops)),
   (UPat(Ops.MULTI, name="multi"), lambda multi: all(x.dtype == multi.dtype for x in multi.src) and isinstance(multi.arg, int)),
 
-  # endrange/reduce for outerworld range work
-  (UPat(Ops.ENDRANGE, src=(UPat(Ops.RANGE),), allow_any_len=True), lambda: True),
-
   # REDUCE with an outerworld range
   (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])),
 ])

tinygrad/uop/symbolic.py

@@ -386,7 +386,7 @@ symbolic_flat = symbolic+PatternMatcher([
 
 # ******** we take a small aside to "simplify_valid" to rewrite valids ********
 
-def parse_valid(valid:UOp) -> tuple[UOp, bool, int]:
+def parse_valid(valid:UOp) -> tuple[UOp, bool, int]|None:
   # if it's X <= c, returns X, True, c
   # if it's X >= c, returns X, False, c
 
@@ -395,7 +395,7 @@ def parse_valid(valid:UOp) -> tuple[UOp, bool, int]:
     (s0:=valid.src[0]).op is Ops.CMPLT and dtypes.is_int(s0.src[0].dtype): return s0.src[0], False, int(s0.src[1].vmin)
   # X < c -> X <= c-1
   if valid.op is Ops.CMPLT and dtypes.is_int(valid.src[0].dtype): return valid.src[0], True, int((valid.src[1]).vmax)-1
-  raise ValueError(f"not able to parse {valid=}")
+  return None
 
 def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
   # return simplified uop (might be the same as input)
@@ -403,8 +403,8 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
   # first, parse valid into {expr: (lower_bound, upper_bound)}
   bounds:defaultdict[UOp, list[ConstType|None]] = defaultdict(lambda: [None, None])
   for stmt in valid.split_uop(Ops.AND):
-    try: expr, is_upper, c = parse_valid(stmt)
-    except ValueError: continue  # give up if we cannot parse the valid
+    if (res:=parse_valid(stmt)) is None: continue
+    expr, is_upper, c = res
     bounds[expr][int(is_upper)] = c
 
   # don't simplify any other gates, can lead to OOB, we substitute them back later
@@ -444,8 +444,7 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
 
 def _valid_priority(v: UOp, valids:list[UOp]):
   # we want valid that's in other valids' parents to be first, so it's more likely the other valids get simplified
-  try: return sum(-1 if parse_valid(v)[0] in other.toposort() else 0 for other in valids)
-  except ValueError: return 0
+  return sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.toposort() else 0 for other in valids)
 
 def simplify_valid(valid:UOp) -> UOp|None:
   if valid.op_in_backward_slice_with_self(Ops.LOAD): return None  # this should only be for indexing, skip if there's a LOAD

tinygrad/viz/serve.py

@@ -287,8 +287,8 @@ def reloader():
       os.execv(sys.executable, [sys.executable] + sys.argv)
     time.sleep(0.1)
 
-def load_pickle(path:pathlib.Path|None) -> list:
-  if path is None or not path.exists(): return []
+def load_pickle(fp:str) -> list:
+  if not (path:=pathlib.Path(fp)).exists(): return []
   with path.open("rb") as f: return pickle.load(f)
 
 # NOTE: using HTTPServer forces a potentially slow socket.getfqdn
@@ -296,8 +296,8 @@ class TCPServerWithReuse(socketserver.TCPServer): allow_reuse_address = True
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  parser.add_argument('--kernels', type=pathlib.Path, help='Path to kernels', default=pathlib.Path(temp("rewrites.pkl", append_user=True)))
-  parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  parser.add_argument('--kernels', type=load_pickle, help='Path to kernels', default=pathlib.Path(temp("rewrites.pkl", append_user=True)))
+  parser.add_argument('--profile', type=load_pickle, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
   args = parser.parse_args()
 
   with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
@@ -308,9 +308,8 @@ if __name__ == "__main__":
   st = time.perf_counter()
   print("*** viz is starting")
 
-  ctxs = get_metadata(load_pickle(args.kernels))
-
-  profile_ret = get_profile(load_pickle(args.profile))
+  ctxs = get_metadata(args.kernels)
+  profile_ret = get_profile(args.profile)
 
   server = TCPServerWithReuse(('', PORT), Handler)
   reloader_thread = threading.Thread(target=reloader)