From f82b16a0e9b7c05e73a2420b871bdbd7e88bfabd Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Thu, 18 Sep 2025 10:35:43 -0400
Subject: [PATCH] RANGEIFY test_tensor (#12235)

---
 .github/workflows/test.yml |  5 +++--
 test/test_tensor.py        | 19 +++++++++++++------
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 55d235623c..27437a0134 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -522,11 +522,12 @@ jobs:
       # test_embedding issue with jit
       # test_load_state_dict_sharded_model_dict_same_axis issue with multi
       # test_instancenorm_3d is very slow
+      # test_copy_from_disk issue with DISK
       run: |
         CPU=1 CPU_LLVM=0 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \
-          -k "not test_embedding and not test_load_state_dict_sharded_model_dict_same_axis and not test_instancenorm_3d" \
+          -k "not test_embedding and not test_load_state_dict_sharded_model_dict_same_axis and not test_instancenorm_3d and not test_copy_from_disk" \
           test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_symbolic_ops.py test/test_tensor_variable.py \
-          test/test_outerworld_range.py test/test_sample.py test/test_randomness.py test/test_nn.py test/test_arange.py
+          test/test_outerworld_range.py test/test_sample.py test/test_randomness.py test/test_nn.py test/test_arange.py test/test_tensor.py
     - name: Test const folding
       run: CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_const_folding.py -k "not test_cast_padded and not TestReduceOpsConstFolding and not TestMultiConstFolding"
     - name: Test multitensor
diff --git a/test/test_tensor.py b/test/test_tensor.py
index 3c839c7f8c..defeaac580 100644
--- a/test/test_tensor.py
+++ b/test/test_tensor.py
@@ -4,7 +4,7 @@ import torch
 import unittest, copy, mmap, random, math, array
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _METADATA
-from tinygrad.helpers import getenv, temp, mv_address
+from tinygrad.helpers import getenv, temp, mv_address, RANGEIFY
 from extra.gradcheck import numerical_jacobian, jacobian, gradcheck
 from hypothesis import given, settings, strategies as strat
 from tinygrad.device import is_dtype_supported
@@ -871,11 +871,18 @@ class TestTensorMetadata(unittest.TestCase):
     self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid")
     self.assertTrue(y.grad.uop.metadata[0].backward)
     si = Tensor.schedule(out, x.grad, y.grad)[-1]
-    self.assertEqual(len(si.metadata), 4, f"failed with {si.metadata}")
-    self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "__mul__", "relu"})
-    bw = [m for m in si.metadata if m.backward]
-    self.assertEqual(len(bw), 2)
-    self.assertEqual(bw[0].name, "sigmoid")
+    if not RANGEIFY:
+      self.assertEqual(len(si.metadata), 4, f"failed with {si.metadata}")
+      self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "__mul__", "relu"})
+      bw = [m for m in si.metadata if m.backward]
+      self.assertEqual(len(bw), 2)
+      self.assertEqual(bw[0].name, "sigmoid")
+    else:
+      self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}")
+      self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "relu"})
+      bw = [m for m in si.metadata if m.backward]
+      self.assertEqual(len(bw), 1)
+      self.assertEqual(bw[0].name, "sigmoid")
 
 class TestIdxUpcast(unittest.TestCase):
   def _find_op(self, ast: UOp, op: Ops):