diff --git a/test/test_uops.py b/test/test_uops.py index f520135ecb..f9484a909b 100644 --- a/test/test_uops.py +++ b/test/test_uops.py @@ -393,6 +393,10 @@ class TestUOpStr(unittest.TestCase): uop = UOp(UOps.CONST, dtypes.int, (), arg=Variable("a",1,10)) assert str(eval(str(uop))) == str(uop) + def test_vectorized_str(self): + vec = UOp(UOps.VECTORIZE, dtypes.int.vec(4), tuple(UOp.const(dtypes.int, x) for x in range(4))) + assert str(eval(str(vec))) == str(vec) + class TestIndexingOrdering(unittest.TestCase): # NOTE: these tests skip type_verify since they add dtype to STORE @unittest.expectedFailure diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index ca1df1f006..ca92e6adee 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -13,7 +13,7 @@ class DType: name: str fmt: Optional[str] count: int - def __repr__(self): return f"dtypes.{'_'*(c:=self.count!=1)}{INVERSE_DTYPES_DICT[self.name if not c else self.scalar().name]}{str(self.count)*c}" + def __repr__(self): return f"dtypes.{INVERSE_DTYPES_DICT[self.scalar().name]}"+(f".vec({self.count})" if self.count > 1 else "") def vec(self, sz:int): assert self.count == 1, f"can't vectorize {self} with size {sz}" if sz == 1 or self.name == 'void': return self # void doesn't vectorize, and sz=1 is scalar