trax icon indicating copy to clipboard operation
trax copied to clipboard

Published 20 hours ago • verified publisher icon google

Fix issues related to new behavior of JAX DeviceArray.copy()

Open copybara-service[bot] opened this issue 2 years ago • 0 comments

Fix issues related to new behavior of JAX DeviceArray.copy()

In https://github.com/google/jax/pull/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).

copybara-service[bot] avatar Apr 01 '22 01:04 copybara-service[bot]