tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 15:36:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

use clone instead of unique_const in allreduce [pr] (#16297)

Browse Source
This commit is contained in:
chenyu
2026-05-20 18:58:47 -04:00
committed by GitHub
parent 9e88b08f93
commit d548f8d0f3
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/schedule/allreduce.py
+1 -2
View File
@@ -55,8 +55,7 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
return UOp.usum(*[c.pad(((s,numel-e),)) for (s,e),c in zip(chunks, copied_chunks)]).reshape(shape)
def create_allreduce_function(buf:UOp, red:UOp, output:UOp|None=None) -> UOp|None:
# BUFFER without unique have unique added later
if output is None: output = UOp.unique_const(Invalid, red.dtype, red.device, red.shape).contiguous()
if output is None: output = UOp.const(red.dtype, Invalid, red.device, red.shape).clone()
to = red.param_like(0)
src = buf.param_like(1)
red = src.allreduce(red.arg, red.src[1])