Skip to content

← JAX Performance & Distributed step 7 of 25

Medium Primitives

Host Roundtrip via device_get

Jump to the task

Use a host library, then return to JAX

A NumPy-only library expects host-accessible data. For the concrete JAX arrays in this exercise, jax.device_get(x) returns a NumPy array. After doing host work, jnp.asarray(result) creates a JAX array again.

This is a round trip between array interfaces. On an accelerator it can require transfers in both directions and waiting for input values. On a CPU, memory sharing and copying can differ. Cost depends on the array size, device, runtime and pending work; there is no fixed millisecond cost for every call.

Follow the types through the calculation

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

x = jnp.array([1., 2., 3.], dtype=jnp.float32)
host = jax.device_get(x)
assert isinstance(host, np.ndarray)
doubled = host * 2               # NumPy arithmetic: [2., 4., 6.]
result = jnp.asarray(doubled)     # a JAX array again

Multiplication creates the doubled host result without modifying x. Do not assume the array returned by device_get is writable. If a host library needs to mutate its input, give it an explicit NumPy copy.

Converting a concrete JAX array with np.asarray or np.array can also bring values to the host. That conversion alone is not a full round trip: the return to JAX is a separate step. Likewise, jnp.asarray does not promise to restore the original array's device or sharding.

Keep this boundary in ordinary Python

The type sequence above describes eager execution on concrete arrays. Inside a JIT trace, device_get is not a way to turn a tracer into runtime NumPy data. A NumPy-only operation that needs those values can fail with a tracer-conversion error. A simple expression such as * 2 may instead trace as JAX arithmetic, bypassing the intended host work.

Keep this exercise's wrapper outside JIT. For a real application, isolate host-library calls between compiled sections or use an appropriate callback API with its own transformation rules. A custom derivative rule alone does not perform a host transfer.

Doubling is already available in JAX as 2 * x, so production code has no reason to make this particular round trip. Here the simple arithmetic lets you focus on the conversion steps. In a hot path, first check whether the needed operation exists in JAX, then profile any necessary host work rather than assuming a universal slowdown.

Problem

Implement host_roundtrip(x) as an ordinary Python function:

  1. Obtain the host array with jax.device_get(x).
  2. Double that NumPy array without modifying it in place.
  3. Return jnp.asarray of the doubled values.

X is a one-dimensional float32 JAX array. Return a float32 JAX array of the same shape. For [5, 10], return [10, 20]. An empty input stays empty. Do not return a NumPy array or a Python list.

Numerical output tests check the result, not where multiplication occurred. Inspect the host/JAX conversion steps as well.