jax-ml/jax · error · ValueError

reduce_sum on operand {operand.str_short(True)} is not allow

Error message

reduce_sum on operand {operand.str_short(True)} is not allowed when jax_allow_f16_reductions=False.

What it means

JAX refuses reduce_sum on float16/bfloat16 inputs when the config flag jax_allow_f16_reductions is False, unless every reduced axis has size 1. This is because f16 accumulation on many backends loses precision or is unsupported, so it is gated behind an explicit opt-in flag.

Source

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

def _reduce_op_reduced_rule(operand, out_sharding, name):
  if out_sharding is not None and out_sharding.spec.reduced:
    raise ValueError(
        f'out_sharding passed to {name} cannot be reduced. Got {out_sharding=}')
  return getr(operand)

def _reduce_sum_ur_rule(operand, *, axes, out_sharding):
  out_unreduced, kind = _reduce_op_unreduced_rule(
      operand, axes, out_sharding, UnreducedKind.sum, 'reduce_sum')
  out_reduced = _reduce_op_reduced_rule(operand, out_sharding, 'reduce_sum')
  return out_unreduced, out_reduced, kind

def _reduce_sum_dtype_rule(operand, *, axes, **_):
  dt = _reduce_number_dtype_rule('reduce_sum', operand)
  if (operand.dtype in [np.float16, dtypes.bfloat16] and
      not config.allow_f16_reductions.value and
      not all(core.definitely_equal(operand.shape[d], 1) for d in axes)):
    raise ValueError(f"reduce_sum on operand {operand.str_short(True)} is not "
                     "allowed when jax_allow_f16_reductions=False.")
  return dt

reduce_sum_p = standard_primitive(
  _reduce_op_shape_rule, _reduce_sum_dtype_rule,
  'reduce_sum', sharding_rule=_reduce_op_sharding_rule_with_out_sharding,
  vma_rule=partial(core.standard_vma_rule, 'reduce_sum'),
  ur_rule=_reduce_sum_ur_rule)
ad.deflinear2(reduce_sum_p, _reduce_sum_transpose_rule)
batching.defreducer(reduce_sum_p)

def _reduce_prod_jvp_rule(primals, tangents, *, axes):
  reducer = lambda x, y: [mul(x, y)]
  primals_out, tangents_out = _reduce_jvp(reducer, [_const(primals[0], 1)],
                                          primals, tangents, axes)
  return primals_out[0], tangents_out[0]

def _reduce_op_sharding_rule(operand, *, axes):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast the operand to float32 before reducing: jnp.sum(x.astype(jnp.float32), axis).astype(x.dtype).
Opt in explicitly: set jax.config.update('jax_allow_f16_reductions', True) (with JAX_JAX_ALLOW_F16_REDUCTIONS=True env var) if your backend's f16 reductions are acceptable.
Check which layer produced the half-precision sum (e.g., a loss function) and keep reductions in fp32 there.

Example fix

# before
loss = jnp.sum(squared_err_bf16)
# after
loss = jnp.sum(squared_err_bf16.astype(jnp.float32))

Defensive patterns

Strategy: validation

Validate before calling

def safe_sum(x, axis=None):
    if x.dtype in (jnp.float16, jnp.bfloat16):
        x = x.astype(jnp.float32)
    return jnp.sum(x, axis=axis)

Type guard

def needs_f32_upcast(x):
    return x.dtype in (jnp.float16, jnp.bfloat16)

Prevention

Keep loss/gradient reductions in fp32 in mixed-precision pipelines.
Set jax_allow_f16_reductions explicitly if you rely on half reductions.

When it happens

Trigger: lax.reduce_sum(x16, axes) where x16 has dtype float16 or bfloat16 and any reduced axis is larger than 1, under default config. Often reached indirectly via jnp.sum on a half-precision array during tracing.

Common situations: Mixed-precision training pipelines that sum bf16 losses or gradients; TPU/GPU half-precision training; upgrading JAX where the flag default or enforcement changed. The flag defaults differently per backend/version.

Related errors

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