jax-ml/jax · error · ValueError

The denominator cannot be unreduced passed to `div`. Got {y=

Error message

The denominator cannot be unreduced passed to `div`. Got {y=}

What it means

In sharding propagation, the denominator of lax.div must be fully reduced: an unreduced y has no well-defined per-element division under partial shardings, so _div_ur_rule raises ValueError showing y.

Source

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

dispatch.simple_impl(mulhi_p)
ad.defjvp_zero(mulhi_p)
mlir.register_lowering(mulhi_p, partial(_nary_lower_hlo, chlo.mulhi))


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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Constrain the denominator to a fully reduced sharding: with_sharding_constraint(y, P()) before dividing
2. Reshard/reduce y appropriately (all-reduce the denominator) so it is reduced on the numerator's unreduced axes
3. Reorder ops: normalize the denominator outside the sharded region or use lax.div after matching shardings

Example fix

// before
out = x / y  # y unreduced under sharding propagation
// after
y_red = jax.lax.with_sharding_constraint(y, P())
out = x / y_red

y = jax.lax.with_sharding_constraint(y, P()) if propagating_shardings else y
out = lax.div(x, y)

try:
    out = lax.div(x, y)
except ValueError as e:
    if 'denominator cannot be unreduced' in str(e):
        out = lax.div(x, jax.lax.with_sharding_constraint(y, P()))
    else:
        raise

When it happens

Trigger: Sharded division where the denominator operand is unreduced along some axes (partial/reduced sharding mismatch) under NamedSharding/GSPMD jit.

Common situations: Dividing partially-sharded activations by partially-sharded normalizers (RMSNorm/LayerNorm-style) in sharded training; mixing reduced sums with unreduced tensors in elementwise division inside sharded jits.

Related errors

• {name} cannot accept args which are unreduced. Got {a.str_sh
• unreduced rule for {name} is not implemented. Please file an
• `out_sharding` passed to `broadcast_in_dim` can only contain
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m

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