Skip to content

Commit 7d0efaa

Browse files
authored
[Kernels] Add a few redundant tensor attribute to be compatible with torch.Tensor (#7688)
1 parent c2cfe83 commit 7d0efaa

File tree

1 file changed

+8
-0
lines changed
  • python/triton_kernels/triton_kernels

1 file changed

+8
-0
lines changed

python/triton_kernels/triton_kernels/tensor.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,14 @@ def data(self):
149149
t = self.storage
150150
return t.data if isinstance(t, Storage) else t
151151

152+
def dim(self):
153+
return self.ndim
154+
155+
def size(self, i=None):
156+
if i is None:
157+
return self.shape
158+
return self.shape[i]
159+
152160

153161
@dataclass
154162
class Bitmatrix(Tensor):

0 commit comments

Comments
 (0)