github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

do not cache subgraphs in toposort [pr] (#8403)

Browse Source 
* is it correct?

* style

* set
This commit is contained in:
nimlgen
2024-12-25 20:48:30 +03:00
committed by GitHub
parent 313bdfa43f
commit 393d39da58
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tinygrad/ops.py
View File

@@ -258,15 +258,15 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
@property
def toposort(self) -> dict[UOp, None]:
def _toposort(u:UOp, cache:dict[UOp, dict[UOp, None]]):
if (cret:=cache.get(u)) is not None: return cret
def _toposort(u:UOp, cache:set[UOp]):
if u in cache: return {}
nodes: dict[UOp, None] = {}
# NOTE: this is a lot faster than the comprehension in parents
for parent in u.src: nodes.update(_toposort(parent, cache))
nodes[u] = None
cache[u] = nodes
cache.add(u)
return nodes
return _toposort(self, cache={})
return _toposort(self, cache=set())
@functools.cached_property
def tuplize(self:UOp) -> tuple[int, Any, Optional[DType], tuple]: return (self.op.value, self.arg, self.dtype, tuple(x.tuplize for x in self.src))