Skip to content

← JAX Performance & Distributed step 5 of 25

Easy Primitives

Explicit device_put

Jump to the task

Make array placement explicit

A JAX device can be a CPU, GPU or TPU. This exercise works on a CPU; it does not require an accelerator.

jax.device_put(x) with no target has two useful cases:

  • A NumPy array becomes a JAX array at the default placement.
  • An existing JAX array keeps its current placement. It is not forced back to the default device.

To choose a particular device, supply it as the second argument. A sharding can instead describe placement across several devices. Without an explicit target, newly placed host data is uncommitted: it can follow explicitly placed data when used together in a calculation.

From NumPy to JAX

import jax
import numpy as np

host = np.array([1., 2., 3.], dtype=np.float32)
placed = jax.device_put(host)
doubled = 2 * placed
# doubled is a JAX array with values [2., 4., 6.]

chosen = jax.devices()[0]
committed = jax.device_put(host, chosen)
# chosen can be a CPU; the result is committed to that device.

Placement may involve a transfer, but it does not guarantee a fresh physical copy in every case. In particular, avoid reasoning about CPU memory sharing from the array's Python type alone.

device_put can return before a transfer finishes. JAX tracks the dependency when the next calculation uses the result. To measure transfer completion separately, wait on the placed array with .block_until_ready() before stopping your timer.

Transferring a new batch in a training loop can be necessary. The avoidable cost is repeatedly transferring unchanged data that could stay on the device. Measure the workload rather than assuming that every call copies data or that placement is useful only with many devices.

Problem

Implement device_put_then_double(x):

  1. Call jax.device_put(x) without an explicit target.
  2. Multiply the returned JAX array by two.
  3. Return that result.

X is a one-dimensional float32 NumPy or JAX array. Return a float32 JAX array of the same shape. For [5, 6], return [10, 12]; for an empty array, return shape (0,). Do not modify the input.

This task asks you to use the placement API; it does not ask you to select a GPU or force an already placed array onto another device. Numerical tests alone cannot prove that you called device_put or identify physical transfers, so also check your implementation.