zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-03 14:24:44 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: Added dlpack device (#1165)

Browse Source 
* feat: Added dlpack device

* feat: Added device_id to dlpack device

* feat: Added device_id to dlpack device

* doc: updated conversion docs

* doc: updated numpy.rst dlpack information

* doc: updated numpy.rst dlpack information

* Update docs/src/usage/numpy.rst

* Update docs/src/usage/numpy.rst

---------

Co-authored-by: Venkat Ramnan Kalyanakumar <venkatramnankalyanakumar@Venkats-MacBook-Air.local>
Co-authored-by: Awni Hannun <awni.hannun@gmail.com>
This commit is contained in:
K Venkat Ramnan
2024-05-31 12:29:01 -07:00
committed by GitHub
parent fd1c08137b
commit ab977109db
3 changed files with 31 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/tests/test_array.py
View File

@@ -161,6 +161,19 @@ class TestInequality(mlx_tests.MLXTestCase):
self.assertTrue(a != b)
self.assertTrue(a != c)
def test_dlx_device_type(self):
a = mx.array([1, 2, 3])
device_type, device_id = a.__dlpack_device__()
self.assertIn(device_type, [1, 8])
self.assertEqual(device_id, 0)
if device_type == 8:
# Additional check if Metal is supposed to be available
self.assertTrue(mx.metal.is_available())
elif device_type == 1:
# Additional check if CPU is the fallback
self.assertFalse(mx.metal.is_available())
def test_tuple_not_equals_array(self):
a = mx.array([1, 2, 3])
b = (1, 2, 3)