import tempfile, unittest, math from tinygrad import Tensor, dtypes, TinyJit from tinygrad.helpers import Context from tinygrad.dtype import least_upper_float from tinygrad.uop.ops import UOp, Ops, GroupOp, dtype_from_uop, graph_rewrite from tinygrad.uop.weak import pm_commit_weak from tinygrad.uop.symbolic import symbolic_simple from tinygrad.uop.spec import spec_shared, type_verify from tinygrad.engine.jit import JitError from test.helpers import full_rewrite class TestWeakPromotion(unittest.TestCase): def test_rand_requires_concrete(self): with self.assertRaises(ValueError): Tensor.rand(2, dtype=dtypes.weakfloat) with self.assertRaises(ValueError): Tensor.const(1.0).rand_like() with self.assertRaises(ValueError): Tensor.const(1.0).randn_like() def test_reduce_strips_weakness(self): for weak, value, strong in ((dtypes.weakint, 1, dtypes.default_int), (dtypes.weakfloat, 1.0, dtypes.default_float)): t = Tensor.const(value, weak).expand(3) for out in (t.sum(), t.max(), t.prod(), t.cumsum(0), t.cummax(0)[0]): self.assertEqual(out.dtype, strong) self.assertEqual((Tensor.const(1.0).expand(3).sum() + Tensor([1], dtype=dtypes.float16)).dtype, dtypes.float32) def test_materialize_at_default_dtype(self): for weak, value, strong in ((dtypes.weakfloat, 0.5, dtypes.default_float),): t = Tensor.const(value, weak) self.assertEqual(t.dtype, weak) self.assertEqual(t.data().itemsize, strong.itemsize) self.assertEqual(t.numpy().dtype.itemsize, strong.itemsize) # materializing commits at the kind default; contiguous has no layout to fix so it stays weak self.assertEqual((c := t.clone("CPU")).dtype, strong) self.assertEqual(c.item(), value) self.assertEqual(t.contiguous().dtype, weak) def test_assign_into_weak_commits(self): t = Tensor.const(0.5) t.assign(Tensor(1.0, dtype=dtypes.default_float)) self.assertEqual((t.dtype, t.item()), (dtypes.default_float, 1.0)) def test_float_unary_on_weakint_stays_weak(self): self.assertIs(least_upper_float(dtypes.weakint), dtypes.weakfloat) def test_copysign_meets_operands(self): r = Tensor([2], dtype=dtypes.uint8, device="CPU").copysign(Tensor([1], dtype=dtypes.uint32, device="CPU")) self.assertEqual((r.dtype, r.tolist()), (dtypes.uint32, [2])) def test_minimum_reflects_weak_operand(self): r = Tensor(1).minimum(Tensor([2], dtype=dtypes.uint8, device="CPU")) self.assertEqual((r.dtype, r.tolist()), (dtypes.uint8, [1])) for dt in dtypes.uints: r = Tensor([dt.max], dtype=dt, device="CPU").minimum(1) self.assertEqual((r.dtype, r.tolist()), (dt, [1])) self.assertNotIn(Ops.CAST, [u.op for u in r._uop.toposort()]) def test_promote_keeps_shape_args(self): # the shape arg is the same CONST as the value, only the value lifts self.assertEqual((Tensor(5).expand(5) + 1.5).tolist(), [6.5]*5) self.assertEqual((Tensor(2).reshape(1,1).expand(2,2).pad(((0,2),(0,0))) + 0.5).tolist(), [[2.5,2.5],[2.5,2.5],[0.5,0.5],[0.5,0.5]]) x, _ = Tensor(5).reshape(1).pad((1,1))._broadcasted(0.5) self.assertEqual((x._uop.op, x._uop.base.dtype, x._uop.src[1].dtype), (Ops.PAD, dtypes.weakfloat, dtypes.weakint)) def test_broadcasted_keeps_const_weak(self): # a python scalar stays a bare weak CONST through _broadcasted, lifted only to the KIND of the lub x, y = Tensor([1], dtype=dtypes.int8)._broadcasted(3) self.assertEqual((y._uop.base.op, y.dtype, x.dtype), (Ops.CONST, dtypes.weakint, dtypes.int8)) x, y = Tensor([1], dtype=dtypes.int8)._broadcasted(0.5) self.assertEqual((y._uop.base.op, y.dtype, x.dtype), (Ops.CONST, dtypes.weakfloat, dtypes.weakfloat)) x, y = Tensor.const(1).reshape(1)._broadcasted(Tensor([1.0], dtype=dtypes.float32)) self.assertEqual((x._uop.base.op, x._uop.base.val, x.dtype, x.shape, y.dtype), (Ops.CONST, 1, dtypes.weakfloat, (1,), dtypes.float32)) def test_weak_expression_anchors_at_strong_lub(self): # regression test for the HALF bert nan (#17408, reverted in #17409): lub(int32, weakfloat)==weakfloat makes # `loss_mask.sum() + 1e-5` a weakfloat EXPRESSION. Meeting a strong float in a binop must pin it at the lub denom = (Tensor.zeros(912, dtype=dtypes.int32) != Tensor.zeros(912, dtype=dtypes.float32)).sum() + 1e-5 self.assertIs(denom.dtype, dtypes.weakfloat) # the setup: the denominator expression itself is weak x, y = Tensor([2048.0], dtype=dtypes.float32)._broadcasted(denom) self.assertIs(y.dtype, dtypes.float32) recips = [u for u in (x / y)._uop.toposort() if u.op is Ops.RECIPROCAL] self.assertEqual([(u.dtype, u.src[0].dtype) for u in recips], [(dtypes.float32, dtypes.float32)]) with Context(DEFAULT_FLOAT=dtypes.float16): committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_commit_weak) self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32]) def test_div_sub_operand_kept_weak(self): a = Tensor.empty(4, dtype=dtypes.float32) for t in (a / 1, a - 0): self.assertEqual(t.uop.src[1].dtype, dtypes.weakfloat) def test_cast_weak_expression_commits_at_cast_floor(self): # the floor never narrows: a cast BELOW the default does not pull the compute width down with it with Context(DEFAULT_FLOAT=dtypes.float32): narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_commit_weak) self.assertEqual((narrowed.dtype, narrowed.src[0].dtype), (dtypes.float16, dtypes.float32)) def test_cast_weak_expression_value_uses_cast_floor(self): with Context(DEFAULT_FLOAT=dtypes.float16): denom = Tensor.ones(1, dtype=dtypes.int32, device="CPU").sum() * 70000 + 1e-5 out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10) def test_stacked_weak_casts_convert_each_kind(self): # each weak cast is a kind conversion: weakint truncates before weakfloat re-lifts (neither is only a marker) x = Tensor([2.5, -3.7], dtype=dtypes.float32, device="CPU") stacked = x.cast(dtypes.weakint).cast(dtypes.weakfloat) self.assertIs(stacked.dtype, dtypes.weakfloat) self.assertEqual(stacked.tolist(), [2.0, -3.0]) def test_uop_scalar_const_lifts_kind(self): for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint), (dtypes.int32, 1, dtypes.int32, dtypes.weakint), (dtypes.int32, 0.5, dtypes.weakfloat, dtypes.weakfloat), (dtypes.float32, 1, dtypes.float32, dtypes.weakfloat)): out = UOp.variable("x", 0.0 if dtype == dtypes.float32 else 0, 10.0 if dtype == dtypes.float32 else 10, dtype) + value self.assertEqual((out.dtype, out.src[1].op, out.src[1].dtype), (out_dtype, Ops.CONST, const_dtype)) # the kind lift converts the VALUE too (the arg is the only dtype carrier once UOp.const loses its dtype arg), # and a bare weak const UOp is the same spelling as the python scalar: both lift to the same node x = UOp.variable("x", 0.0, 1.0, dtypes.float32) self.assertIsInstance((x + 2).src[1].val, float) self.assertIs(x + UOp.const(2), x + 2) 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)) gate = UOp.const(True) out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_commit_weak) # a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it self.assertEqual((out.src[1], out.src[2]), (UOp.const(5.0, dtypes.bfloat16), gate)) def test_weak_srcs_commit_only_at_a_concrete_lub(self): weak_lub = UOp(Ops.ADD, src=(UOp.const(1), UOp.const(1.0))) self.assertIs(graph_rewrite(weak_lub, pm_commit_weak), weak_lub) concrete = UOp.const(2.0).cast(dtypes.float16) # the weak arm stays bare: its sibling states the width, so the WHERE already derives float16 for it where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(True), concrete, UOp.const(1.0))), pm_commit_weak) self.assertEqual((where.dtype, tuple(x.dtype for x in where.src)), (dtypes.float16, (dtypes.bool, dtypes.float16, dtypes.weakfloat))) def test_derivable_const_rounds_at_the_derived_width(self): # re-rounds a derivable const in place (still bare) so value-keyed folds (x*1 -> x, x*-1 -> NEG) still fire x = UOp.param(0, dtypes.float32, 1).index(UOp.const(0).cast(dtypes.int32)).load() mul = graph_rewrite(x * UOp.const(-0.9999999893980771), symbolic_simple+pm_commit_weak) self.assertIs(mul.src[1], UOp.const(-1.0)) self.assertIs(graph_rewrite(x * UOp.const(1.0000000106), symbolic_simple+pm_commit_weak), x) def test_committed_const_conversion_folds_for_native_format(self): folded = graph_rewrite(UOp.const(16256, dtypes.ushort).cast(dtypes.uint), symbolic_simple) self.assertIs(folded, UOp.const(16256, dtypes.uint)) # fmt-less targets are lowered by renderer rewrites, where collapsing this pair would cycle with float-intermediate insertion. emulated = UOp.const(1.0, dtypes.float).cast(dtypes.bfloat16) self.assertIs(graph_rewrite(emulated, symbolic_simple), emulated) def test_weak_shift_lhs_commits_the_node(self): # a shift derives its lhs's dtype, so committing the lhs restates the root (WGSL's packed store writes `mask << shift_am`) shl = graph_rewrite(UOp.const(0xFFFF) << UOp.variable("x", 0, 16, dtypes.uint), symbolic_simple+pm_commit_weak) self.assertEqual((shl.dtype, shl.src[0]), (dtypes.uint, UOp.const(0xFFFF, dtypes.uint))) @unittest.expectedFailure # TODO: a weak const defers to its consumer (JAX): these dtypes change once python scalars are weak consts def test_changed_rows(self): t_i8, t_f16, t_bf16 = Tensor([1], dtype=dtypes.int8), Tensor([1], dtype=dtypes.float16), Tensor([1], dtype=dtypes.bfloat16) t_bool, t_u16 = Tensor([True]), Tensor([1], dtype=dtypes.uint16) self.assertEqual((t_i8 + 0.5).dtype, dtypes.weakfloat) self.assertEqual(((t_i8 + 0.5) + t_f16).dtype, dtypes.float16) self.assertEqual(((t_i8 + 0.5) + t_bf16).dtype, dtypes.bfloat16) self.assertEqual(((t_bool + 1) + t_i8).dtype, dtypes.int8) self.assertEqual(((t_bool + 1) + t_u16).dtype, dtypes.uint16) self.assertEqual((Tensor(3) + t_i8).dtype, dtypes.int8) # zeros/ones are full with a python fill value, so they are weak too (jnp.zeros pins float32; deliberate divergence) self.assertEqual((Tensor.zeros(3) + t_f16).dtype, dtypes.float16) def test_unchanged_rows(self): t_i8, t_f16, t_f32 = Tensor([1], dtype=dtypes.int8), Tensor([1], dtype=dtypes.float16), Tensor([1], dtype=dtypes.float32) self.assertEqual((t_i8 + 1).dtype, dtypes.int8) self.assertEqual((t_f16 + 0.5).dtype, dtypes.float16) self.assertEqual((t_f32 + t_f16).dtype, dtypes.float32) self.assertEqual(Tensor([2], dtype=dtypes.uint8).pad(((1, 1),), value=1).dtype, dtypes.uint8) def test_concrete_pair_promotes_weak(self): out = Tensor([-1], dtype=dtypes.int64, device="CPU") + Tensor([3], dtype=dtypes.uint64, device="CPU") + Tensor(0.5) self.assertEqual((out.dtype, out.tolist()), (dtypes.weakfloat, [2.5])) def test_dot_defers_weak(self): weak = Tensor([True, False]).where(Tensor(1), 2) self.assertEqual(weak.dot(Tensor([1, 1], dtype=dtypes.int8)).dtype, dtypes.int8) def test_weak_int_binop(self): v = UOp.variable("i", 0, 10, dtypes.weakint) self.assertEqual((v << 1).dtype, dtypes.weakint) self.assertEqual(dtype_from_uop(Ops.SHL, (UOp.const(1, dtypes.int8), UOp.const(1, dtypes.uint32)), None), dtypes.int8) self.assertEqual(UOp.const(1).alu(Ops.SHL, UOp.const(1, dtypes.uint)).dtype, dtypes.weakint) self.assertEqual((v & 3).dtype, dtypes.weakint) with self.assertRaises(RuntimeError): (Tensor.const(1.0) << Tensor.const(1.0)).dtype with self.assertRaises(RuntimeError): UOp.const(1, dtypes.int32).alu(Ops.SHL, UOp.const(1, dtypes.float64)).dtype for op in (Ops.SHL, Ops.SHR): with self.assertRaises(RuntimeError): UOp.const(1, dtypes.float32).alu(op, UOp.const(1, dtypes.int32)).dtype # float bitwise builds, the spec rejects it with Context(SPEC=1): f32, wf = UOp.const(1.0, dtypes.float32), UOp.const(1.0) for bad in (f32.alu(Ops.AND, f32), UOp(Ops.AND, (f32, f32)), UOp(Ops.AND, (wf, wf))): with self.assertRaises(RuntimeError): type_verify([bad], spec_shared) def test_integer_values(self): x = Tensor.full((1,), 1, dtype=dtypes.int64, device="CPU") self.assertEqual((x + 2**40).item(), 2**40 + 1) self.assertEqual((x << 3).item(), 8) self.assertTrue((x < 2**40).item()) def test_float64_precision(self): value = 1.0 + 2**-40 x64 = Tensor.full((1,), 1.0, dtype=dtypes.float64, device="CPU") self.assertEqual((x64 + value).item(), 2.0 + 2**-40) x32 = Tensor.full((1,), 0.0, dtype=dtypes.float32, device="CPU") self.assertEqual((x32 + value).item(), 1.0) def test_weak_transcendentals(self): t_f16 = Tensor([1], dtype=dtypes.float16) for out in (Tensor(2).exp(), Tensor(2).cos(), Tensor(2).sigmoid()): self.assertEqual((out.dtype, (out + t_f16).dtype), (dtypes.weakfloat, dtypes.float16)) def test_null_lowering(self): for t in (Tensor.full((1,), 1, dtype=dtypes.int64, device="NULL") + 2**40, Tensor.full((1,), 1.0, dtype=dtypes.float64, device="NULL") + (1.0 + 2**-40)): t.realize() self.assertNotIn(t.uop.buffer.dtype, dtypes.weaks) def test_computed_float_index_lowers(self): # a half-pixel nearest index resolves its float-scaled range before the gather idx = (Tensor.arange(8) + 0.5) / 4 - 0.5 idx = (idx.clip(0, 1) - 0.5).ceil().int() out = Tensor([0, 1], device="NULL")[idx].contiguous().realize() self.assertNotIn(out.uop.buffer.dtype, dtypes.weaks) class TestWeakBounds(unittest.TestCase): def test_bounds_survive_movement(self): moved = Tensor(5).reshape(1).expand(2).pad((1, 1)).detach().contiguous_backward() self.assertEqual((moved.uop.vmin, moved.uop.vmax, moved.uop.bufferize().vmax), (0, 5, 5)) self.assertEqual(moved.numpy().dtype, Tensor(5).numpy().dtype) # a moved weak int reads at the same dtype as the bare one def test_wide_src_keeps_its_width(self): # the node's result fits int32, its variable does not: the shift runs at long, only the result narrows v = UOp.variable("v", 0, 2**40).bind(2**35+7) for t in (Tensor(v) // 2**31, (Tensor(v) - 1) // 2**31, Tensor(v).reshape(1) // 2**31): self.assertEqual(t.item(), 16) def test_padded_weak_const_keeps_its_zeros(self): self.assertEqual(Tensor(1).expand(1).cat(Tensor(2).expand(2), Tensor(3).expand(3)).tolist(), [1, 2, 2, 3, 3, 3]) self.assertEqual((Tensor(5).reshape(1).pad((1, 1)) == 5).tolist(), [False, True, False]) self.assertEqual((Tensor(5).reshape(1,1).expand(1,2).pad(((0,2),(0,0))) + Tensor([[1],[2],[3]])).tolist(), [[6,6],[2,2],[3,3]]) class TestWeakStorageBoundary(unittest.TestCase): # weak has no storage: a weak assignment source casts when it defers to the destination, everything else raises def test_weak_source(self): w05 = Tensor.const(0.5).reshape(1) dst = Tensor.zeros(2, dtype=dtypes.int8, device="CPU").contiguous().realize() with self.assertRaises(RuntimeError): dst.assign(w05.expand(2)) # weakfloat into int does not defer with self.assertRaises(RuntimeError): dst[0:1] = w05 fdst = Tensor.zeros(2, dtype=dtypes.float32, device="CPU").contiguous().realize() fdst[0:1] = w05 # weakfloat defers to float self.assertEqual(fdst.tolist(), [0.5, 0.0]) with tempfile.TemporaryDirectory() as td: # the DISK path checks the same ddst = Tensor.empty(2, dtype=dtypes.int32, device=f"DISK:{td}/t") with self.assertRaises(RuntimeError): ddst.assign(w05.expand(2)) def test_weak_commits_by_bounds(self): big = Tensor(2**40) edges = (big.clone(), big.sum(), big.reshape(1).max(), big.reshape(1).mean(), Tensor.stack(big, Tensor(1)).sum() - 1, Tensor([2**40]), big.full_like(2**40)) for t in edges: self.assertEqual(t.item(), 2**40) self.assertEqual(Tensor(UOp.variable("b", 0, 2**40).bind(2**35+3)).clone().item(), 2**35+3) self.assertEqual(Tensor([10, 20, 30])[[2**32+1]].tolist(), [0]) # a wide list index is out of range, not wrapped with Context(DEFAULT_INT=dtypes.int64): self.assertEqual(Tensor(2).clone().dtype, dtypes.int64) def test_literal_beyond_any_int_raises(self): for make in (lambda: Tensor(2**64).item(), lambda: Tensor([2**64]), lambda: Tensor.full((2,), -2**63-1)): with self.assertRaises(OverflowError): make() def test_weak_sentinels_commit_first(self): # max_pool2d, scatter_reduce and cummax pad with the dtype's min/max, which a weak dtype does not have self.assertEqual(Tensor(-5).expand(1, 1, 2, 2).max_pool2d(2, padding=1).dtype, Tensor(-5).clone().dtype) self.assertEqual(Tensor(-5).expand(2).scatter_reduce(0, Tensor([0]), Tensor(-5).expand(1), "amax", include_self=False).tolist(), [-5, -5]) self.assertEqual(Tensor(2**40).expand(3).cummax(0)[0].tolist(), [2**40]*3) def test_weak_has_no_storage(self): import numpy as np with self.assertRaises(RuntimeError): Tensor(np.ones(2, dtype=np.float32), dtype=dtypes.weakfloat) with self.assertRaises(RuntimeError): Tensor(bytes(8), dtype=dtypes.weakfloat) class TestWeakMaterializationEntries(unittest.TestCase): # everything that creates storage from a weak value raises def test_reads_commit_storage_raises(self): for weak, value, strong in ((dtypes.weakfloat, 0.5, dtypes.default_float),): def weak_val(): return Tensor([True], device="CPU").where(Tensor.const(value, weak), Tensor.const(value, weak)) self.assertEqual(weak_val().dtype, weak) self.assertEqual(weak_val().to("CPU").dtype, weak) self.assertEqual(weak_val().data().format, strong.fmt) self.assertEqual(weak_val().numpy().dtype.itemsize, strong.itemsize) self.assertEqual(weak_val().tolist(), [value]) self.assertEqual(weak_val().cast(strong).realize().uop.buffer.dtype, strong) self.assertEqual(weak_val().contiguous().dtype, weak) # no layout to fix, stays weak self.assertEqual(weak_val().realize().dtype, weak) # no width to store, stays weak self.assertEqual(weak_val().clone().dtype, strong) # storage commits at the default for entry in (lambda t: t.to("CPU:1").realize(), lambda t: t.as_param(0)): with self.assertRaises(RuntimeError): entry(weak_val()) def test_weak_is_virtual(self): # NOTE: int64 lub uint64 is weakfloat, so this is device-ful weak from promotion, never from a cast to weak devful = Tensor([1], dtype=dtypes.int64, device="CPU") + Tensor([1], dtype=dtypes.uint64, device="CPU") for t in (Tensor.const(0.5), devful): self.assertTrue(t.uop.is_virtual) # realize is a no-op, so a weak input can never become the real buffer TinyJit needs with self.assertRaises(JitError): TinyJit(lambda x: (x+1).realize())(t) # callify must not silently commit a weak CONTIGUOUS to storage c = devful.alu(Ops.CONTIGUOUS) c.callify() self.assertIs(c.dtype, dtypes.weakfloat) def test_empty_reads_commit(self): for weak, strong in ((dtypes.weakfloat, dtypes.default_float),): empty = Tensor.const(0, weak).reshape(1).shrink(((0, 0),)) self.assertEqual(empty.data().format, strong.fmt) self.assertEqual(empty.numpy().dtype.itemsize, strong.itemsize) self.assertEqual(empty.tolist(), []) class TestSignedUint64Weakfloat(unittest.TestCase): # int64 and uint64 have no common integer supertype (JAX JEP), so the join defers to weakfloat instead of wrapping def test_no_wrap(self): r = Tensor([-1], dtype=dtypes.int64, device="CPU") + Tensor([1], dtype=dtypes.uint64, device="CPU") self.assertEqual((r.dtype, r.item()), (dtypes.weakfloat, 0.0)) def test_weakfloat_lowers(self): i64, u64 = Tensor([-1], dtype=dtypes.int64, device="CPU"), Tensor([3], dtype=dtypes.uint64, device="CPU") r = i64 + u64 + Tensor([2], dtype=dtypes.float16, device="CPU") # a concrete consumer takes the join self.assertEqual((r.dtype, r.cast(dtypes.float32).item()), (dtypes.half, 4.0)) self.assertEqual((i64 < u64).item(), True) # comparison meets at float self.assertAlmostEqual((i64 + u64).sin().item(), math.sin(2), places=5) # Unary lowers before transcendental class TestNoRedundantWide(unittest.TestCase): def wide_alu(self, t:Tensor) -> int: return sum(sum(1 for u in full_rewrite(call.src[0]).toposort() if u.op in GroupOp.ALU and u.dtype in {dtypes.long, dtypes.ulong}) for call in t.schedule_linear().src if call.src[0].op is Ops.SINK) def test_unbounded_long_stays_long(self): self.assertGreater(self.wide_alu(Tensor.empty(16, dtype=dtypes.long)*3 + 1), 0) def test_fancy_index_has_no_wide_alu(self): j, o = Tensor([0, 1, 2]).reshape(3, 1), Tensor([0, 1]).reshape(1, 2) self.assertEqual(self.wide_alu(Tensor.empty(8, 9, 10, 11, 12)[1, j, 2, o, 2]), 0) if __name__ == "__main__": unittest.main()