Add support for MLX arrays #301
Conversation
gabrieldemarmiesse
changed the title
jaxtyping/_array_types.py
Outdated
| @@ -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)) | |||
There was a problem hiding this comment.
I was afraid things might look something like this!
FWIW we could maybe just remove subclassing altogether -- I'm not sure how much this buys us, and we're really threading the needle here to do so -- WDYT?
(I think I might have introduced this subclassing to make things work better when using type annotations with plum, although I might have that wrong. I think that's a use-case that can be supported in other ways, though.)
There was a problem hiding this comment.
I have no opinion on the matter. I'll remove the subclassing entirely and see where it goes
There was a problem hiding this comment.
______________________________________________________ test_subclass _______________________________________________________
def test_subclass():
> assert issubclass(Float[Array, ""], Array)
E AssertionError: assert False
E + where False = issubclass(<class 'jaxtyping.Float[Array, '']'>, Array)
test/test_array.py:607: AssertionError
The only test failing is this one: test/test_array.py::test_subclass - AssertionError: assert False. Every assertion on this test is failing.
I find it strange that a type hint need to be a subclass of a type. If that's not needed, we can remove this completely. How should I proceed?
There was a problem hiding this comment.
Okay. I think let's go ahead and remove it then!
There was a problem hiding this comment.
Done, I'll let you review one more time!
| else _make_metaclass(type(array_type)) | ||
| ) | ||
|
|
||
| metaclass = _make_metaclass(type) | ||
| out = metaclass( |
There was a problem hiding this comment.
Can we skip creating the metaclass too? And have this line just be type(...
Fixes #299