materialize weak dtype casts to default (#17051)

in clone and _buffer

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

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

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]:

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