sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

delete dtype_from_uop INDEX exempt (#17674)

Browse Source
This commit is contained in:
chenyu
2026-08-21 21:08:06 -04:00
committed by GitHub
parent af242819d8
commit 7204d46786
2 changed files with 2 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_dtype_weak.py
-7
View File
@@ -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))
tinygrad/uop/ops.py
+2 -5
View File
@@ -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))