jax-ml/jax · error · ValueError

reduce_precision: exponent_bits must be positive; got {expon

Error message

reduce_precision: exponent_bits must be positive; got {exponent_bits}

What it means

reduce_precision's exponent_bits parameter must be >= 1 after operator.index() conversion. exponent_bits controls the target float format's exponent field; zero or negative values describe no representable format.

Source

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

mlir.register_lowering(reduce_prod_p, partial(_unary_reduce_lower, hlo.MulOp,
                                              _get_prod_identity))
mlir.register_lowering(reduce_or_p, partial(_unary_reduce_lower, hlo.OrOp,
                                            _get_bitwise_or_identity))
mlir.register_lowering(reduce_and_p, partial(_unary_reduce_lower, hlo.AndOp,
                                             _get_bitwise_and_identity))
mlir.register_lowering(reduce_xor_p, partial(_unary_reduce_lower, hlo.XorOp,
                                             _get_bitwise_or_identity))
mlir.register_lowering(reduce_min_p, partial(_unary_reduce_lower, mlir.min_hlo,
                                             _get_min_identity))
mlir.register_lowering(reduce_max_p, partial(_unary_reduce_lower, mlir.max_hlo,
                                             _get_max_identity))


def _reduce_precision_shape_rule(operand, *, exponent_bits, mantissa_bits):
  exponent_bits = operator.index(exponent_bits)
  mantissa_bits = operator.index(mantissa_bits)
  if exponent_bits < 1:
    raise ValueError(f"reduce_precision: exponent_bits must be positive; got {exponent_bits}")
  if mantissa_bits < 0:
    raise ValueError(f"reduce_precision: mantissa_bits must be non-negative; got {mantissa_bits}")
  return operand.shape

def _reduce_precision_sharding_rule(operand, *, exponent_bits, mantissa_bits):
  return operand.sharding

def _reduce_precision_memory_space_rule(operand, *, exponent_bits, mantissa_bits):
  return operand.memory_space

def _reduce_precision_ur_rule(operand, *, exponent_bits, mantissa_bits):
  out_unreduced = core.getu(operand)
  kind = UnreducedKind.sum if out_unreduced else None
  return out_unreduced, core.getr(operand), kind

reduce_precision_p = standard_primitive(
    _reduce_precision_shape_rule,
    partial(unop_dtype_rule, _identity, _float, 'reduce_precision'),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Validate exponent_bits >= 1 before calling; clamp or skip invalid configurations.
2. For bf16 simulation use exponent_bits=8, mantissa_bits=7.
3. Check the arithmetic that splits a total bit budget between exponent and mantissa (reserve at least 1 exponent bit).

Example fix

# before
y = lax.reduce_precision(x, exponent_bits=0, mantissa_bits=7)
# after
y = lax.reduce_precision(x, exponent_bits=8, mantissa_bits=7)

assert exponent_bits >= 1, exponent_bits
y = lax.reduce_precision(x, exponent_bits, mantissa_bits)

def valid_reduce_precision_params(e, m):
    return operator.index(e) >= 1 and operator.index(m) >= 0

When it happens

Trigger: lax.reduce_precision(x, exponent_bits=0, mantissa_bits=10), or passing a float/negative exponent_bits that indexes to a non-positive int. Fires in the shape rule during tracing.

Common situations: Simulating low-precision formats (bf16/fp8-style) with computed bit budgets where the exponent allocation math yields 0; parameter sweeps that include invalid corners.

Related errors

• reduce_precision: mantissa_bits must be non-negative; got {m
• {full_name} must be a pytree prefix with bool leaves or a tu
• Axes mentioned in `manual_axis_type` field of ShapedArray sh
• varying and unreduced cannot have common mesh axes. Got vary
• varying and reduced cannot have common mesh axes. Got varyin

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