EagerJAXArrayContext: return mutable numpy array in to_numpy

[matthiasdiener]

AFAICS, there is no way to make an immutable numpy array mutable again except for copying it.

[Copilot]

Pull Request Overview

This PR modifies the to_numpy method in the JAX array context to return a mutable NumPy array by explicitly copying the result of jax.device_get.

• Wraps jax.device_get result in np.copy to produce a writable array.

Comments suppressed due to low confidence (1)

arraycontext/impl/jax/init.py:104

• The code uses np.copy but numpy is not imported in this scope, which will raise a NameError. Add import numpy as np at the module level or within the function.

return np.copy(jax.device_get(ary))

[alexfikl]

I'm a bit out of the loop, so I have questions 😁

1. Why is this required? To my knowledge, most of the code doesn't do in-place modifications (e.g. DOFArray in-place operations are very deprecated).
2. Haven't tested, but jax.device_get docs mention copies being done. Is that not the case?
3. Would this be useful when using the CPU backend for jax? I image the GPU/TPU backends do a copy anyway?

[matthiasdiener]

Thanks for the review @alexfikl !

1. Why is this required? To my knowledge, most of the code doesn't do in-place modifications (e.g. DOFArray in-place operations are very deprecated).

There are in-place modifications in grudge, e.g. grp_field in https://github.com/inducer/grudge/blob/193100abd566719b26001b85feb3d7d2171f5b8d/grudge/geometry/metrics.py#L567-L577

Without this PR, some tests in grudge with Jax fail (see e.g. inducer/grudge#380).

2. Haven't tested, but jax.device_get docs mention copies being done. Is that not the case?
3. Would this be useful when using the CPU backend for jax? I image the GPU/TPU backends do a copy anyway?

I think there are copies done, both for CPU and GPU, the problem is that jax.device_get returns an immutable numpy array (i.e., ary.flags.WRITEABLE is set to False). As far as I can tell, the only way around that is to make another copy, as in this PR. If ary.flags.WRITEABLE is already set to False, numpy refuses to set the flag back to True.

[alexfikl]

There are in-place modifications in grudge, e.g. grp_field in inducer/grudge@193100a/grudge/geometry/metrics.py#L567-L577

I remember writing some version of that code.. it should really be rewritten in a nicer way :(

But fair enough, I can see how that breaks. Are there any other places? (the tests in inducer/grudge#380 seem to be passing nicely at the moment)

EDIT: It wasn't running tests with jax, I see them now! 😁

[alexfikl]

I would personally prefer if this gets fixed in grudge, since it's just that one place that seems to fail (?). Doing two copies on to_numpy here seems like a hack.

If we do go for it, we should also document it somewhere that the arrays returned by to_numpy are writable, since it's apparently not a given!

[matthiasdiener]

I would personally prefer if this gets fixed in grudge, since it's just that one place that seems to fail (?). Doing two copies on to_numpy here seems like a hack.

I agree, and I've just added a (sketchy) reimplementation in inducer/grudge#380, feel free to take a look. Setting this PR back to draft.

If we do go for it, we should also document it somewhere that the arrays returned by to_numpy are writable, since it's apparently not a given!

Right, either way of what happens with this PR, we should document this in to_numpy.