Add support for copy kwarg in astype to match Array API

[Micky774]

Towards #20200

cf. data-apis/array-api#665

Note

This PR includes a placeholder utility function for managing device transfer, which will be developed in a follow-up PR

Default behavior is preserved when copy=None

[jakevdp]

Suggestion: let's implement these changes in jax.numpy.astype here:

jax/jax/_src/numpy/lax_numpy.py

Line 2212 in f8af4ef

def astype(x: ArrayLike, dtype: DTypeLike | None, /, *, copy: bool = True) -> Array:

Then we can just call jax.numpy.astype in the current location.

Also we could add some test coverage in lax_numpy_test.py

[jakevdp]

One small suggestion: it may be helpful to separate the copy change from the device change, since the way forward on the former is clear, and the way forward on the latter is less certain. Do you want to do a separate PR changing the semantics of the copy parameter in astype?

[Micky774]

One small suggestion: it may be helpful to separate the copy change from the device change, since the way forward on the former is clear, and the way forward on the latter is less certain. Do you want to do a separate PR changing the semantics of the copy parameter in astype?

Agreed. For now I'm just factoring out the device semantics as mentioned here and leaving only the changes to copy so I'll update the PR title accordingly. I'm hoping eventually we can have a common utility for general memory management for both device/copy semantics.

[Micky774]

I've created this PR in my fork to track discussion regarding device movement semantics using jnp.astype as an example. Now this PR passes through the device kwarg until it is ignored in the _place_array utility, which will eventually house the genuine device moment functionality.

[jakevdp]

Thanks!

[jakevdp]

Oh, it looks like this branch has conflicts that need to be resolved before we can merge.

[jakevdp]

Actually, I'm a bit worried about this change. In a lot of places we implicitly depend on x.astype(x.dtype) being a no-op, and now that the default semantics are copy=True, I think this assumption will no longer hold.

Maybe we should default to copy=None so that the default is backward-compatible?

NumPy has made a similar exception for copy in astype: https://numpy.org/neps/nep-0056-array-api-main-namespace.html#copy-keyword-semantics

[Micky774]

If it is causing that many problems, then I agree going the gentler deprecation route as opposed to the "bug-fix" route should be preferred.

With that being said, I'm suspecting that the array API standard for astype is actually "mistaken" about their copy semantics. The original array API specification was based on NumPy's -- and other contemporary libraries' -- implementations and copy semantics. For reference, I've opened an issue discussing updating the copy semantics to be made similar to the other functions/methods: data-apis/array-api#788.

Maybe we wait to see what the community consensus on the proposal is, that way we can tune the numpy namespace to the False/None/True semantics if its accepted, and then we can stick to the "old" semantics in the array_api namespace just for technical compliance during the short term.

[jakevdp]

It seems like the array API will be sticking with copy=False/True.

I think what we should do here is follow NumPy and the Array API on the semantics of the copy argument for astype, except we should default to copy=False (which behaves like copy=None in other contexts).

In JAX, there's almost never a reason to force a copy, because arrays are immutible and copies are expensive. In NumPy, it's the opposite: you almost never want to silently return a view, because arrays are mutable.

[Micky774]

@jakevdp Updated w/ a default of copy=False

[jakevdp]

Thanks - I'm still a bit confused about the state of things upstream. It sounds like NumPy does support copy in {True, False, None} for this function? Is that your understanding too?

[Micky774]

Thanks - I'm still a bit confused about the state of things upstream. It sounds like NumPy does support copy in {True, False, None} for this function? Is that your understanding too?

That seems not to be the case in NumPy 1.26, but is indeed true in NumPy 2.0.0rc1 with the caveat that copy=None is equivalent to copy=False. This is the same functionality that is enabled in this PR, except we do not publicly document the fact that copy=None is a valid choice. FWIW, I prefer not documenting it since I think having two different-yet-equivalent values for a keyword argument makes for a confusing API.

I'm personally hoping we see a change in NumPy and the array API to converge the astype copy semantics to be equivalent to the other copy-enabled functions/methods, but that probably won't be too soon, if it happens at all.