jax-ml/jax · error · ValueError

'devices' argument to pmap must be non-empty, or None.

Error message

'devices' argument to pmap must be non-empty, or None.

What it means

`jax.pmap`'s `devices` parameter must be None (use all devices) or a non-empty sequence. An empty list/tuple is rejected because pmap needs at least one device to map over.

Source

Thrown at jax/_src/pmap.py:280

  which runs on the first six devices and one on the remaining two:

  >>> from functools import partial
  >>> @partial(pmap, axis_name='i', devices=jax.devices()[:6])
  ... def f1(x):
  ...   return x / jax.lax.psum(x, axis_name='i')
  >>>
  >>> @partial(pmap, axis_name='i', devices=jax.devices()[-2:])
  ... def f2(x):
  ...   return jax.lax.psum(x ** 2, axis_name='i')
  >>>
  >>> print(f1(jnp.arange(6.)))  # doctest: +SKIP
  [0.         0.06666667 0.13333333 0.2        0.26666667 0.33333333]
  >>> print(f2(jnp.array([2., 3.])))  # doctest: +SKIP
  [ 13.  13.]
  """
  if devices is not None:
    if not devices:
      raise ValueError("'devices' argument to pmap must be non-empty, or None.")
    devices = tuple(devices)
  axis_name, static_broadcasted_tuple, donate_tuple = _prepare_pmap(
      fun, axis_name, static_broadcasted_argnums, donate_argnums, in_axes,
      out_axes)
  wrapped_fun = _pmap_wrap_init(fun, static_broadcasted_tuple)
  out_axes_flat, out_axes_tree = tree_flatten(out_axes)
  out_axes_flat = tuple(out_axes_flat)

  def infer_params(*args, **kwargs):
    process_count = xb.process_count(backend)
    trace_state_clean = core.trace_state_clean()
    dyn_f, dyn_argnums, dyn_args = _get_dyn_args(
        wrapped_fun, static_broadcasted_tuple, args)
    dyn_args_flat, dyn_args_tree = tree_flatten((dyn_args, kwargs))
    in_axes_flat = _get_in_axes_flat(
        in_axes, dyn_argnums, dyn_args, kwargs, len(dyn_args_flat),
        dyn_args_tree)
    local_axis_size = _mapped_axis_size(dyn_args_flat, in_axes_flat)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check the device list is non-empty before passing it (fall back to None)
2. Fix backend visibility: install correct CUDA/JAX version so `jax.devices('gpu')` returns devices
3. Pass `devices=None` to use all available devices

Example fix

# before
f = jax.pmap(fn, devices=jax.devices('gpu'))  # empty if no GPU
# after
devs = jax.devices('gpu') or None
f = jax.pmap(fn, devices=devs)

import jax
devs = jax.devices('gpu')
if devs:
    f = jax.pmap(fn, devices=devs)
else:
    f = jax.jit(fn)  # fallback

When it happens

Trigger: Calling `jax.pmap(f, devices=[])` or passing an empty device list computed dynamically, e.g. `devices=jax.devices('gpu')` when no GPU backend is visible.

Common situations: Selecting devices by backend string that returns nothing (no GPU/TPU visible); filtering devices and getting an empty result; CI environments without accelerators.

Related errors

• The devices() method was called on {self._error_repr()}.{sel
• unbound axis name: {axis_name}
• Cannot lower jaxpr with effects: {closed_jaxpr.effects}
• Expected num_classes to match the size of axis {axis}, but {
• Named reductions not implemented for jnp.{name}()

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/7af541d7ad10cfd7. Report an issue: GitHub.