Add support for the __array_namespace_info__

[aaishwarymishra]

Proposed changes

Please include a description of the problem or feature this PR is addressing. If there is a corresponding issue, include the issue #.
This pull request adds support for the __array_namespace_info__ to make it complaint with the array api spec
https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.__array_namespace_info__.html#array-namespace-info

Ai was used for writing test :)
closes #4239

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[aaishwarymishra]

Hi @zcbenz , I had few questions like is there max dimenstions limit in mlx? also is there some default dtypes mlx has for different devices, like in the case of the gpu, default would be float32, what about in cpu, also are there specific dtypes for integral and indexing, I would appreciate your input :)

[zcbenz]

For cuda backend the max ndim is 10, there is no real restrictions on other backends but ndim larger than 10 would make things slower so I think we can just return 10. There is no real default dtype, the dtype is just set based on the python type.

[aaishwarymishra]

I think its ready for review :)

[ev-br]

I'll note that the lack of __array_namespace_info__ is currently the main blocker for running the Array API test suite (tracked in data-apis/array-api-compat#450). The latter has been advancing a lot recently, thanks to great efforts of @aaishwarymishra and @prady0t -- and @zcbenz obviously, thank you for responding to our PRs!

I left two comments below from the Array API perspective. Regardless of the MLX decision on the default integral dtype (as detailed below), I think landing this PR would be great, as it would unblock further incremental work on Array API compatibility for us downstream.

[ev-br]

This is the dtype to match the C type for indexing, and this basically says that arrays of size >=INT_MAX are not really supported, and that indexing them is UB for integer overflow. Is this really the case (and if it is, is this the intended state, or is supporting 64-bit integer sized arrays somewhere on the roadmap?)

[aaishwarymishra]

Hi @ev-br #4238 recently fixed that, I forgot that it got merged, I think we can use int64 for the indexing atleast.

[zcbenz]

The array shape in MLX is int32[] so I think using in32 for indexing is fine.

[ev-br]

There is no real default dtype, the dtype is just set based on the python type.

It looks like MLX is using value-based semantics, not just type-based semantics (and so it sort of follows NumPy 1.x, pre NEP 50 rules):

>>> int_max = mx.iinfo(mx.int32).max
>>> mx.asarray(int_max)
array(2147483647, dtype=int32)
>>> mx.asarray(int_max+1)
array(2147483648, dtype=int64)

There is no perfect way to fully express this for Array API. As a piece of prior art, jax (which implements Array API compatibility in its main namespace) avoids the value-based semantics via its ENABLE_X64 switch:

• in the enable_x64=1 mode, the dtype of asarray(python_int) is always int64
• in the default mode, the dtype of asarray(python_int) is always int32 , and if a value is too large, it raises a ValueError

I wonder if MLX's is open to moving away from value-based semantics @zcbenz? (long-term, not in this PR, obviously).
Experience in the numpy land is that value-based semantics brings more harm than good.

In this PR, I think this deserves a code comment @aaishwarymishra .

[prady0t]

Link to NEP 50. Thanks for sharing!

[zcbenz]

I don't like the value-based semantics, it is a historical decision that I just follow. I'm open to move away from it but it would require some discussions before making the decision.

[zcbenz]

Thanks for the PR and the reviews!