github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

unit tests for Device.canonicalize (#3055)

Browse Source
This commit is contained in:
chenyu
2024-01-09 12:47:20 -05:00
committed by GitHub
parent 374f7659a7
commit dab8214103
2 changed files with 20 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
test/test_lazybuffer.py
View File

@@ -45,7 +45,7 @@ class TestLazyBuffer(unittest.TestCase):
z = Tensor([1, np.e]).numpy()
np.testing.assert_allclose(y, z)
def test_device_canonicalize(self):
def test_device_0_is_the_same_device(self):
a = Tensor([1, 2, 3], f"{Device.DEFAULT}")
b = Tensor([1, 2, 3], f"{Device.DEFAULT}:0")
assert a.device == b.device

19
test/unit/test_device.py Normal file
View File

@@ -0,0 +1,19 @@
#!/usr/bin/env python
import unittest
from tinygrad.device import Device
class TestDevice(unittest.TestCase):
def test_canonicalize(self):
assert Device.canonicalize(None) == Device.DEFAULT
assert Device.canonicalize("CPU") == "CPU"
assert Device.canonicalize("cpu") == "CPU"
assert Device.canonicalize("GPU") == "GPU"
assert Device.canonicalize("GPU:0") == "GPU"
assert Device.canonicalize("gpu:0") == "GPU"
assert Device.canonicalize("GPU:1") == "GPU:1"
assert Device.canonicalize("gpu:1") == "GPU:1"
assert Device.canonicalize("GPU:2") == "GPU:2"
assert Device.canonicalize("disk:/dev/shm/test") == "DISK:/dev/shm/test"
if __name__ == "__main__":
unittest.main()