github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-02-08 05:35:11 -05:00
Code Issues Packages Projects Releases Wiki Activity

add dtype.ptr() [pr] (#6839)

Browse Source
This commit is contained in:
qazal
2024-10-02 15:03:05 +08:00
committed by GitHub
parent be12409b51
commit 29363fb85e
3 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tinygrad/dtype.py
View File

@@ -19,6 +19,7 @@ class DType:
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
return DType(self.priority, self.itemsize*sz, f"{INVERSE_DTYPES_DICT[self.name]}{sz}", None, sz)
def ptr(self) -> Union[PtrDType, ImageDType]: return PtrDType(self)
def scalar(self) -> DType: return DTYPES_DICT[self.name[:-len(str(self.count))]] if self.count > 1 else self
# dependent typing?
@@ -29,6 +30,7 @@ class ImageDType(DType):
local: bool = False # images are never local
def scalar(self) -> DType: return self.base
def vec(self, sz:int): return self.base.vec(sz)
def ptr(self) -> Union[PtrDType, ImageDType]: return self
def __repr__(self): return f"dtypes.{self.name}({self.shape})"
# @dataclass(frozen=True, init=False, repr=False, eq=False)