forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Release Commit @ b6534c0
This commit is contained in:
17
artifacts/package_sources/tinygrad/test/external/external_jit_failure.py
vendored
Normal file
17
artifacts/package_sources/tinygrad/test/external/external_jit_failure.py
vendored
Normal file
@@ -0,0 +1,17 @@
|
||||
from tinygrad import Tensor, TinyJit, Device
|
||||
import numpy as np
|
||||
|
||||
GPUS = 4
|
||||
N = 128
|
||||
ds = tuple([Device.canonicalize(f"{Device.DEFAULT}:{i}") for i in range(GPUS)])
|
||||
t = Tensor.rand(N, N, N).shard(ds, 0)
|
||||
n = t.numpy()
|
||||
|
||||
@TinyJit
|
||||
def allreduce(t:Tensor) -> Tensor:
|
||||
return t.sum(0) #.realize()
|
||||
|
||||
for i in range(10):
|
||||
print(i)
|
||||
tn = allreduce(t).numpy()
|
||||
np.testing.assert_allclose(tn, n.sum(0), atol=1e-4, rtol=1e-4)
|
||||
Reference in New Issue
Block a user