liveBook · Manning

Jax.device_get()

As explained in Deep Learning with JAX

A JAX function that transfers data from a device back to the Python host process as a NumPy ndarray.

Need to inspect or use a device-resident value in ordinary Python? jax.device_get() transfers it from the device to the host process and returns a NumPy ndarray. section 3.2.3

Bring an array back to the host

arr_host = jax.device_get(arr)\ntype(arr_host)\n>>> numpy.ndarray\n\narr_host\n>>> array([1, 42, 31337], dtype=int32)\n\nUse this boundary when host-side Python or NumPy needs the values. Device placement remains a separate concern handled by jax.device_put(). section 3.2.3

How Deep Learning with JAX covers it

How deep it goes
Defines the reverse of device placement and specifies the NumPy return type, without detailing transfer costs or synchronization.
Examples
Demonstrates converting an array into a NumPy ndarray and inspecting available CPU and GPU devices.
What it assumes
Assumes JAX arrays and device placement while establishing the host-versus-device distinction.
What you can do after
Apply it via jax.device_get() to obtain device data as a NumPy ndarray in the host process.
How it connects
Connects device placement with jax.device_put(), device discovery, committed data, and NumPy ndarray values.
Cover of Deep Learning with JAX
liveBookmanning.com
Read Deep Learning with JAX on liveBook Get the book at manning.com

Related concepts

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.