jax-ml/jax · error · TypeError

logical reduction requires operand dtype bool or int, got {o

Error message

logical reduction requires operand dtype bool or int, got {operand.dtype}.

What it means

Logical reductions (reduce_and, reduce_or, reduce_xor) require the operand to have bool or integer dtype; float or complex operands are rejected because bitwise logical ops are undefined for them.

Source

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

    mlir.lower_fun(
        partial(_compute_argminmax, lt, _get_min_identity),
        multiple_results=False,
    ),
    inline=False,
)

mlir.register_lowering(
    argmax_p,
    mlir.lower_fun(
        partial(_compute_argminmax, gt, _get_max_identity),
        multiple_results=False,
    ),
    inline=False,
)

def _reduce_logical_shape_rule(operand, *, axes):
  if operand.dtype != np.bool_ and not np.issubdtype(operand.dtype, np.integer):
    raise TypeError(f"logical reduction requires operand dtype bool or int, got {operand.dtype}.")
  return tuple(np.delete(operand.shape, axes))

def _reduce_logical_sharding_rule(operand, *, axes):
  return operand.sharding.update(spec=tuple_delete(operand.sharding.spec, axes))

def _reduce_or_lin(_is_vjp, nzs, x, *, axes):
  nz, = nzs
  y = reduce_or_p.bind(x, axes=axes)
  aval = typeof(y).to_tangent_aval()
  return y, False, (), None, lambda _, __, t: ad_util.Zero(aval)

reduce_or_p = standard_primitive(
    _reduce_logical_shape_rule, input_dtype, 'reduce_or',
    weak_type_rule=_strip_weak_type, sharding_rule=_reduce_logical_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'reduce_or'))
batching.defreducer(reduce_or_p)
ad.primitive_linearizations[reduce_or_p] = _reduce_or_lin

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the operand to bool before the reduction: lax.reduce_or(x != 0, axes) or x.astype(jnp.bool_).
2. Fix upstream so the mask stays boolean instead of being promoted to float.
3. Use jnp.all / jnp.any for logical semantics on any dtype-comparable input.

Example fix

# before
out = lax.reduce_or(mask * 1.0, (0,))
# after
out = lax.reduce_or((mask * 1.0).astype(jnp.bool_), (0,))

if x.dtype != jnp.bool_:
    x = x != 0
out = lax.reduce_or(x, (0,))

def bool_or_int_dtype(x):
    import numpy as np
    return x.dtype == np.bool_ or np.issubdtype(x.dtype, np.integer)

When it happens

Trigger: lax.reduce_and(x_float, axes), jnp.logical_and-style reductions applied via lax on a float array, or calling the internal _reduce_or_lin transpose path with non-integer input.

Common situations: Applying all()/any()-style reductions implemented with bitwise ops to boolean masks that were accidentally cast to float (e.g., after arithmetic like mask * 1.0); data-type drift in a preprocessing pipeline.

Related errors

• {} does not accept dtype {}. Accepted dtypes are subtypes of
• logaddexp2 requires floating-point or complex inputs; got {x
• index arguments to dynamic_update_slice must be integers of
• indices must have an integer type
• primal and tangent arguments to jax.jvp do not match; dtypes

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