diff --git a/test/null/test_tensor.py b/test/null/test_tensor.py index 32ccc255e3..560ce3f5e9 100644 --- a/test/null/test_tensor.py +++ b/test/null/test_tensor.py @@ -171,6 +171,11 @@ class TestTensorConstLike(unittest.TestCase): t = Tensor.ones(8, 4).shard(("NULL:0", "NULL:1"), axis=0) with self.assertRaises(RuntimeError): t.full_like(5, device="NULL") +class TestTensorShape(unittest.TestCase): + def test_float_shape_raises(self): + for dim in (2.0, 2.5): + with self.subTest(dim=dim), self.assertRaisesRegex(RuntimeError, "shape must be int"): Tensor.ones(dim) + class TestTensorDevice(unittest.TestCase): def test_create_from_single_device_tuple(self): (Tensor([1.0], device=(Device.DEFAULT,)) + Tensor([2.0])).realize() diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 49a7429d0a..9a30dbe312 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -94,11 +94,10 @@ def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str: return ret def shape_to_shape_arg(arg:tuple[sint, ...]) -> UOp: - for x in arg: - if isinstance(x, UOp) and not dtypes.is_int(x.dtype): raise RuntimeError(f"shape must be int, got {x.dtype} in {arg}") - if len(arg) == 0: return UOp(Ops.STACK) - elif len(arg) == 1: return UOp.const(arg[0], dtypes.weakint) - else: return UOp(Ops.STACK, src=tuple(UOp.const(x) if isinstance(x, int) else x for x in arg)) + src = tuple(x if isinstance(x, UOp) else UOp.const(x) for x in arg) + for x in src: + if not dtypes.is_int(x.dtype): raise RuntimeError(f"shape must be int, got {x.dtype} in {arg}") + return src[0] if len(src) == 1 else UOp(Ops.STACK, src=src) def consumer_map_from_toposort(lst:Iterable[UOp]): ret: dict[UOp, dict[UOp, None]] = {}