Skip to content

← JAX Performance & Distributed step 21 of 25

Hard Primitives

JIT with Input and Output Sharding

Jump to the task

Why this matters

jax.jit can compile a function with specified input and output shardings. You write operations on global arrays, and the compiler distributes the computation over the mesh. Some operations work independently on each shard; others need communication between devices.

Older examples call this API pjit. For new code, use jax.jit with in_shardings and out_shardings. The exercise keeps its original pjit_double function name but uses the current API.

Worked mini-example

import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec

mesh = Mesh(np.array(jax.devices()[:1]), ('data',))
sharding = NamedSharding(mesh, PartitionSpec('data'))
double = jax.jit(
    lambda x: 2 * x,
    in_shardings=sharding,
    out_shardings=sharding,
)
x = jnp.array([2.0, -1.0, 0.5], dtype=jnp.float32)
result = double(x)
print(result)  # [ 4. -2.  1.]
assert result.sharding.is_equivalent_to(sharding, ndim=1)

The mesh selects only the first device, even if more are available. It teaches the API without requiring multiple GPUs. With a larger data mesh, the same elementwise computation can double separate pieces of a vector. Sharding controls layout; multiplying by two changes the values.

Problem

Implement pjit_double(x):

  1. Create a mesh from jax.devices()[:1], with axis name 'data'.
  2. Build NamedSharding(mesh, PartitionSpec('data')).
  3. Compile a function that doubles its argument, passing that sharding to both in_shardings and out_shardings of jax.jit.
  4. Call the compiled function on x and return its result.

x is a one-dimensional float32 JAX array. Return a float32 array with the same shape and each value doubled. Numerical fixtures check the values; they do not prove that your implementation used JIT or the requested layout. Inspect the result's .sharding separately.

Common pitfalls

  • For equal partitions, the vector length must be divisible by the size of the mesh axis assigned to it. Count the data axis, not every device in an unrelated mesh dimension. A size-one data axis accepts any vector length.
  • A single Sharding can apply across multiple positional arguments. For different layouts, provide a matching tuple or nested structure. Input and output shardings need not be identical; they describe different boundaries of the computation.
  • An already committed input with an incompatible layout may be rejected. Use jax.device_put(x, sharding) to explicitly place it in the expected layout before calling the compiled function.
  • One-device placement is not necessarily metadata-only. Data may need to move to that device, and compilation still has a cost.
  • In a repeated workload, create the compiled function once and reuse it. Recreating a fresh lambda inside every call can defeat compilation caching. This small exercise includes setup inside the function for convenience.

The diagram below is a separate matrix example with two mesh axes. It splits rows across data and replicates each row block across model. The submitted function uses a vector and a one-device mesh.

Loading visualization…