From 6cb74bb630b3c824b6a8ceeddda7d305dc30a249 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Thu, 23 Jan 2025 01:28:07 -0500
Subject: [PATCH] fix using clone with shrink [pr] (#8724)

* fix using clone with shrink [pr]

* remove extra arg, add test_clone_with_shrink_realized
---
 test/test_tensor.py | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/test/test_tensor.py b/test/test_tensor.py
index 0e9c20d022..d73e55e36a 100644
--- a/test/test_tensor.py
+++ b/test/test_tensor.py
@@ -652,19 +652,26 @@ class TestZeroShapeTensor(unittest.TestCase):
 
   def test_clone(self):
     a = Tensor.rand(16, 16).realize()
-    self.assertIsNot(a.lazydata, a.clone().lazydata)
-    np.testing.assert_allclose(a.numpy(), a.clone().numpy())
+    b = a.clone()
+    np.testing.assert_allclose(a.numpy(), b.numpy())
+    self.assertIsNot(a.lazydata.base.buffer, b.lazydata.base.buffer)
 
     a = Tensor.rand(16, 16).mul(5.0).add(5.0)
-    self.assertIsNot(a.lazydata, a.clone().lazydata)
-    np.testing.assert_allclose(a.numpy(), a.clone().numpy())
+    b = a.clone()
+    np.testing.assert_allclose(a.numpy(), b.numpy())
+    self.assertIsNot(a.lazydata.base.buffer, b.lazydata.base.buffer)
 
   def test_clone_with_shrink(self):
-    a = Tensor.empty(16, 16)
-    self.assertIsNot(a.lazydata, a.clone().lazydata)
+    a = Tensor.rand(16, 16)
+    b = a.shrink(((2, 10), None)).clone()
+    b.realize()
+    self.assertIsNot(a.lazydata.base.buffer, b.lazydata.base.buffer)
 
-    b = a.shrink(((2, 10), None))
-    self.assertIsNot(b.lazydata, b.clone().lazydata)
+  def test_clone_with_shrink_realized(self):
+    a = Tensor.rand(16, 16).realize()
+    b = a.shrink(((2, 10), None)).clone()
+    b.realize()
+    self.assertIsNot(a.lazydata.base.buffer, b.lazydata.base.buffer)
 
   def test_clone_with_grad(self):
     a = Tensor.rand(16, 16, requires_grad=True)