Added docs and tests for MLX support

Author: patrick-kidger

README.md

@@ -2,15 +2,15 @@
 
 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) arrays; *(Now also supports PyTorch, NumPy, MLX, and TensorFlow!)*
 2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).
 
-
 **For example:**
 ```python
 from jaxtyping import Array, Float, PyTree
 
 # Accepts floating-point 2D arrays with matching axes
+# You can replace `Array` with `torch.Tensor` etc.
 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,11 +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) arrays; *(Now also supports PyTorch, NumPy, MLX, and TensorFlow!)*
 2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).
 
- *(Now also supports PyTorch, NumPy, and TensorFlow!)*
-
 ## Installation
 
 ```bash
@@ -25,6 +23,7 @@ The annotations provided by jaxtyping are compatible with runtime type-checking
 from jaxtyping import Array, Float, PyTree
 
 # Accepts floating-point 2D arrays with matching axes
+# You can replace `Array` with `torch.Tensor` etc.
 def matrix_multiply(x: Float[Array, "dim1 dim2"],
                     y: Float[Array, "dim2 dim3"]
                   ) -> Float[Array, "dim1 dim3"]:

test/requirements.txt

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

test/test_decorator.py

@@ -268,3 +268,23 @@ class X(eqx.Module):
     X(1)
     with pytest.raises(ParamError):
         X("1")
+
+
+def test_mlx(jaxtyp, typecheck):
+    import mlx.core as mx
+    import numpy as np
+
+    @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):
+        hello(np.zeros((8, 16), dtype=np.float32))
+
+    with pytest.raises(ParamError):
+        hello(mx.zeros((8, 16), dtype=mx.int32))