jax-ml/jax · error · ValueError

Denominator should be reduced along the same axes numerator

Error message

Denominator should be reduced along the same axes numerator is unreduced on. Got {x=}, {y=}

What it means

Raised by the unreduced-sharding rule for lax.div: in JAX's named-sharding 'unreduced' mechanism, the numerator of a division may be unreduced along some axes, but the denominator must be reduced along exactly those same axes. This error fires when the numerator's unreduced axes don't match the denominator's reduced axes, which would make the division semantically ill-defined (partial sums in the denominator but not the numerator).

Source

Thrown at jax/_src/lax/lax.py:5156

def _div_transpose_rule(cotangent, x, y):
  assert ad.is_undefined_primal(x)
  if ad.is_undefined_primal(y):
    raise RuntimeError("nonlinear div can't be transposed")
  if type(cotangent) is ad_util.Zero:
    return [ad_util.Zero(x.aval), None]
  else:
    return [_unbroadcast(x.aval, div(cotangent, y)), None]

def _div_ur_rule(x, y):
  out_reduced = default_nary_reduced_rule(x, y)
  x_ur, y_ur = getu(x), getu(y)
  if y_ur:
    raise ValueError(
        f'The denominator cannot be unreduced passed to `div`. Got {y=}')
  if x_ur and x_ur != getr(y):
    raise ValueError(
        'Denominator should be reduced along the same axes numerator is'
        f' unreduced on. Got {x=}, {y=}')
  out_unreduced = x_ur
  if out_unreduced:
    assert out_reduced == out_unreduced, (out_reduced, out_unreduced)
    out_reduced = frozenset()  # if both are equal, set difference is empty.
  kind = UnreducedKind.sum if out_unreduced else None
  return out_unreduced, out_reduced, kind

div_p = standard_naryop([_num, _num], 'div', ur_rule=_div_ur_rule)
ad.defjvp(div_p,
          lambda g, x, y: div(g, y),
          lambda g, x, y: mul(mul(neg(g), x), integer_pow(y, -2)))
ad.primitive_transposes[div_p] = _div_transpose_rule
mlir.register_lowering(div_p, partial(_nary_lower_hlo, hlo.divide))

rem_p = standard_naryop([_int | _float, _int | _float], 'rem')
ad.defjvp(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the denominator is fully reduced (e.g. apply all_gather/reduce on y) before dividing
2. Make the numerator and denominator sharding specs consistent: reduce y along exactly the axes x is unreduced on
3. Do the division after an explicit lax.psum on the denominator
4. Check the sharding specs of both operands with jax.debug.inspect_array_sharding to spot the mismatch

Example fix

// before
z = jnp.sum(x_unreduced, axis=0) / y  # y not reduced on same axes

// after
y_red = lax.psum(y, axis=0) if is_unreduced(y) else y
z = jnp.sum(x_unreduced, axis=0) / y_red

def safe_div(x, y):
    xu, yr = get_unreduced(x), get_reduced(y)
    if yr:
        raise ValueError('unreduced denominator')
    if xu and xu != yr:
        y = lax.psum(y, axis=tuple(sorted(xu - yr))) if (xu - yr) else y
    return lax.div(x, y)

def is_div_sharding_ok(x, y) -> bool:
    return not getu(y) and (not getu(x) or getu(x) == getr(y))

When it happens

Trigger: Calling jax.lax.div (or '/' on arrays) where operands carry NamedSharding specs with unreduced axes, e.g. after pjit/jit with a sharding where x is unreduced on axis a but y is not reduced on axis a; also passing an unreduced denominator directly (a sibling error).

Common situations: Migrating from pmap/pjit collectives to automatic partially-reduced sharding; specifying out_specs / in_specs with unreduced markers inconsistently between operands of a division; sum-then-divide patterns where only one side went through a reduced collective.

Related errors

• Mapped away dimension of inputs passed to vmap should be sha
• The 'sharding' attribute is not available on {self._error_re
• The is_fully_addressable property was called on {self._error
• The error code state and the predicate must be on the same m
• `sharding` passed to `convert_element_type` can only contain

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