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

all2all (#13902)

Browse Source 
* all2all

* um

* fix

* x

* um

* simler

* mypy

* fix

* t

* cmnts
This commit is contained in:
nimlgen
2025-12-31 16:38:32 +03:00
committed by GitHub
parent f7ee644950
commit 25440f0f72
6 changed files with 66 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
test/test_multitensor.py
View File

@@ -256,6 +256,11 @@ class TestMultiTensor(unittest.TestCase):
a,b = _test_allreduce(Tensor.rand(256, 256))
np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
def test_allreduce_all2all(self):
with Context(ALL2ALL=2):
a,b = _test_allreduce(Tensor.rand(256, 256))
np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
def test_copy_jit(self):
@TinyJit
def copy_tensor(x:Tensor): return (x.to(f"{x.device.split(':')[0]}:1") + 1)