github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 23:48:01 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix vectorized dtype repr [run_process_replay] (#6535)

Browse Source
This commit is contained in:
qazal
2024-09-16 13:42:55 +08:00
committed by GitHub
parent 9b9b83b8b0
commit 607113fcdf
2 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
test/test_uops.py
View File

@@ -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

2
tinygrad/dtype.py
View File

@@ -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