Skip to content

← JAX Conceptual Deep-Dives step 7 of 15

Easy Primitives

JAX and NumPy Array Conversion

Jump to the task

Why this matters

Sometimes a JAX result needs processing by code that works with NumPy arrays on the host. jax.device_get retrieves a concrete array for that work, and jnp.asarray converts the result back to a JAX array.

This is a boundary between array systems, not necessarily between a GPU and CPU. JAX can run on a CPU too. Transfer and copying costs depend on placement and whether memory can be shared; neither conversion always implies a new copy or an accelerator transfer.

Worked mini-example

import jax
import jax.numpy as jnp
import numpy as np

x = jnp.array([1.0, -2.0, 3.0], dtype=jnp.float32)
host_x = jax.device_get(x)
assert isinstance(host_x, np.ndarray)
host_doubled = host_x * 2
result = jnp.asarray(host_doubled)  # JAX array [2, -4, 6].

For an accelerator-backed array, retrieving values on the host can wait for pending work and transfer data. Keep such boundaries intentional, for example when exporting results or calling a host-only library.

Problem

Implement jax_to_numpy_to_jax(x) using the three steps in the example:

  1. Retrieve x with jax.device_get(x).
  2. Multiply the resulting NumPy array by 2.
  3. Return jnp.asarray of that result.

x is a one-dimensional finite float32 JAX array. Assume doubling stays finite. Return a JAX array of the same shape and dtype, including for an empty input. Do not mutate x.

This is an eager interoperability exercise. Its output fixtures check the doubled values, so they cannot prove that a solution actually made the host round trip; 2 * x produces the same output.

Common pitfalls

  • Many JAX operations accept NumPy arrays directly. Explicit conversion can make a boundary clearer, but is not a prerequisite for every call.
  • Ordinary NumPy operations that request concrete values cannot consume a dynamic tracer inside jax.jit. However, device_get itself is not a guarantee of a host transfer during tracing: in this small example, the multiply can remain a traced JAX operation and still compile. A compiled result therefore does not demonstrate host interoperability.
  • jnp.asarray creates a JAX array on the placement selected by JAX; it does not necessarily send it to a GPU. Use explicit placement APIs when a particular device is required.
  • np.asarray can convert a concrete JAX array to NumPy, but is not a general substitute for device_get, which also handles pytrees.
  • Check dtype at the boundary. NumPy often constructs float64 arrays; JAX normally disables 64-bit dtypes and may narrow those inputs. This exercise starts with float32 and keeps it throughout.