github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-08 22:48:25 -05:00
Code Issues Packages Projects Releases Wiki Activity

support DEV= to specify device (#11351)

Browse Source
This commit is contained in:
chenyu
2025-07-23 17:40:55 -04:00
committed by GitHub
parent 76a2ddbd78
commit 5b570196e4
2 changed files with 21 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
test/test_tensor.py
View File

@@ -512,6 +512,24 @@ class TestTinygrad(unittest.TestCase):
subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
shell=True, check=True)
if Device.DEFAULT != "CPU":
# setting multiple devices fail
with self.assertRaises(subprocess.CalledProcessError):
subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
shell=True, check=True)
# setting device via DEV
subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
shell=True, check=True)
subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
shell=True, check=True)
subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
shell=True, check=True)
with self.assertRaises(subprocess.CalledProcessError):
subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
shell=True, check=True)
def test_no_attributeerror_after_apply_uop_exception(self):
try:
Tensor.arange(4).reshape(3,2)