tinygrad 0.8.0__tar.gz → 0.9.1__tar.gz

{tinygrad-0.8.0 → tinygrad-0.9.1}/LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2023 George Hotz
+Copyright (c) 2024, the tiny corp
 
 Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 

{tinygrad-0.8.0 → tinygrad-0.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tinygrad
-Version: 0.8.0
+Version: 0.9.1
 Summary: You like pytorch? You like micrograd? You love tinygrad! <3
 Author: George Hotz
 License: MIT
@@ -10,8 +10,6 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy
-Requires-Dist: tqdm
-Requires-Dist: gpuctypes
 Requires-Dist: pyobjc-framework-Metal; platform_system == "Darwin"
 Requires-Dist: pyobjc-framework-libdispatch; platform_system == "Darwin"
 Provides-Extra: llvm
@@ -32,10 +30,11 @@ Requires-Dist: torch; extra == "testing"
 Requires-Dist: pillow; extra == "testing"
 Requires-Dist: pytest; extra == "testing"
 Requires-Dist: pytest-xdist; extra == "testing"
-Requires-Dist: onnx==1.15.0; extra == "testing"
+Requires-Dist: onnx==1.16.0; extra == "testing"
 Requires-Dist: onnx2torch; extra == "testing"
 Requires-Dist: opencv-python; extra == "testing"
 Requires-Dist: tabulate; extra == "testing"
+Requires-Dist: tqdm; extra == "testing"
 Requires-Dist: safetensors; extra == "testing"
 Requires-Dist: transformers; extra == "testing"
 Requires-Dist: sentencepiece; extra == "testing"
@@ -43,16 +42,30 @@ Requires-Dist: tiktoken; extra == "testing"
 Requires-Dist: librosa; extra == "testing"
 Requires-Dist: networkx; extra == "testing"
 Requires-Dist: hypothesis; extra == "testing"
+Requires-Dist: nibabel; extra == "testing"
+Requires-Dist: bottle; extra == "testing"
+Provides-Extra: docs
+Requires-Dist: mkdocs-material; extra == "docs"
+Requires-Dist: mkdocstrings[python]; extra == "docs"
+Requires-Dist: markdown-callouts; extra == "docs"
+Requires-Dist: markdown-exec[ansi]; extra == "docs"
+Requires-Dist: black; extra == "docs"
+Provides-Extra: testing-tf
+Requires-Dist: tensorflow==2.15.1; extra == "testing-tf"
+Requires-Dist: tensorflow_addons; extra == "testing-tf"
 
 <div align="center">
 
-[![logo](https://raw.githubusercontent.com/tinygrad/tinygrad/master/docs/logo.png)](https://tinygrad.org)
+<picture>
+  <source media="(prefers-color-scheme: light)" srcset="/docs/logo_tiny_light.svg">
+  <img alt="tiny corp logo" src="/docs/logo_tiny_dark.svg" width="50%" height="50%">
+</picture>
 
 tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) and [karpathy/micrograd](https://github.com/karpathy/micrograd). Maintained by [tiny corp](https://tinygrad.org).
 
 <h3>
 
-[Homepage](https://github.com/tinygrad/tinygrad) | [Documentation](/docs) | [Examples](/examples) | [Showcase](/docs/showcase.md) | [Discord](https://discord.gg/ZjZadyC7PK)
+[Homepage](https://github.com/tinygrad/tinygrad) | [Documentation](https://docs.tinygrad.org/) | [Discord](https://discord.gg/ZjZadyC7PK)
 
 </h3>
 
@@ -122,17 +135,15 @@ See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full vers
 
 tinygrad already supports numerous accelerators, including:
 
-- [x] [CPU](tinygrad/runtime/ops_cpu.py)
 - [x] [GPU (OpenCL)](tinygrad/runtime/ops_gpu.py)
-- [x] [C Code (Clang)](tinygrad/runtime/ops_clang.py)
+- [x] [CLANG (C Code)](tinygrad/runtime/ops_clang.py)
 - [x] [LLVM](tinygrad/runtime/ops_llvm.py)
 - [x] [METAL](tinygrad/runtime/ops_metal.py)
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
-- [x] [PyTorch](tinygrad/runtime/ops_torch.py)
-- [x] [HIP](tinygrad/runtime/ops_hip.py)
+- [x] [AMD](tinygrad/runtime/ops_amd.py)
+- [x] [NV](tinygrad/runtime/ops_nv.py)
 
 And it is easy to add more! Your accelerator of choice only needs to support a total of ~25 low level ops.
-More information can be found in the [documentation for adding new accelerators](/docs/adding_new_accelerators.md).
 
 ## Installation
 
@@ -154,7 +165,7 @@ python3 -m pip install git+https://github.com/tinygrad/tinygrad.git
 
 ## Documentation
 
-Documentation along with a quick start guide can be found in the [docs/](/docs) directory.
+Documentation along with a quick start guide can be found on the [docs website](https://docs.tinygrad.org/) built from the [docs/](/docs) directory.
 
 ### Quick example comparing to PyTorch
 
@@ -193,13 +204,14 @@ We'll start with what will get your PR closed with a pointer to this section:
 - All docs and whitespace changes will be closed unless you are a well-known contributor. The people writing the docs should be those who know the codebase the absolute best. People who have not demonstrated that shouldn't be messing with docs. Whitespace changes are both useless *and* carry a risk of introducing bugs.
 - Anything you claim is a "speedup" must be benchmarked. In general, the goal is simplicity, so even if your PR makes things marginally faster, you have to consider the tradeoff with maintainablity and readablity.
 - In general, the code outside the core `tinygrad/` folder is not well tested, so unless the current code there is broken, you shouldn't be changing it.
+- If your PR looks "complex", is a big diff, or adds lots of lines, it won't be reviewed or merged. Consider breaking it up into smaller PRs that are individually clear wins. A common pattern I see is prerequisite refactors before adding new functionality. If you can (cleanly) refactor to the point that the feature is a 3 line change, this is great, and something easy for us to review.
 
 Now, what we want:
 
 - Bug fixes (with a regression test) are great! This library isn't 1.0 yet, so if you stumble upon a bug, fix it, write a test, and submit a PR, this is valuable work.
 - Solving bounties! tinygrad [offers cash bounties](https://docs.google.com/spreadsheets/d/1WKHbT-7KOgjEawq5h5Ic1qUWzpfAzuD_J06N1JwOCGs/edit?usp=sharing) for certain improvements to the library. All new code should be high quality and well tested.
 - Features. However, if you are adding a feature, consider the line tradeoff. If it's 3 lines, there's less of a bar of usefulness it has to meet over something that's 30 or 300 lines. All features must have regression tests. In general with no other constraints, your feature's API should match torch or numpy.
-- Refactors that are clear wins. In general, if your refactor isn't a clear win it will be closed. But some refactors are amazing! Think about readability in a deep core sense. A whitespace change or moving a few functions around is useless, but if you realize that two 100 line functions can actually use the same 110 line function with arguments while also improving readability, this is a big win.
+- Refactors that are clear wins. In general, if your refactor isn't a clear win it will be closed. But some refactors are amazing! Think about readability in a deep core sense. A whitespace change or moving a few functions around is useless, but if you realize that two 100 line functions can actually use the same 110 line function with arguments while also improving readability, this is a big win. Refactors should pass [process replay](#process-replay-tests).
 - Tests/fuzzers. If you can add tests that are non brittle, they are welcome. We have some fuzzers in here too, and there's a plethora of bugs that can be found with them and by improving them. Finding bugs, even writing broken tests (that should pass) with `@unittest.expectedFailure` is great. This is how we make progress.
 - Dead code removal from core `tinygrad/` folder. We don't care about the code in extra, but removing dead code from the core library is great. Less for new people to read and be confused by.
 
@@ -215,3 +227,9 @@ python3 -m pip install -e '.[testing]'  # install extra deps for testing
 python3 test/test_ops.py                # just the ops tests
 python3 -m pytest test/                 # whole test suite
 ```
+
+#### Process replay tests
+
+[Process replay](https://github.com/tinygrad/tinygrad/blob/master/test/external/process_replay/process_replay.py) detects changes in the generated kernels of CI tests by comparing them against tinygrad master. If your PR is a refactor or speedup without any expected behavior change, it should include a green process replay pass to get merged.
+
+You can enable process replay by adding [run_process_replay] to your PR title. [example](https://github.com/tinygrad/tinygrad/pull/4995). Note that you should keep your branch up-to-date with master.

{tinygrad-0.8.0 → tinygrad-0.9.1}/README.md

@@ -1,12 +1,15 @@
 <div align="center">
 
-[![logo](https://raw.githubusercontent.com/tinygrad/tinygrad/master/docs/logo.png)](https://tinygrad.org)
+<picture>
+  <source media="(prefers-color-scheme: light)" srcset="/docs/logo_tiny_light.svg">
+  <img alt="tiny corp logo" src="/docs/logo_tiny_dark.svg" width="50%" height="50%">
+</picture>
 
 tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) and [karpathy/micrograd](https://github.com/karpathy/micrograd). Maintained by [tiny corp](https://tinygrad.org).
 
 <h3>
 
-[Homepage](https://github.com/tinygrad/tinygrad) | [Documentation](/docs) | [Examples](/examples) | [Showcase](/docs/showcase.md) | [Discord](https://discord.gg/ZjZadyC7PK)
+[Homepage](https://github.com/tinygrad/tinygrad) | [Documentation](https://docs.tinygrad.org/) | [Discord](https://discord.gg/ZjZadyC7PK)
 
 </h3>
 
@@ -76,17 +79,15 @@ See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full vers
 
 tinygrad already supports numerous accelerators, including:
 
-- [x] [CPU](tinygrad/runtime/ops_cpu.py)
 - [x] [GPU (OpenCL)](tinygrad/runtime/ops_gpu.py)
-- [x] [C Code (Clang)](tinygrad/runtime/ops_clang.py)
+- [x] [CLANG (C Code)](tinygrad/runtime/ops_clang.py)
 - [x] [LLVM](tinygrad/runtime/ops_llvm.py)
 - [x] [METAL](tinygrad/runtime/ops_metal.py)
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
-- [x] [PyTorch](tinygrad/runtime/ops_torch.py)
-- [x] [HIP](tinygrad/runtime/ops_hip.py)
+- [x] [AMD](tinygrad/runtime/ops_amd.py)
+- [x] [NV](tinygrad/runtime/ops_nv.py)
 
 And it is easy to add more! Your accelerator of choice only needs to support a total of ~25 low level ops.
-More information can be found in the [documentation for adding new accelerators](/docs/adding_new_accelerators.md).
 
 ## Installation
 
@@ -108,7 +109,7 @@ python3 -m pip install git+https://github.com/tinygrad/tinygrad.git
 
 ## Documentation
 
-Documentation along with a quick start guide can be found in the [docs/](/docs) directory.
+Documentation along with a quick start guide can be found on the [docs website](https://docs.tinygrad.org/) built from the [docs/](/docs) directory.
 
 ### Quick example comparing to PyTorch
 
@@ -147,13 +148,14 @@ We'll start with what will get your PR closed with a pointer to this section:
 - All docs and whitespace changes will be closed unless you are a well-known contributor. The people writing the docs should be those who know the codebase the absolute best. People who have not demonstrated that shouldn't be messing with docs. Whitespace changes are both useless *and* carry a risk of introducing bugs.
 - Anything you claim is a "speedup" must be benchmarked. In general, the goal is simplicity, so even if your PR makes things marginally faster, you have to consider the tradeoff with maintainablity and readablity.
 - In general, the code outside the core `tinygrad/` folder is not well tested, so unless the current code there is broken, you shouldn't be changing it.
+- If your PR looks "complex", is a big diff, or adds lots of lines, it won't be reviewed or merged. Consider breaking it up into smaller PRs that are individually clear wins. A common pattern I see is prerequisite refactors before adding new functionality. If you can (cleanly) refactor to the point that the feature is a 3 line change, this is great, and something easy for us to review.
 
 Now, what we want:
 
 - Bug fixes (with a regression test) are great! This library isn't 1.0 yet, so if you stumble upon a bug, fix it, write a test, and submit a PR, this is valuable work.
 - Solving bounties! tinygrad [offers cash bounties](https://docs.google.com/spreadsheets/d/1WKHbT-7KOgjEawq5h5Ic1qUWzpfAzuD_J06N1JwOCGs/edit?usp=sharing) for certain improvements to the library. All new code should be high quality and well tested.
 - Features. However, if you are adding a feature, consider the line tradeoff. If it's 3 lines, there's less of a bar of usefulness it has to meet over something that's 30 or 300 lines. All features must have regression tests. In general with no other constraints, your feature's API should match torch or numpy.
-- Refactors that are clear wins. In general, if your refactor isn't a clear win it will be closed. But some refactors are amazing! Think about readability in a deep core sense. A whitespace change or moving a few functions around is useless, but if you realize that two 100 line functions can actually use the same 110 line function with arguments while also improving readability, this is a big win.
+- Refactors that are clear wins. In general, if your refactor isn't a clear win it will be closed. But some refactors are amazing! Think about readability in a deep core sense. A whitespace change or moving a few functions around is useless, but if you realize that two 100 line functions can actually use the same 110 line function with arguments while also improving readability, this is a big win. Refactors should pass [process replay](#process-replay-tests).
 - Tests/fuzzers. If you can add tests that are non brittle, they are welcome. We have some fuzzers in here too, and there's a plethora of bugs that can be found with them and by improving them. Finding bugs, even writing broken tests (that should pass) with `@unittest.expectedFailure` is great. This is how we make progress.
 - Dead code removal from core `tinygrad/` folder. We don't care about the code in extra, but removing dead code from the core library is great. Less for new people to read and be confused by.
 
@@ -169,3 +171,9 @@ python3 -m pip install -e '.[testing]'  # install extra deps for testing
 python3 test/test_ops.py                # just the ops tests
 python3 -m pytest test/                 # whole test suite
 ```
+
+#### Process replay tests
+
+[Process replay](https://github.com/tinygrad/tinygrad/blob/master/test/external/process_replay/process_replay.py) detects changes in the generated kernels of CI tests by comparing them against tinygrad master. If your PR is a refactor or speedup without any expected behavior change, it should include a green process replay pass to get merged.
+
+You can enable process replay by adding [run_process_replay] to your PR title. [example](https://github.com/tinygrad/tinygrad/pull/4995). Note that you should keep your branch up-to-date with master.

{tinygrad-0.8.0 → tinygrad-0.9.1}/setup.py

@@ -8,19 +8,19 @@ with open(directory / 'README.md', encoding='utf-8') as f:
   long_description = f.read()
 
 setup(name='tinygrad',
-      version='0.8.0',
+      version='0.9.1',
       description='You like pytorch? You like micrograd? You love tinygrad! <3',
       author='George Hotz',
       license='MIT',
       long_description=long_description,
       long_description_content_type='text/markdown',
-      packages = ['tinygrad', 'tinygrad.codegen', 'tinygrad.nn', 'tinygrad.renderer',
-                  'tinygrad.runtime', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.features'],
+      packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.codegen', 'tinygrad.nn', 'tinygrad.renderer', 'tinygrad.engine',
+                  'tinygrad.runtime', 'tinygrad.runtime.driver', 'tinygrad.runtime.graph', 'tinygrad.shape'],
       classifiers=[
         "Programming Language :: Python :: 3",
         "License :: OSI Approved :: MIT License"
       ],
-      install_requires=["numpy", "tqdm", "gpuctypes",
+      install_requires=["numpy",
                         "pyobjc-framework-Metal; platform_system=='Darwin'",
                         "pyobjc-framework-libdispatch; platform_system=='Darwin'"],
       python_requires='>=3.8',
@@ -36,15 +36,17 @@ setup(name='tinygrad',
             "ruff",
             "types-tqdm",
         ],
+        #'mlperf': ["mlperf-logging @ git+https://github.com/mlperf/logging.git@4.0.0-rc2"],
         'testing': [
             "torch",
             "pillow",
             "pytest",
             "pytest-xdist",
-            "onnx==1.15.0",
+            "onnx==1.16.0",
             "onnx2torch",
             "opencv-python",
             "tabulate",
+            "tqdm",
             "safetensors",
             "transformers",
             "sentencepiece",
@@ -52,6 +54,19 @@ setup(name='tinygrad',
             "librosa",
             "networkx",
             "hypothesis",
+            "nibabel",
+            "bottle",
+        ],
+        'docs': [
+            "mkdocs-material",
+            "mkdocstrings[python]",
+            "markdown-callouts",
+            "markdown-exec[ansi]",
+            "black"
+        ],
+        'testing_tf': [
+            "tensorflow==2.15.1",
+            "tensorflow_addons",
         ]
       },
       include_package_data=True)

tinygrad-0.9.1/test/test_arange.py

@@ -0,0 +1,19 @@
+import unittest
+from tinygrad import Tensor, GlobalCounters
+from tinygrad.helpers import Context
+
+class TestArange(unittest.TestCase):
+  def _get_flops(self, N):
+    GlobalCounters.reset()
+    with Context(NOOPT=1):
+      Tensor.arange(N).realize()
+    return GlobalCounters.global_ops
+
+  def test_complexity(self):
+    f1 = self._get_flops(256)
+    f2 = self._get_flops(2560)
+    print(f"{f1=}, {f2=}")
+    assert f2 / f1 < 15, f"bad complexity, flops {f2/f1:.1f}X while inputs 10X"
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad-0.9.1/test/test_assign.py

@@ -0,0 +1,378 @@
+#!/usr/bin/env python
+import unittest
+import numpy as np
+from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
+
+N = 200  # has to be bigger than the cache to fail
+
+class TestAssign(unittest.TestCase):
+  def test_simple_assignment(self):
+    a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
+    b = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
+    a.realize()
+    b.realize()
+    ba1 = a.lazydata.base.realized
+    bb1 = b.lazydata.base.realized
+    a += b
+    a.realize()
+    ba2 = a.lazydata.base.realized
+    assert ba1 == ba2 and ba1 != bb1
+    np.testing.assert_allclose(a.numpy(), (np.arange(N*N)*2).reshape((N,N)))
+
+  def test_assign_zeros_good(self):
+    a = Tensor.zeros(10,10).contiguous()
+    a.assign(Tensor.ones(10,10))
+    b = Tensor.zeros(10,10).contiguous()
+    a.realize()
+    np.testing.assert_allclose(b.numpy(), 0)
+
+  def test_assign_zeros(self):
+    a = Tensor.zeros(10,10).contiguous()
+    b = Tensor.zeros(10,10).contiguous()
+    a.assign(Tensor.ones(10,10))
+    a.realize()
+    np.testing.assert_allclose(b.numpy(), 0)
+
+  def test_assign_add(self):
+    def f(x):
+      x += 1
+      x.realize()
+    x = Tensor([0])
+    f(x)
+    assert x.item() == 1
+
+  def test_assign_add_twice(self):
+    # NOTE: this has two kernels
+    def f(x):
+      x += 1
+      x += 1
+      x.realize()
+    x = Tensor([0])
+    f(x)
+    assert x.item() == 2
+
+  def test_assign_add_double(self):
+    def f(x):
+      x += 1
+      x.realize()
+    x = Tensor([0])
+    f(x)
+    assert (out:=x.item()) == 1, f"expected 1, got {out}"
+    x = Tensor([0])
+    f(x)
+    assert (out:=x.item()) == 1, f"expected 1, got {out}"
+
+  def test_assign_add_jit(self):
+    @TinyJit
+    def f(x):
+      x += 1
+      x.realize()
+    x = Tensor([0])
+    for _ in range(5): f(x)
+    assert x.item() == 5
+
+  def test_assign_add_jit_other(self):
+    @TinyJit
+    def f(x):
+      x += 1
+      x.realize()
+    x = Tensor([0])
+    for _ in range(5): f(x)
+    assert x.item() == 5
+
+    y = Tensor([0])
+    for _ in range(4): f(y)
+    assert y.item() == 4
+
+  def test_assign_other_jit(self):
+    @TinyJit
+    def f(x, a):
+      x.assign(a)
+      x.realize()
+    x = Tensor([0])
+    for i in range(1, 6):
+      f(x, x.full_like(i).contiguous())  # const would be implicitly folded without contiguous
+      assert x.item() == i
+
+  def test_assign_add_other_jit(self):
+    @TinyJit
+    def f(x, a):
+      x += a
+      x.realize()
+    x = Tensor([0])
+    a = 0
+    for i in range(1, 6):
+      a += i
+      f(x, x.full_like(i).contiguous())
+      assert x.item() == a
+
+  def test_assign_changes(self):
+    a = Tensor.ones(4).contiguous().realize()
+    old_a = a
+    a.assign(Tensor.full((4,), 2.).contiguous())
+    # NOTE: old_a is now 2, and this would match the behavior of pytorch
+    new = a + old_a
+    np.testing.assert_allclose(new.numpy(), 4)
+
+  def test_assign_diamond_cycle(self):
+    # NOTE: should *not* raise AssertionError from numpy
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      a = Tensor.ones(4).contiguous().realize()
+      times_a = a*3
+      a.assign(Tensor.full((4,), 2.).contiguous())
+      new = a + (times_a-1)
+      np.testing.assert_allclose(new.numpy(), 4)
+
+  def test_assign_diamond_contiguous_cycle(self):
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      a = Tensor.ones(4).contiguous().realize()
+      times_a = a*3
+      a.assign(Tensor.full((4,), 2.))
+      new = a.contiguous() + times_a-1
+      np.testing.assert_allclose(new.numpy(), 4)
+
+  def test_assign_diamond_possible(self):
+    a = Tensor.ones(4).contiguous().realize()
+    times_a = a*3
+    a.assign(Tensor.full((4,), 2.))
+    new = a + (times_a-1).contiguous()
+    np.testing.assert_allclose(new.numpy(), 4)
+
+  def test_assign_diamond_possible_contiguous(self):
+    a = Tensor.ones(4).contiguous().realize()
+    times_a = a*3
+    a.assign(Tensor.full((4,), 2.).contiguous())
+    new = a + (times_a-1).contiguous()
+    np.testing.assert_allclose(new.numpy(), 4)
+
+  def test_assign_diamond_both_contiguous(self):
+    a = Tensor.ones(4).contiguous().realize()
+    times_a = a*3
+    a.assign(Tensor.full((4,), 2.))
+    new = a.contiguous() + (times_a-1).contiguous()
+    np.testing.assert_allclose(new.numpy(), 4)
+
+  def test_assign_diamond_alt(self):
+    a = Tensor.ones(4).contiguous().realize()
+    a.assign(Tensor.full((4,), 2.).contiguous())
+    times_a = a*3
+    new = a + times_a
+    np.testing.assert_allclose(new.numpy(), 8)
+
+  def test_double_assign(self):
+    a = Tensor.ones(4).contiguous().realize()
+    a += 1
+    a += 1
+    np.testing.assert_allclose(a.numpy(), 3)
+
+  def test_crossover_assign(self):
+    a = Tensor.full((4,), 2).contiguous().realize()
+    b = Tensor.full((4,), 3).contiguous().realize()
+    a += b
+    b += a
+    Tensor.realize(a,b)
+    np.testing.assert_allclose(a.numpy(), 5)
+    np.testing.assert_allclose(b.numpy(), 8)
+
+  def test_assign_double_diamond(self):
+    a = Tensor.full((4,), 2).contiguous().realize()
+    b = Tensor.full((4,), 3).contiguous().realize()
+    a_prev = a*4
+    b_prev = b+3
+    b += a_prev.contiguous()
+    a += b_prev.contiguous()
+    Tensor.realize(a, b)
+    np.testing.assert_equal(b.numpy(), 11)
+    np.testing.assert_equal(a.numpy(), 8)
+
+  def test_assign_double_diamond_reduce(self):
+    a0 = Tensor.full((16, 16), 10).contiguous().realize()
+    a1 = Tensor.full((16, 16), 20).contiguous().realize()
+    b0 = Tensor.full((16, ), 1).contiguous().realize()
+    b1 = Tensor.full((16, ), 2).contiguous().realize()
+
+    r0 = (a0 - b1.contiguous()).sum(1)
+    r1 = (a1 - b0.contiguous()).sum(1)
+    b0.assign(r0 * b0)
+    b1.assign(r1 * b1)
+    Tensor.realize(b0, b1)
+    np.testing.assert_equal(b0.numpy(), 128)
+    np.testing.assert_equal(b1.numpy(), 608)
+
+  def test_crossunder_assign(self):
+    # NOTE: should *not* raise AssertionError from numpy
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      a = Tensor.full((4,), 2).contiguous().realize()
+      b = Tensor.full((4,), 3).contiguous().realize()
+      c = a+9
+      a += b
+      b += c
+      Tensor.realize(a,b)
+      np.testing.assert_allclose(a.numpy(), 2+3)
+      np.testing.assert_allclose(b.numpy(), 3+2+9)
+
+  def test_assign_kv_cache(self):
+    bsz, max_context = 2, 8
+
+    class Attn:
+      @TinyJit
+      def __call__(self, xk:Tensor, start_pos:Variable):
+        seqlen = xk.shape[1]
+        if not hasattr(self, "cache_k"):
+          self.cache_k = Tensor.zeros(bsz, max_context, 1, 1).contiguous()
+        keys = self.cache_k.shrink((None, (0, start_pos), None, None)).cat(xk, dim=1).contiguous() if start_pos > 0 else xk
+        self.cache_k.assign(keys.pad((None,(0,max_context-start_pos-seqlen),None,None)).contiguous()).realize()
+
+    attn = Attn()
+    xk = Tensor.ones(bsz, 3, 1, 1).contiguous()
+    attn(xk, 0)
+    for i in range(3,6):
+      # copied from LLaMA
+      start_pos = Variable("start_pos", 1, max_context).bind(i)
+      xk = Tensor.ones(bsz, 1, 1, 1).contiguous()
+      attn(xk, start_pos)
+
+    out = attn.cache_k.flatten().numpy()
+    np.testing.assert_allclose(out, [1.,1.,1.,1.,1.,1.,0.,0.,1.,1.,1.,1.,1.,1.,0.,0.])
+
+  def test_assign_contiguous(self):
+    b = Tensor.rand(4,4).realize()
+    a = (Tensor.rand(4,4).realize() + 1)
+    kc = GlobalCounters.kernel_count
+    b.assign(a.contiguous()).realize()
+    assert GlobalCounters.kernel_count - kc == 2
+
+  def test_assign_contiguous_permute(self):
+    b = Tensor.rand(4,4).realize()
+    a = (Tensor.rand(4,4).realize() + 1).permute((1,0))
+    kc = GlobalCounters.kernel_count
+    b.assign(a.contiguous()).realize()
+    assert GlobalCounters.kernel_count - kc == 2
+
+  def test_permuted_assignment(self):
+    a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
+    b = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
+    a.realize()
+    b.realize()
+    ba1 = a.lazydata.base.realized
+    bb1 = b.lazydata.base.realized
+    with self.assertRaises((RuntimeError, AssertionError)):
+      a = a.permute(1,0)
+      a += b
+      a.realize()
+      ba2 = a.lazydata.base.realized
+      assert ba1 != ba2 and ba1 != bb1
+      np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
+
+  def test_post_permuted_assignment(self):
+    a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
+    b = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
+    a.realize()
+    b.realize()
+    #GlobalCounters.cache = []
+    ba1 = a.lazydata.base.realized # noqa: F841
+    bb1 = b.lazydata.base.realized # noqa: F841
+    with self.assertRaisesRegex(RuntimeError, "contiguous"):
+      a.assign(a.permute(1,0) + b)   # this should not work!
+      a.realize()
+      ba2 = a.lazydata.base.realized # noqa: F841
+      # NOTE: don't test that it's assigned
+      #assert ba1 == ba2 and ba1 != bb1
+      np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
+
+  def test_simple_assignment_multioutput(self):
+    a = Tensor.randn(32, 32).realize()
+    b = Tensor.full((32, ), 1.).contiguous().realize()
+    c = Tensor.full((32, ), 2.).contiguous().realize()
+    d = Tensor.full((32, ), 3.).contiguous().realize()
+
+    r = a.sum(axis=1)
+    b.assign(r + b)
+    c.assign(r + c)
+    d.assign(r + d)
+
+    kc = GlobalCounters.kernel_count
+    Tensor.realize(b, c, d)
+    assert GlobalCounters.kernel_count - kc == 1
+    np.testing.assert_allclose(b.numpy(), a.sum(1).numpy()+1)
+    np.testing.assert_allclose(c.numpy(), a.sum(1).numpy()+2)
+    np.testing.assert_allclose(d.numpy(), a.sum(1).numpy()+3)
+
+  # NOTE: if the assign target is read/write in a single kernel, it should be contiguous
+
+  def test_permuted_assignment_correct(self):
+    a = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
+    b = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
+    # TODO: scheduler limitation, should NOT raise AssertionError from numpy.
+    with self.assertRaisesRegex(RuntimeError, "contiguous"):
+      a = a.permute(1, 0)
+      new_val = a + b
+      a.assign(new_val)
+      np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))
+
+  def test_permuted_reduceop_child_dual_use(self):
+    a = Tensor.randn(32, 32, 32).realize()
+    b = Tensor.full((32, 32), 1.).contiguous().realize()
+    with self.assertRaisesRegex(RuntimeError, "contiguous"):
+      r = a.sum(axis=1)
+      b.assign(r + b.permute(1, 0))
+      b.realize()
+
+  def test_permuted_reduceop_multioutput_dual_use(self):
+    a = Tensor.randn(32, 32, 32).realize()
+    b = Tensor.full((32, 32), 1.).contiguous().realize()
+    c = Tensor.full((32, 32), 2.).contiguous().realize()
+
+    with self.assertRaisesRegex(RuntimeError, "contiguous"):
+      r = a.sum(axis=1)
+      b_perm = b.permute(1, 0)
+      b.assign(r + b)
+      c.assign(r + b_perm)
+      Tensor.realize(b, c)
+
+  def test_permuted_reduceop_multioutput_dual_use_possible(self):
+    a = Tensor.randn(32, 32, 32, dtype=dtypes.int).realize()
+    b = Tensor.arange(32 * 32).reshape(32, 32).realize()
+    c = Tensor.arange(32 * 32).reshape(32, 32).realize()
+
+    kc = GlobalCounters.kernel_count
+    r = a.sum(axis=1)
+    b_perm = b.permute(1, 0)
+    b.assign(r + b)
+    c.assign(r + b_perm.contiguous())
+    Tensor.realize(b, c)
+    assert GlobalCounters.kernel_count - kc == 2
+    np.testing.assert_equal(b.numpy(), a.numpy().sum(1) + np.arange(32 * 32).reshape(32, 32))
+    np.testing.assert_equal(c.numpy(), a.numpy().sum(1) + np.arange(32 * 32).reshape(32, 32).transpose(1, 0))
+
+  def test_permuted_assignment_masked_view_possible(self):
+    a = Tensor.ones(4, 4).contiguous().realize()
+    b = a.shrink((None, (0, 2))).pad((None, (0, 2)), 2)
+    a.assign(a + b)
+    kc = GlobalCounters.kernel_count
+    a.realize()
+    assert GlobalCounters.kernel_count - kc == 1
+    np.testing.assert_equal(a.numpy(), np.ones((4, 4))+np.pad(np.ones((4, 4))[:, 0:2], ((0, 0), (0, 2)), constant_values=2))
+
+  def test_permuted_assignment_masked_view_not_contiguous(self):
+    a = Tensor.ones(4, 4).contiguous().realize()
+    with self.assertRaisesRegex(RuntimeError, "contiguous"):
+      b = a.shrink((None, (0, 2))).pad((None, (0, 2)), 2).permute(1, 0)
+      a.assign(a + b)
+      a.realize()
+
+  # TODO: is there a way to sneak in a permute such that it returns the wrong answer?
+
+  @unittest.skip("don't use output buffer, and mismatch dtype no longer supported")
+  def test_cast_assignment(self):
+    a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
+    a.realize()
+    oba1 = a.lazydata.base.output_buffer
+    a.assign(a.cast(dtypes.int32).realize())
+    a.realize()
+    oba2 = a.lazydata.base.output_buffer
+    assert oba1 is None and oba2 is None
+    np.testing.assert_allclose(a.numpy(), np.arange(N*N,dtype=np.int32).reshape((N,N)))
+
+if __name__ == "__main__":
+  unittest.main()