Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

make device_array.copy() return a device array #10069

Merged
merged 1 commit into from
Mar 30, 2022

Conversation

jakevdp
Copy link
Collaborator

@jakevdp jakevdp commented Mar 29, 2022

Addresses part of #2632

In #9732 we made jnp.copy return a device array rather than a numpy array. This gives the same treatment to the copy() method of device arrays.

I think this is the behavior that we want – note that it will not chgne the behavior of copy.copy() or copy.deepcopy(), which current return numpy arrays due to their use of the reduce() method (similar to pickle). We could address those separately via the __copy__ and __deepcopy__ methods if we wish.

Copy link
Member

@mattjj mattjj left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This sounds right to me. It's reversing an ancient API decision, so it might break some very old (and hopefully easily fixable) code, but I think jax.device_get or np.array(x) is the right way to spell "give me a NumPy array".

@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Mar 29, 2022
@copybara-service copybara-service bot merged commit ef2efec into google:main Mar 30, 2022
@jakevdp jakevdp deleted the devicearray-copy branch March 30, 2022 14:01
jakevdp added a commit to jakevdp/jax that referenced this pull request Mar 30, 2022
copybara-service bot pushed a commit that referenced this pull request Mar 30, 2022
PiperOrigin-RevId: 438323845
copybara-service bot pushed a commit that referenced this pull request Mar 31, 2022
PiperOrigin-RevId: 438323845
copybara-service bot pushed a commit that referenced this pull request Mar 31, 2022
PiperOrigin-RevId: 438323845
copybara-service bot pushed a commit that referenced this pull request Mar 31, 2022
PiperOrigin-RevId: 438323845
copybara-service bot pushed a commit that referenced this pull request Mar 31, 2022
PiperOrigin-RevId: 438323845
copybara-service bot pushed a commit to google/trax that referenced this pull request Apr 1, 2022
In google/jax#10069, JAX changes the behavior of DeviceArray.copy() so that it returns a DeviceArray rather than returning a numpy array. For converting a DeviceArray to numpy, the preferred method is now np.asarray(device_array).

PiperOrigin-RevId: 438711926
copybara-service bot pushed a commit to google/trax that referenced this pull request Apr 1, 2022
In google/jax#10069, JAX changes the behavior of DeviceArray.copy() so that it returns a DeviceArray rather than returning a numpy array. For converting a DeviceArray to numpy, the preferred method is to explicitly call np.asarray(device_array).

PiperOrigin-RevId: 438718784
copybara-service bot pushed a commit to google/trax that referenced this pull request Apr 4, 2022
In google/jax#10069, JAX changes the behavior of DeviceArray.copy() so that it returns a DeviceArray rather than returning a numpy array. For converting a DeviceArray to numpy, the preferred method is to explicitly call np.asarray(device_array).

PiperOrigin-RevId: 438718784
copybara-service bot pushed a commit to google/trax that referenced this pull request Apr 4, 2022
In google/jax#10069, JAX changes the behavior of DeviceArray.copy() so that it returns a DeviceArray rather than returning a numpy array. For converting a DeviceArray to numpy, the preferred method is to explicitly call np.asarray(device_array).

PiperOrigin-RevId: 439338425
copybara-service bot pushed a commit that referenced this pull request Apr 4, 2022
PiperOrigin-RevId: 438323845
copybara-service bot pushed a commit that referenced this pull request Apr 4, 2022
PiperOrigin-RevId: 438323845
copybara-service bot pushed a commit that referenced this pull request Apr 4, 2022
PiperOrigin-RevId: 439381161
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants