mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

simpler idxs_to_idx (#3071)

Browse source
This commit is contained in:
chenyu 2024-01-10 00:30:10 -05:00 • committed by GitHub
commit 023f5df0e9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

8
tinygrad/shape/shapetracker.py
View file

@ -49,11 +49,9 @@ def simplify(views:Tuple[View, ...]) -> Tuple[View, ...]:
@functools.lru_cache(maxsize=None)
def idxs_to_idx(shape:Tuple[int, ...], idxs:Tuple[Node, ...]) -> Node:
assert len(idxs) == len(shape), "need an idx for all dimensions"
acc, ret = 1, []
for tidx,d in zip(reversed(idxs), reversed(shape)):
ret.append(tidx * acc)
acc *= d
return Node.sum(ret)
# idxs[-1] * 1 + idxs[-2] * shape[-1] + idxs[-3] * shape[-1] * shape[-2] + ...
accs = itertools.accumulate(reversed(shape[1:]), operator.mul, initial=1)
return Node.sum([idx * acc for idx, acc in zip(reversed(idxs), accs)])
@dataclass(frozen=True)
class ShapeTracker: