geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

gptoss: fix sharded invalids (#17450)

Browse Source
This commit is contained in:
wozeparrot
2026-08-07 08:42:02 -07:00
committed by GitHub
parent b6189db8e9
commit 1827ec57f7
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/moe_routing.py
+2 -1
View File
@@ -5,7 +5,8 @@ BLOCK_ROW = 256
def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
if isinstance(device, tuple):
return Tensor.invalids(*shape, dtype=dtype, device=device[0]).shard(device, axis=0)
per = Tensor.invalids(shape[0]//len(device), *shape[1:], dtype=dtype, device=device)
return Tensor(per.uop.unshard(0), device=device)
return Tensor.invalids(*shape, dtype=dtype, device=device)
def _atomic_add(device:str) -> str: