From 7204d467866cd338be5018cf0c3444fb7802b05d Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 21 Aug 2026 21:08:06 -0400 Subject: [PATCH] delete dtype_from_uop INDEX exempt (#17674) --- test/unit/test_dtype_weak.py | 7 ------- tinygrad/uop/ops.py | 7 ++----- 2 files changed, 2 insertions(+), 12 deletions(-) diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index 59cc8f40f3..7851310666 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -114,13 +114,6 @@ class TestWeakPromotion(unittest.TestCase): self.assertIsInstance((x + 2).src[1].val, float) self.assertIs(x + UOp.const(2), x + 2) - def test_index_dtype_ignores_weakness(self): - with Context(SPEC=2): - idx = UOp.const(0).cast(dtypes.int32) - weak = UOp.const(1.0).expand((1,)) - self.assertEqual(UOp(Ops.INDEX, dtypes.float32, (weak, idx)).dtype, dtypes.float32) - with self.assertRaisesRegex(RuntimeError, "bad dtype"): UOp(Ops.INDEX, dtypes.int32, (weak, idx)) - def test_store_weak_value_uses_destination_dtype(self): with Context(DEFAULT_FLOAT=dtypes.float16): dst = UOp.param(0, dtypes.bfloat16, (1,)).index(UOp.const(0).cast(dtypes.int32)) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 75a8e26a3e..9102f2af5b 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, replace from enum import Enum, auto from tinygrad.uop import Ops, GroupOp from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, Invalid, AddrSpace, strong_dtype -from tinygrad.dtype import PyConst, InvalidType, weak_dtype, bitcast +from tinygrad.dtype import PyConst, InvalidType, bitcast from tinygrad.device import Buffer, MultiBuffer, canonicalize_device, TinyELF from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, floordiv, floormod, diskcache_put, to_function_name, cpu_profile, TracingKey @@ -190,10 +190,7 @@ class UOpMetaClass(type): if dtype is None: dtype = dtype_from_uop(op, src, arg) or dtypes.void # CONST derives its dtype by value only when the constructor omits one # TODO: delete this once the dtype field is removed, for now it just re-implements spec.py - # an INDEX presents its access dtype, which a still-weak source matches up to weakness - if SPEC == 2 and op is not Ops.CONST and \ - (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype and \ - not (op is Ops.INDEX and weak_dtype(expected_dtype) == weak_dtype(dtype)): + if SPEC == 2 and op is not Ops.CONST and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype: raise RuntimeError(f"bad dtype {dtype}, expected {expected_dtype} on {op}") if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret UOpMetaClass.ucache[key] = weakref.ref(created:=super().__call__(*key))