gh-1021: Add decorator for np conversions

[prady0t]

Description

Adds a NumPy decorator to perform the internal NumPy conversion and change back.

Closes: #1021

Checks

• Is your code passing linting?
• Is your code passing tests?
• Have you added additional tests (if required)?
• Have you modified/extended the documentation (if required)?
• Have you added a one-liner changelog entry above (if required)?

[prady0t]

@lucascolley This could be more generalised for the standard. Currently, I feel the limitation is the inability to specify parameters such as dtype for xp.asarray.

@ntessore @paddyroddy Would love some feedback.

[lucascolley]

cc @betatim, I wonder whether there is any overlap between this and scikit-learn/scikit-learn#34324 which would be nice to upstream into xpx.numpy.

[ntessore]

This is pretty much out of my purview, and @paddyroddy is on leave currently, so it might be a bit before he can answer.

In my opinion, which as I said doesn't carry much weight here, I don't think this is something where GLASS should have to innovate. We're in the cosmological simulations business, not the array business 🙂

[betatim]

There isn't a lot of context in the PR description, what I figure from a quick browse: this PR adds a decorator that converts all "arrays" that are passed to a function to numpy arrays, then calls the wrapped function and converts any "arrays" in the output back to the original namespace.

I'm not sure how much the work in scikit-learn/scikit-learn#34324 is useful for others. It feels somewhat scikit-learn specific. The tricky (imho) part is deciding when to convert to numpy, which I guess is a library specific thing? At least in scikit-learn the decision to convert to numpy is not a simple "this function is hard to support". It depends on the global config, the specific constructor arguments of the class and what the type of the input array is. Once you know you want to convert to numpy it feels like you are "home free".

For this library it looks like most of the code of the decorator deals with finding arrays in the input arguments and such infrastructure. One thing I'm wondering is what to do if the inputs use different array namespaces or if one of the inputs is a torch array and another a list. But maybe these are things aren't an issue here? In scikit-learn we allow a mix of things :-/

Conclusion, I wonder if we can make a generic fallback decorator for xpx that doesn't get mega complex and at the same time supports most of the different behaviours users would want. If the answer is "yes we can" then I think it would be a useful tool (even if I don't think we'd use it in scikit-learn).

[paddyroddy]

This seems like a really nice approach and is the kind of thing I was thinking of. Long-term I definitely wouldn't want the code in this repo as it is complex, and our current approach would be easier for contributors to understand. Do we think we've captured all edge cases here? I assume the decorator is smart enough to detect which variables need to be converted and which don't?

[prady0t]

Currently, the state of the decorator is:

• Goes through all the function parameters (args and kwargs) and collects those which have __array_namespace__ attribute.
• If xp is in kwargs, it uses its value. If value is none, then use default_xp(). xp value overwrites parameter's __array_namespace__ attribute value.
• If no xp parameter and no parameter has __array_namespace__ attribute, returns the function as is.
• Parameters are converted to numpy
• Depending on what the original function returns, the wrapper recursively converts (so that it can support tuple, list and dict) back to the xp.

There are two places where we cannot use this decorator: pixwin and the query_strip function, as we are specifying dtype during conversion. Something that this decorator cannot do yet.

[paddyroddy]

Have fixed the linting in b9aacdb, hope you don't mind.

[paddyroddy]

@connoraird looks like the regression tests don't work on PRs from forks.

Fixed in #1176.

[prady0t]

Have fixed the linting in b9aacdb, hope you don't mind.

Ofc not. I've applied the changes; thanks for all the suggestions. I must have forgotten to remove some of np.asarray which now I have removed.

[paddyroddy]

This looks great! Sorry it has taken so long. We've discussed this offline, and we are happy to merge this. Long-term we're wondering if it would make sense in something like https://github.com/data-apis/array-api-extra, perhaps as the array_api_extra.numpy module suggested above.

[paddyroddy]

Having some regression test failures, but I'm 99% sure its noise. If I keep running them then hopefully it'll fix itself.

[paddyroddy]

I'm merging this now (@connoraird), we need to get the regression tests stable again. Thanks @prady0t!