L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests 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 <[email protected]>
Co-authored-by: Awni Hannun <[email protected]>
This commit is contained in:
K Venkat Ramnan
2024-05-31 12:29:01 -07:00
committed by GitHub
co-authored by Venkat Ramnan Kalyanakumar Awni Hannun
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
python/tests/test_array.py
+13
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)