diff --git a/test/backend/test_tensor_variable.py b/test/backend/test_tensor_variable.py index 84311f850d..37d29266e2 100644 --- a/test/backend/test_tensor_variable.py +++ b/test/backend/test_tensor_variable.py @@ -28,16 +28,16 @@ class TestTensorVariable(unittest.TestCase): def test_variable_tensor_dtype_arg(self): vv = Variable("a", 1, 10).bind(2) - # TODO: dtype arg is silently dropped for a symbolic int, should be honored (or rejected) - try: - self.assertEqual(Tensor(vv, dtype=dtypes.float32).dtype, dtypes.float32) - except AssertionError: pass + t = Tensor(vv, dtype=dtypes.float32) + self.assertEqual(t.dtype, dtypes.float32) + self.assertEqual(t.item(), 2.0) def test_unbound_variable_tensor(self): - # TODO: Tensor creation from unbound variable should assert - # with self.assertRaises(AssertionError): Tensor(Variable("u", 1, 10)) - t = Tensor(Variable("u", 1, 10)) - self.assertRaises(KeyError, t.item) # today it builds silently and fails at execution + # an unbound variable schedules fine, but can't execute + with self.assertRaisesRegex(RuntimeError, "unbound"): Tensor(Variable("u", 1, 10)).item() + with self.assertRaisesRegex(RuntimeError, "unbound"): (Tensor(Variable("u", 1, 10)) + 1).item() + # bound variables in an expression are fine + self.assertEqual(Tensor(Variable("u", 1, 10).bind(2) + 1).item(), 3) def test_shrink_beyond_buffer_variable(self): # TODO: shrink by a variable whose vmax exceeds the dim should fail at build, today only CHECK_OOB=1 rejects it diff --git a/test/unit/test_multitensor.py b/test/unit/test_multitensor.py index 922dabe15f..b425931132 100644 --- a/test/unit/test_multitensor.py +++ b/test/unit/test_multitensor.py @@ -78,8 +78,9 @@ class TestMultiTensor(unittest.TestCase): self.assertEqual(Y.device, devices_2) np.testing.assert_equal(X.numpy(), Y.numpy()) - with self.assertRaises(AssertionError): - _ = Tensor(X.uop, dtype=dtypes.float) + Z = Tensor(X.uop, dtype=dtypes.float) + self.assertEqual(Z.dtype, dtypes.float) + np.testing.assert_equal(Z.numpy(), [1.0, 2.0]) def test_sharded_arange(self): sharded_arange = Tensor.arange(1000).clone().shard(devices_2, 0) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 958ba23dd8..6d8865e1bc 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -69,9 +69,9 @@ class Tensor(RandMixin): # create a UOp from the different types of inputs if isinstance(data, UOp): - assert _dtype is None or _dtype==data.dtype or data.dtype==dtypes.index, f"dtype mismatch: {_dtype} vs {data.dtype}" # if data is dtype.index that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of if data.dtype == dtypes.index: data = _index_to_concrete_int(data) + if _dtype is not None: data = data.cast(_dtype) elif data is None: data = UOp.const(_dtype or dtypes.default_float, 0) elif isinstance(data, get_args(ConstType)): diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 4927628891..0a394e71c4 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1154,7 +1154,9 @@ class ProgramInfo: local_size = tuple([sym_infer(sz, var_vals) for sz in self.local_size]) if self.local_size is not None else None return global_size, local_size - def vals(self, var_vals:dict[str, int]): return tuple(var_vals[k.expr] if k.expr not in self.runtimevars else None for k in self.vars) + def vals(self, var_vals:dict[str, int]) -> tuple[int|None, ...]: + try: return tuple(var_vals[k.expr] if k.expr not in self.runtimevars else None for k in self.vars) + except KeyError as e: raise RuntimeError(f"unbound Variable {e} used by {self.function_name}") from None @staticmethod def from_sink(sink:UOp, aux:tuple=()) -> ProgramInfo: