added tests

Author: aaishwarymishra

python/mlx/__array_api_info.py

@@ -1,12 +1,9 @@
 class ArrayNamespaceInfo:
-    def __init__(self):
-        pass
-
     def capabilities(self):
         return {
             "boolean indexing": False,
             "data-dependent shapes": False,
-            "max dimensions": None,
+            "max dimensions": 10,
         }
 
     def default_device(self):
@@ -19,19 +16,11 @@ def default_dtypes(self, *, device=None):
 
         if device is not None and not isinstance(device, mx.Device):
             raise TypeError("Expected a mlx Device")
-        device = device if device is not None else self.default_device()
-        if device.type == mx.gpu:
-            return {
-                "real floating": mx.float32,
-                "complex floating": mx.complex64,
-                "integral": mx.int32,
-                "indexing": mx.uint32,
-            }
         return {
-            "real floating": mx.float64,
+            "real floating": mx.float32,
             "complex floating": mx.complex64,
-            "integral": mx.int64,
-            "indexing": mx.uint64,
+            "integral": mx.int32,
+            "indexing": mx.int32,
         }
 
     def devices(self):
@@ -45,7 +34,48 @@ def devices(self):
         return tuple(devices)
 
     def dtypes(self, *, device=None, kind=None):
-        pass
+        import mlx.core as mx
+
+        if device is not None and not isinstance(device, mx.Device):
+            raise TypeError("Expected a mlx Device")
+        device = device if device is not None else self.default_device()
+
+        dtypes = {
+            "bool": mx.bool_,
+            "int8": mx.int8,
+            "int16": mx.int16,
+            "int32": mx.int32,
+            "int64": mx.int64,
+            "uint8": mx.uint8,
+            "uint16": mx.uint16,
+            "uint32": mx.uint32,
+            "uint64": mx.uint64,
+            "float32": mx.float32,
+            "complex64": mx.complex64,
+        }
+        if device.type == mx.cpu:
+            dtypes["float64"] = mx.float64
+        if kind is None:
+            return dtypes
+
+        signed = {"int8", "int16", "int32", "int64"}
+        unsigned = {"uint8", "uint16", "uint32", "uint64"}
+        real = {"float32", "float64"}
+        complex_ = {"complex64"}
+        kinds = {
+            "bool": {"bool"},
+            "signed integer": signed,
+            "unsigned integer": unsigned,
+            "integral": signed | unsigned,
+            "real floating": real,
+            "complex floating": complex_,
+            "numeric": signed | unsigned | real | complex_,
+        }
+        kind = (kind,) if isinstance(kind, str) else kind
+        if not isinstance(kind, tuple) or any(k not in kinds for k in kind):
+            raise ValueError(f"Unsupported dtype kind: {kind!r}")
+        names = {name for k in kind for name in kinds[k]}
+        return {name: dtype for name, dtype in dtypes.items() if name in names}
 
 
 def __array_namespace_info__():

python/tests/test_array.py

@@ -53,6 +53,32 @@ class TestArrayNamespsceInfo(mlx_tests.MLXTestCase):
     def test(self):
         namespace = mx.__array_namespace_info__()
 
+        self.assertEqual(namespace.default_device(), mx.default_device())
+        self.assertEqual(
+            namespace.default_dtypes(),
+            {
+                "real floating": mx.float32,
+                "complex floating": mx.complex64,
+                "integral": mx.int32,
+                "indexing": mx.int32,
+            },
+        )
+        self.assertEqual(
+            namespace.dtypes(device=mx.Device(mx.cpu), kind="real floating"),
+            {"float32": mx.float32, "float64": mx.float64},
+        )
+        if mx.is_available(mx.gpu):
+            self.assertEqual(
+                namespace.dtypes(device=mx.Device(mx.gpu), kind="real floating"),
+                {"float32": mx.float32},
+            )
+        self.assertEqual(
+            namespace.dtypes(kind=("bool", "complex floating")),
+            {"bool": mx.bool_, "complex64": mx.complex64},
+        )
+        with self.assertRaises(ValueError):
+            namespace.dtypes(kind="invalid")
+
 
 class TestDtypes(mlx_tests.MLXTestCase):
     def test_dtypes(self):