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

Gather bugfix (#1561)

Browse Source
This commit is contained in:
geohotstan
2023-08-16 16:53:14 -07:00
committed by GitHub
parent cb62911f6b
commit a293c18d34
2 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tinygrad/tensor.py
View File

@@ -359,7 +359,7 @@ class Tensor:
def gather(self: Tensor, idx: Tensor, dim: int):
assert idx.ndim == self.ndim, "self.ndim must equal idx.ndim"
assert all(s > i for s,i in zip(self.shape, idx.shape)), "all dim of idx.shape must be smaller than self.shape"
assert all(s >= i for s,i in zip(self.shape, idx.shape)), "all dim of idx.shape must be smaller than self.shape"
if dim < 0: dim += self.ndim
idx = idx.transpose(ax1=dim, ax2=0).unsqueeze(-1)
permarg = list(range(self.ndim))