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

use next instead of full list in UOp._device [pr] (#11369)

Browse Source 
prevents exponential fan out
This commit is contained in:
chenyu
2025-07-25 10:04:29 -04:00
committed by GitHub
parent 1bb1f1aee8
commit dab07bcad9
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tinygrad/uop/ops.py
View File

@@ -372,7 +372,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
return self.src[0].device[self.arg]
if self.op is Ops.MSTACK: return tuple(cast(str, x.device) for x in self.src)
if self.op in {Ops.COPY, Ops.BUFFER, Ops.ALLREDUCE}: return self.src[1].device
return dsrcs[0]._device if len(dsrcs:=[x for x in self.src if x._device is not None]) != 0 else None
return next((x._device for x in self.src if x._device is not None), None)
@property
def buf_uop(self) -> UOp:
if self.op is Ops.BUFFER: return self