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

simplify fancy index with negative Tensor entries (#2749)

Browse Source
This commit is contained in:
chenyu
2023-12-13 14:45:50 -05:00
committed by GitHub
parent b229879613
commit 22feb7330e
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tinygrad/tensor.py
View File

@@ -385,7 +385,8 @@ class Tensor:
for tensor_dim in type_dim[Tensor]:
dims_collapsed_, dims_injected = sum(1 for d in dims_collapsed if tensor_dim >= d), sum(1 for d in type_dim[None] if tensor_dim >= d)
tdim.append(td := tensor_dim - dims_collapsed_ + dims_injected)
idx.append((t := indices[tensor_dim + dims_injected]).sign().__neg__().relu() * ret.shape[td] + t) # normalize the negative tensor indices
# normalize the negative tensor indices
idx.append(((t := indices[tensor_dim + dims_injected]) < 0).where(ret.shape[td], 0) + t)
# compute sum_dim, arange, and idx
max_dim = max(i.ndim for i in idx)