forked from tinygrad/tinygrad
explicit error for unbound Variable in program (#16971)
also allow Tensor(UOp, dtype)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
+1
-1
@@ -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)):
|
||||
|
||||
+3
-1
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user