jax-ml/jax · error · TypeError

reduce_window expected init_values to be scalars but init_va

Error message

reduce_window expected init_values to be scalars but init_values have shapes {}.

What it means

Raised by JAX's reduce_window abstract evaluation when the init_value arguments passed to lax.reduce_window are not scalars (rank != 0). init_value acts as the reduction's identity element (e.g. -inf for max) and must be a 0-d array or Python scalar.

Source

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

    consts,
    window_dimensions,
    window_strides,
    padding,
    base_dilation,
    window_dilation,
):
  operand_avals, init_val_avals = util.split_list(avals, [len(avals) // 2])
  if any(
      o.dtype != iv.dtype for o, iv in zip(operand_avals, init_val_avals)
  ):
    msg = ("reduce_window got inconsistent dtypes for operands and init_values:"
           " got operand dtypes {} and init_value dtypes {}.")
    raise TypeError(msg.format([o.dtype for o in operand_avals],
                               [iv.dtype for iv in init_val_avals]))
  if any(len(v.shape) != 0 for v in init_val_avals):
    msg = ("reduce_window expected init_values to be scalars but init_values "
           "have shapes {}.")
    raise TypeError(msg.format([v.shape for v in init_val_avals]))
  out_shape = _common_reduce_window_shape_rule(
      operand_avals[0], window_dimensions, window_strides, padding,
      base_dilation, window_dilation)
  out_sharding = reduce_window_sharding_rule(
      operand_avals[0], window_dimensions, window_strides, padding,
      base_dilation, window_dilation)
  vma = core.standard_vma_rule('reduce_window', *operand_avals)
  if any(core.getu(a) or core.getr(a) for a in operand_avals):
    raise NotImplementedError
  return tuple(ShapedArray(out_shape, op.dtype, sharding=out_sharding,
                           manual_axis_type=op.mat.update(varying=vma))
               for op in operand_avals)


def _generic_reduce_window_batch_rule(
    batched_args,
    batch_dims,
    *,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass a Python scalar or 0-d value: -np.inf, 0, jnp.asarray(-np.inf) with shape ()
If init_value came from an array, index it: init_value[0] or use .item()

Example fix

// before
init = jnp.array([-np.inf])
lax.reduce_window(x, init, jaxpr, dims, strides, padding)
// after
init = jnp.asarray(-np.inf)  # scalar, shape ()
lax.reduce_window(x, init, jaxpr, dims, strides, padding)

Defensive patterns

Strategy: validation

Validate before calling

import numpy as np
assert np.ndim(init_value) == 0, f'init_value must be scalar, got shape {np.shape(init_value)}'

Type guard

def is_scalar(x): return np.ndim(x) == 0

Prevention

Always build init_value from Python scalars (-inf, 0) or jnp.asarray(scalar)
Assert ndim==0 in pooling wrappers

When it happens

Trigger: Calling lax.reduce_window(operand, init_value, jaxpr, consts, ...) where init_value is a list, array of shape (1,), or any non-scalar (e.g. np.full((1,), -np.inf)).

Common situations: Porting NumPy pooling code where identity values were kept in arrays; building generic reduce wrappers that pass vector identities.

Related errors

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