From 6fa0b2b19e88f92b94c0c1d49e78eee3bf907fbd Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 16 Jul 2026 16:12:33 -0400 Subject: [PATCH] materialize weak dtype casts to default (#17051) in clone and _buffer --- test/unit/test_dtype_weak.py | 11 +++++++++++ tinygrad/dtype.py | 2 ++ tinygrad/tensor.py | 11 ++++++----- tinygrad/uop/ops.py | 6 +++--- 4 files changed, 22 insertions(+), 8 deletions(-) diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index 6f43fd1a5a..bceae17e12 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -22,6 +22,17 @@ class TestWeakPromotion(unittest.TestCase): for fn in (lambda: t.bitcast(dtypes.int32), lambda: Tensor.const(dtypes.int32, 2).bitcast(dtypes.weakint), t.element_size, t.nbytes): with self.assertRaises(RuntimeError): fn() + def test_materialize_at_default_dtype(self): + for weak, value, strong in ((dtypes.weakint, 3, dtypes.default_int), (dtypes.weakfloat, 0.5, dtypes.default_float)): + t = Tensor.const(weak, value) + self.assertEqual(t.dtype, weak) + self.assertEqual(t.data().itemsize, strong.itemsize) + self.assertEqual(t.numpy().dtype.itemsize, strong.itemsize) + realized = t.clone("CPU").realize() + self.assertEqual((realized.dtype, realized.uop.buffer.dtype), (strong, strong)) + with patch.object(dtypes, "default_int", dtypes.int64): + self.assertEqual(Tensor.const(dtypes.weakint, 3).numpy().dtype.itemsize, dtypes.int64.itemsize) + def test_uop_scalar_const_unchanged(self): for dtype, value in ((dtypes.index, 1), (dtypes.int32, 1), (dtypes.float32, 0.5)): out = UOp.variable("x", 0.0 if dtype == dtypes.float32 else 0, 10.0 if dtype == dtypes.float32 else 10, dtype) + value diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index 7d6720b6a4..258cabb1e6 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -163,6 +163,8 @@ if (env_default_float := getenv("DEFAULT_FLOAT", "")): DTypeLike = str|DType def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType) else getattr(dtypes, dtype.lower()) +def strong_dtype(dtype:DType) -> DType: + return dtypes.default_int if dtype == dtypes.weakint else dtypes.default_float if dtype == dtypes.weakfloat else dtype # https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html # we don't support complex type diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 5da95e585a..2ad2f5137e 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -3,7 +3,7 @@ from __future__ import annotations import time, functools, sys, inspect, pathlib, hashlib, weakref from typing import Any, Callable, cast, get_args, ParamSpec, TypeVar, Generic, TYPE_CHECKING if TYPE_CHECKING: import numpy -from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, to_dtype, _from_np_dtype, _to_np_dtype, PyConst +from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, to_dtype, strong_dtype, _from_np_dtype, _to_np_dtype, PyConst from tinygrad.helpers import all_int, getenv, fully_flatten, fetch, Metadata, TRACEMETA, is_numpy_ndarray, TracingKey from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, _broadcast_shape @@ -238,7 +238,7 @@ class Tensor(RandMixin): if capturing and not getenv("UNSAFE_ALLOW_JIT_BUFFER"): from tinygrad.engine.jit import JitError raise JitError("cannot access tensor data during JIT capture, the value will be baked in") - x = self.cast(self.dtype).contiguous() + x = self.cast(strong_dtype(self.dtype)).contiguous() if self.uop.device is None or isinstance(self.device, tuple): x = x.clone("CPU") return cast(Buffer, x.realize().uop.buffer).ensure_allocated() @@ -255,10 +255,11 @@ class Tensor(RandMixin): """ if 0 in self.shape: return memoryview(bytearray(0)).cast(self.dtype.fmt) # type: ignore[arg-type,return-value] assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}" - fmt = self.dtype.fmt - assert fmt is not None, f"no fmt dtype for {self.dtype}" + buf = self._buffer() + fmt = buf.dtype.fmt + assert fmt is not None, f"no fmt dtype for {buf.dtype}" assert fmt != "e" or sys.version_info >= (3, 12) - return self._data().cast(fmt, self.shape) # type: ignore[arg-type,return-value] + return buf.as_memoryview().cast(fmt, self.shape) # type: ignore[arg-type,return-value] # NOTE: list[Any] because return type is recursive (list[list[...]] for higher dimensions) def tolist(self) -> PyConst|list[Any]: diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index bae6014cf2..524d7665e9 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -4,7 +4,7 @@ import sys, time, functools, itertools, math, operator, hashlib, os, types, pick 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 +from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, strong_dtype, Invalid, AddrSpace from tinygrad.dtype import ConstFloat, PyConst, InvalidType, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar from tinygrad.device import Buffer, MultiBuffer, canonicalize_device from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA @@ -788,9 +788,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return ret if ret.device == device else ret.copy_to_device(device) def clone(self, device=None) -> UOp: device = device or self.device - ret = self.empty_like(device=device) + ret = self.empty_like(dtype=strong_dtype(self.dtype), device=device) src = self if self.device is None or self.device == device else self.copy_to_device(device) - return ret.after(ret.store(src)) + return ret.after(ret.store(src.cast(ret.dtype))) @recursive_property def device(self) -> str|tuple[str, ...]|None: if self.op is Ops.PARAM: return self.arg.device