Add support for MLX arrays

CONTRIBUTING.md

@@ -30,7 +30,7 @@ Next verify the tests all pass:
 ```bash
 pip install -r test/requirements.txt
 pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
-pytest
+pytest ./test
 ```
 
 Then push your changes back to your fork of the repository:

README.md

@@ -2,7 +2,7 @@
 
 Type annotations **and runtime type-checking** for:
 
-1. shape and dtype of [JAX](https://github.com/google/jax) arrays; *(Now also supports PyTorch, NumPy, and TensorFlow!)*
+1. shape and dtype of [JAX](https://github.com/google/jax), [NumPy](https://github.com/numpy/numpy), [MLX](https://github.com/ml-explore/mlx), [PyTorch](https://github.com/pytorch/pytorch), [Tensorflow](https://github.com/tensorflow/tensorflow) arrays/tensors.
 2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).
 
 
@@ -11,6 +11,7 @@ Type annotations **and runtime type-checking** for:
 from jaxtyping import Array, Float, PyTree
 
 # Accepts floating-point 2D arrays with matching axes
+# Feel free to replace 'Array' by torch.Tensor, np.ndarray, tf.Tensor or mx.array.
 def matrix_multiply(x: Float[Array, "dim1 dim2"],
                     y: Float[Array, "dim2 dim3"]
                   ) -> Float[Array, "dim1 dim3"]:

docs/api/array.md

@@ -90,8 +90,9 @@ jaxtyping.Array / jax.Array / jax.numpy.ndarray  # these are all aliases of one
 np.ndarray
 torch.Tensor
 tf.Tensor
+mx.array
 ```
-That is -- despite the now-historical name! -- jaxtyping also supports NumPy + PyTorch + TensorFlow.
+That is -- despite the now-historical name! -- jaxtyping also supports NumPy + PyTorch + TensorFlow + MLX.
 
 Some other types are also supported here:
 

docs/index.md

@@ -2,10 +2,9 @@
 
 jaxtyping is a library providing type annotations **and runtime type-checking** for:
 
-1. shape and dtype of [JAX](https://github.com/google/jax) arrays;
+1. shape and dtype of [JAX](https://github.com/google/jax), [NumPy](https://github.com/numpy/numpy), [MLX](https://github.com/ml-explore/mlx), [PyTorch](https://github.com/pytorch/pytorch), [Tensorflow](https://github.com/tensorflow/tensorflow) arrays/tensors.
 2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).
 
- *(Now also supports PyTorch, NumPy, and TensorFlow!)*
 
 ## Installation
 
@@ -15,7 +14,7 @@ pip install jaxtyping
 
 Requires Python 3.10+.
 
-JAX is an optional dependency, required for a few JAX-specific types. If JAX is not installed then these will not be available, but you may still use jaxtyping to provide shape/dtype annotations for PyTorch/NumPy/TensorFlow/etc.
+JAX is an optional dependency, required for a few JAX-specific types. If JAX is not installed then these will not be available, but you may still use jaxtyping to provide shape/dtype annotations for PyTorch/NumPy/TensorFlow/MLX/etc.
 
 The annotations provided by jaxtyping are compatible with runtime type-checking packages, so it is common to also install one of these. The two most popular are [typeguard](https://github.com/agronholm/typeguard) (which exhaustively checks every argument) and [beartype](https://github.com/beartype/beartype) (which checks random pieces of arguments).
 

jaxtyping/_array_types.py

@@ -617,15 +617,14 @@ def _make_array(x, dim_str, dtype):
 
     if type(out) is tuple:
         array_type, name, dtypes, dims, index_variadic, dim_str = out
+        # Nanobind classes can't be a base type
+        can_subclass = array_type is Any or "nanobind." in str(type(array_type))
         metaclass = (
-            _make_metaclass(type)
-            if array_type is Any
-            else _make_metaclass(type(array_type))
+            _make_metaclass(type) if can_subclass else _make_metaclass(type(array_type))
         )
-
         out = metaclass(
             name,
-            (AbstractArray,) if array_type is Any else (array_type, AbstractArray),
+            (AbstractArray,) if can_subclass else (array_type, AbstractArray),
             dict(
                 dtype=dtype,
                 array_type=array_type,

test/requirements.txt

@@ -8,3 +8,4 @@ pytest
 pytest-asyncio
 tensorflow
 typeguard<3
+mlx

test/test_array.py

@@ -32,6 +32,11 @@
 except ImportError:
     torch = None
 
+try:
+    import mlx.core as mx
+except ImportError:
+    mx = None
+
 from jaxtyping import (
     AbstractArray,
     AbstractDtype,
@@ -605,6 +610,12 @@ def test_subclass():
     if torch is not None:
         assert issubclass(Float[torch.Tensor, ""], torch.Tensor)
 
+    if mx is not None:
+        # Nanobind classes can't be a base type.
+        # Maybe this can be fixed in the future, but it does not
+        # impact type checking.
+        assert not issubclass(Float[mx.array, ""], mx.array)
+
 
 def test_ignored_names():
     x = Float[np.ndarray, "foo=4"]

test/test_decorator.py

@@ -12,6 +12,12 @@
 from .helpers import assert_no_garbage, ParamError, ReturnError
 
 
+try:
+    import mlx.core as mx
+except ImportError:
+    mx = None
+
+
 class M(metaclass=abc.ABCMeta):
     @jaxtyped(typechecker=None)
     def f(self): ...
@@ -128,6 +134,26 @@ def f(x: int, y: int = 1) -> Float[Array, "x {y}"]:
     f(1, 5)
 
 
+@pytest.mark.skipif(mx is None, reason="MLX is not installed")
+def test_mlx_decorator(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def hello(x: Float[mx.array, "8 16"]):
+        pass
+
+    hello(mx.zeros((8, 16), dtype=mx.float32))
+
+    with pytest.raises(ParamError):
+        hello(mx.zeros((8, 14), dtype=mx.float32))
+
+    with pytest.raises(ParamError):
+        import numpy as np
+
+        hello(np.zeros((8, 16), dtype=np.float32))
+
+    with pytest.raises(ParamError):
+        hello(mx.zeros((8, 16), dtype=mx.int32))
+
+
 class _GlobalFoo:
     pass