30 lines
950 B
Python
30 lines
950 B
Python
import unittest
|
|
from tinygrad.helpers import getenv
|
|
import torch
|
|
import tinygrad.frontend.torch
|
|
torch.set_default_device("tiny")
|
|
import numpy as np
|
|
|
|
@unittest.skipIf(getenv("GPUS",1)<=1, "only single GPU")
|
|
class TestTorchBackendMultiGPU(unittest.TestCase):
|
|
def test_transfer(self):
|
|
a = torch.Tensor([[1,2],[3,4]]).to("tiny:0")
|
|
b = torch.Tensor([[3,2],[1,0]]).to("tiny:1")
|
|
self.assertNotEqual(a.device, b.device)
|
|
np.testing.assert_array_equal(a.cpu(), a.to("tiny:1").cpu())
|
|
np.testing.assert_array_equal(b.cpu(), b.to("tiny:1").cpu())
|
|
|
|
def test_basic_ops(self):
|
|
a = torch.Tensor([[1,2],[3,4]]).to("tiny:0")
|
|
b = torch.Tensor([[3,2],[1,0]]).to("tiny:1")
|
|
c1 = a + b.to("tiny:0")
|
|
c2 = b + a.to("tiny:1")
|
|
np.testing.assert_array_equal(c1.cpu(), torch.full((2,2),4).cpu())
|
|
np.testing.assert_array_equal(c1.cpu(), c2.cpu())
|
|
|
|
# TODO: torch.distributed functions
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|