jax-ml/jax · error · TypeError

operand to reduce_window_sum must have a number dtype, got {

Error message

operand to reduce_window_sum must have a number dtype, got {}

What it means

reduce_window_sum (and lax.reduce_window with the sum reducer) requires a numeric dtype. The shape rule rejects operands with non-number dtypes such as bool or complex-incompatible types.

Source

Thrown at jax/_src/lax/windowed_reductions.py:522

      init_values=init_values,
      init_values_avals=init_value_avals,
      out_avals=ctx.avals_out,
      window_dimensions=window_dimensions,
      window_strides=window_strides,
      base_dilation=base_dilation,
      window_dilation=window_dilation,
      padding=padding,
  )


mlir.register_lowering(reduce_window_p, _generic_reduce_window_lower)


def _reduce_window_sum_shape_rule(operand, *, window_dimensions, window_strides,
                                  padding, base_dilation, window_dilation):
  if not dtypes.issubdtype(operand.dtype, np.number):
    msg = "operand to reduce_window_sum must have a number dtype, got {}"
    raise TypeError(msg.format(np.dtype(operand.dtype).name))
  return _common_reduce_window_shape_rule(operand, window_dimensions,
                                          window_strides, padding,
                                          base_dilation, window_dilation)

def _reduce_window_sum_transpose_rule(cotangent, operand, *, window_dimensions,
                                      window_strides, padding, base_dilation,
                                      window_dilation):
  assert ad.is_undefined_primal(operand)
  input_shape = operand.aval.shape
  pads = convolution._conv_general_vjp_lhs_padding(
      input_shape, window_dimensions, window_strides, cotangent.shape, padding,
      base_dilation, window_dilation)
  ones = [1] * len(input_shape)
  padding_config = [(lo, hi, stride - 1)
                    for (lo, hi), stride in zip(pads, window_strides)]
  pad_cotangent = lax.pad(cotangent, lax._zero(cotangent), padding_config)
  result = _reduce_window_sum(pad_cotangent, window_dimensions, base_dilation,
                              [(0, 0)] * len(input_shape),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the operand: x.astype(jnp.float32) or x.astype(jnp.int32)
2. For bool masks, use reduce_window_max or cast before summing

Example fix

# before
lax.reduce_window_sum(bool_mask, dims, strides, padding)
# after
lax.reduce_window_sum(bool_mask.astype(jnp.int32), dims, strides, padding)

assert jnp.issubdtype(x.dtype, jnp.number), f'requires number dtype, got {x.dtype}'

def is_number_dtype(x): return jnp.issubdtype(x.dtype, jnp.number)

When it happens

Trigger: Calling jax.lax.reduce_window_sum on a boolean array or an array with dtype like np.str_ / object.

Common situations: Summing a boolean mask over windows (e.g. counting valid pixels) without converting to an integer type first.

Related errors

• All objects to concatenate must be arrays, got {}.
• mulhi requires integer inputs, got {dtype}
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}

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