Support for MLX

[gabrieldemarmiesse]

Hello,
I'm opening an issue to request support for MLX. Currently we get TypeError: type 'nanobind.nb_type_0' is not an acceptable base type. Hopefully it will gather interest from the community.

import mlx.core as mx
from beartype import beartype
from jaxtyping import jaxtyped, Int


@jaxtyped(typechecker=beartype)
def hello(a: Int[mx.array, "8"]):
    pass

hello(mx.zeros((8,), dtype=mx.int32))

Traceback (most recent call last):
  File "/Users/gabrieldemarmiesse/projects/V2LLMs/moshi_mlx/./moshi_mlx/a.py", line 7, in <module>
    def hello(a: Int[mx.array, "8"]):
                 ~~~^^^^^^^^^^^^^^^
  File "/Users/gabrieldemarmiesse/projects/V2LLMs/moshi_mlx/.venv/lib/python3.12/site-packages/jaxtyping/_array_types.py", line 678, in __getitem__
    out = _make_array(array_type, dim_str, cls)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/gabrieldemarmiesse/projects/V2LLMs/moshi_mlx/.venv/lib/python3.12/site-packages/jaxtyping/_array_types.py", line 617, in _make_array
    else _make_metaclass(type(array_type))
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/gabrieldemarmiesse/projects/V2LLMs/moshi_mlx/.venv/lib/python3.12/site-packages/jaxtyping/_array_types.py", line 327, in _make_metaclass
    class MetaAbstractArray(_MetaAbstractArray, base_metaclass):
TypeError: type 'nanobind.nb_type_0' is not an acceptable base type

Cheers!

[patrick-kidger]

Looks like this might be an MLX issue, but happy to take a PR here if it would help.

[gabrieldemarmiesse]

Do you have a hint on how to get started with this? I'm noticing that jaxtyping is trying to subclass anything passed as the type hint: https://github.com/patrick-kidger/jaxtyping/blob/main/jaxtyping/_array_types.py#L620 but mx.array can't be subclassed, likely because it's coming from nanobind. Is there any way to skip this step without using Any? Using Any would allow numpy arrays to be passed for example, and it's not what we want.

[patrick-kidger]

FWIW I don't completely remember why I went to the effort of subclassing the provided type, instead of treating this as a completely independent type. Does removing that break any tests?

[gabrieldemarmiesse]

@patrick-kidger #301 should fix this