jax-ml/jax · error · ValueError

reduce_precision: mantissa_bits must be non-negative; got {m

Error message

reduce_precision: mantissa_bits must be non-negative; got {mantissa_bits}

What it means

reduce_precision's mantissa_bits parameter must be >= 0. A negative mantissa means no significand bits, which describes no representable floating-point format, so the shape rule rejects it.

Source

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

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'),
    name='reduce_precision', sharding_rule=_reduce_precision_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'reduce_precision'),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Guard the sweep: only call reduce_precision when mantissa_bits >= 0 and exponent_bits >= 1.
  2. For FP8 E5M2 simulation use exponent_bits=5, mantissa_bits=2.
  3. Fix the bit-allocation arithmetic that produced a negative mantissa.

Example fix

# before
y = lax.reduce_precision(x, 5, -1)
# after
y = lax.reduce_precision(x, 5, 2)
Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

def valid_fp_format(e, m):
    return e >= 1 and m >= 0

Prevention

When it happens

Trigger: lax.reduce_precision(x, exponent_bits=5, mantissa_bits=-1); typically from computed bit budgets in quantization experiments.

Common situations: Sweeps over (exponent_bits, mantissa_bits) that include negative mantissa values; off-by-one errors in format-conversion code (e.g., FP8 E4M3 handled as mantissa=3 but E5M2 mishandled).

Related errors


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