jax-ml/jax · error · NotImplementedError

Named reductions not implemented for jnp.{name}()

Error message

Named reductions not implemented for jnp.{name}()

What it means

The reduction was invoked with named axes (e.g. axis='batch' from pmap/jit named axes) but this particular reduction has no named-axis (collective/pmap) implementation registered in JAX.

Source

Thrown at jax/_src/numpy/reductions.py:153

    if promote_integers:
      result_dtype = _promote_integer_dtype(result_dtype)
  else:
    result_dtype = dtypes.check_and_canonicalize_user_dtype(dtype, name)

  if upcast_f16_for_computation and dtypes.issubdtype(result_dtype, np.inexact):
    computation_dtype = _upcast_f16(result_dtype)
  else:
    computation_dtype = result_dtype
  a = lax.convert_element_type(a, computation_dtype)
  op = op if computation_dtype != np.bool_ else bool_op
  # NB: in XLA, init_val must be an identity for the op, so the user-specified
  # initial value must be applied afterward.
  init_val = _reduction_init_val(a, init_val)
  if where_ is not None:
    a = _where(where_, a, init_val)
  if pos_dims is not dims:
    if parallel_reduce is None:
      raise NotImplementedError(f"Named reductions not implemented for jnp.{name}()")
    result = parallel_reduce(a, dims)
  else:
    result = lax.reduce(a, init_val, op, dims)
  if initial is not None:
    initial_arr = lax.convert_element_type(initial, lax.asarray(a).dtype)
    if initial_arr.shape != ():
      raise ValueError("initial value must be a scalar. "
                       f"Got array of shape {initial_arr.shape}")
    result = op(initial_arr, result)
  if keepdims:
    result = lax.expand_dims(result, pos_dims)
  return lax.convert_element_type(result, dtype or result_dtype)

def _canonicalize_axis_allow_named(x, rank):
  return maybe_named_axis(x, lambda i: canonicalize_axis(i, rank), lambda name: name)

def _reduction_dims(a: ArrayLike, axis: Axis):
  if axis is None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use positional integer axes instead of named axes
2. Use lax.psum / the collective equivalent (e.g. jax.lax.pmax) for named-axis reductions
3. Move the reduction off the named axis: reduce positional axes first, then use psum on the remainder

Example fix

// before
jax.pmap(lambda x: jnp.logsumexp(x, axis='dev'), axis_name='dev')(x)
// after
jax.pmap(lambda x: jax.lax.pmax(x, 'dev').sum(-1), axis_name='dev')(x)

if isinstance(axis, str):  # named axis unsupported
    axis = None  # or use lax.psum/pmax collectives instead

When it happens

Trigger: Calling jnp.<reduce>(x, axis='dev') inside jax.pmap where the reduction lacks a parallel_reduce implementation; mixing named axes with an uncommon reduction op.

Common situations: Using pmap with named axes and reductions that only support positional axes; migrating multi-device code that worked with lax.psum but not this jnp reduction.

Related errors

• unbound axis name: {axis_name}
• Cannot lower jaxpr with effects: {closed_jaxpr.effects}
• vary_unreduced_cast only accepts inputs that are varying. Go
• vary_unreduced_cast is a Varying->Unreduced collective. This
• vary_unreduced_cast input cannot be unreduced across the axi

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