From 823bc17fb567ae5c2da65ef2c6aeaad6ad804dab Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Mon, 26 Jan 2026 11:58:45 -0500
Subject: [PATCH] failed test case for shrink overlap assigns (#14350)

* failed test case for shrink overlap assigns

current logic can create a race resulted in wrong output

* skip for now
---
 test/unit/test_assign.py | 52 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 51 insertions(+), 1 deletion(-)

diff --git a/test/unit/test_assign.py b/test/unit/test_assign.py
index 0d7c0df71c..9cb4bb7b4a 100644
--- a/test/unit/test_assign.py
+++ b/test/unit/test_assign.py
@@ -3,7 +3,7 @@ import unittest
 import numpy as np
 from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import temp, CI, CPU_LVP
+from tinygrad.helpers import temp, CI, CPU_LVP, Context
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -309,6 +309,20 @@ class TestAssign(unittest.TestCase):
     a.assign(a.T+b)
     np.testing.assert_allclose(a.numpy(), new_a)
 
+  def test_post_flipped_assignment(self):
+    a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
+    b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
+    new_a = (a.flip(0)+b).numpy()
+    a.assign(a.flip(0)+b)
+    np.testing.assert_allclose(a.numpy(), new_a)
+
+  def test_post_flipped_assignment_axis1(self):
+    a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
+    b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
+    new_a = (a.flip(1)+b).numpy()
+    a.assign(a.flip(1)+b)
+    np.testing.assert_allclose(a.numpy(), new_a)
+
   def test_post_reshape_assignment_fine(self):
     a = Tensor.arange(N*N).reshape(N, N).contiguous().realize()
     b = Tensor.arange(N*N).reshape(N, N).contiguous().realize()
@@ -401,6 +415,42 @@ class TestAssign(unittest.TestCase):
 
   # TODO: is there a way to sneak in a permute such that it returns the wrong answer?
 
+  # NOTE: overlapping shrink assign tests are WIP, behavior depends on backend/thread ordering
+  @unittest.skip("WIP: not a stable test, relies on undefined behavior")
+  def test_overlapping_shrink_assignment_forward(self):
+    # Forward shift: read index > write index in overlap - works by thread ordering luck
+    N = 100000
+    shift = 1000
+    a = Tensor.arange(N).float().contiguous().realize()
+    expected = np.arange(N, dtype=np.float32)
+    expected[:N-shift] = expected[shift:].copy()
+    with Context(NOOPT=1): a[0:N-shift].assign(a[shift:N]).realize()
+    np.testing.assert_allclose(a.numpy(), expected)
+
+  @unittest.skip("WIP: not a stable test, relies on undefined behavior")
+  @unittest.expectedFailure
+  def test_overlapping_shrink_assignment_reverse(self):
+    # Reverse shift: write index > read index in overlap - race condition!
+    # This fails because find_permutes excludes SHRINK from hazard detection
+    N = 100000
+    shift = 1000
+    a = Tensor.arange(N).float().contiguous().realize()
+    expected = np.arange(N, dtype=np.float32)
+    expected[shift:] = expected[:N-shift].copy()
+    with Context(NOOPT=1): a[shift:N].assign(a[0:N-shift]).realize()
+    np.testing.assert_allclose(a.numpy(), expected)
+
+  @unittest.skip("WIP: not a stable test, relies on undefined behavior")
+  def test_overlapping_shrink_assignment_reverse_with_contiguous(self):
+    # Adding .contiguous() forces a copy, fixing the race
+    N = 100000
+    shift = 1000
+    a = Tensor.arange(N).float().contiguous().realize()
+    expected = np.arange(N, dtype=np.float32)
+    expected[shift:] = expected[:N-shift].copy()
+    with Context(NOOPT=1): a[shift:N].assign(a[0:N-shift].contiguous()).realize()
+    np.testing.assert_allclose(a.numpy(), expected)
+
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
   def test_setitem_half(self):
     a = Tensor.full((8,), 1.0, dtype=dtypes.half).contiguous().realize()