Fix array api integration for additive_chi2_kernel with torch mps

[OmarManzoor]

Reference Issues/PRs

Follow up of #29144

What does this implement/fix? Explain your changes.

• Fixes a test that was failing for the additive_chi2_kernel array api with torch mps

Any other comments?

[github-actions]

✔️ Linting Passed

All linting checks passed. Your pull request is in excellent shape! ☀️

_{Generated for commit: 43d860c. Link to the linter CI: here}

[Tialo]

Oh, I overlooked it has to be on the same device. Maybe common tests should also assert it? It should spot the error at least in gpu workflows

[OmarManzoor]

Oh, I overlooked it has to be on the same device. Maybe common tests should also assert it? It should spot the error at least in gpu workflows

I think this might be specifically an issue with mps because it seems to work fine with cuda.

[EdAbati]

I confirm that this fix is needed for the mps device 😕

[EdAbati]

I think we could also use get_namespace_and_device and substitute it to xp, _ = get_namespace(X, Y) below

[EdAbati]

FYI I have updated to the latest torch (2.3.1) and it looks like specifying the device here is not required anymore when using mps. 👀

[ogrisel]

FYI I have updated to the latest torch (2.3.1) and it looks like specifying the device here is not required anymore when using mps. 👀

Nice that they did this. It was verbose to have to specify the device even for scalar arrays.

I am -0 for merging this PR. Since array API is very new, one can expect people who want to try it to use the latest stable versions of libraries and I would rather keep our code based not too verbose if not necessary.

What other people think? We could make that requirement explicit in the _array_api.rst file.

[OmarManzoor]

FYI I have updated to the latest torch (2.3.1) and it looks like specifying the device here is not required anymore when using mps. 👀

Nice that they did this. It was verbose to have to specify the device even for scalar arrays.

I am -0 for merging this PR. Since array API is very new, one can expect people who want to try it to use the latest stable versions of libraries and I would rather keep our code based not too verbose if not necessary.

What other people think? We could make that requirement explicit in the _array_api.rst file.

I think it makes sense if it is fixed. No need for the change. @ogrisel should we close this PR?

[ogrisel]

Let's wait a bit to let others have time to express their opinion.

[ogrisel]

Let's close, we can reopen if others disagree.

[ogrisel]

Note: we would have to check with other libraries. For instance I think jax support in array-api-compat is improving so we might include it in our test suite at some point and we might end-up needing this change.