mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

add a test case that triggers CALL passthrough_multi (#14887)

Browse source
This commit is contained in:
chenyu 2026-02-19 10:45:40 -05:00 • committed by GitHub
commit 06ef8a26b7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 8 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

8
test/unit/test_call.py
View file

@ -92,5 +92,13 @@ class TestCall(unittest.TestCase):
np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)
def test_call_plus_sharded(self):
devs = ("CPU:0", "CPU:1")
a = Tensor.ones(10, 10).shard(devs, axis=0)
b = Tensor.ones(10, 10).shard(devs, axis=0)
Tensor.realize(a, b)
c = Tensor.call(a, b, fxn=a.as_param(0) + b.as_param(1))
np.testing.assert_equal(c.numpy(), 2 * np.ones((10, 10)))
if __name__ == '__main__':
unittest.main()