jax-ml/jax · error · ValueError

initial value must be a scalar. Got array of shape {initial_

Error message

initial value must be a scalar. Got array of shape {initial_arr.shape}

What it means

The initial value passed to a JAX reduction must be a scalar; arrays of any non-empty shape are rejected because there is no defined broadcasting of a per-element initial into the reduction.

Source

Thrown at jax/_src/numpy/reductions.py:160

  else:
    computation_dtype = result_dtype
  a = lax.convert_element_type(a, computation_dtype)
  op = op if computation_dtype != np.bool_ else bool_op
  # NB: in XLA, init_val must be an identity for the op, so the user-specified
  # initial value must be applied afterward.
  init_val = _reduction_init_val(a, init_val)
  if where_ is not None:
    a = _where(where_, a, init_val)
  if pos_dims is not dims:
    if parallel_reduce is None:
      raise NotImplementedError(f"Named reductions not implemented for jnp.{name}()")
    result = parallel_reduce(a, dims)
  else:
    result = lax.reduce(a, init_val, op, dims)
  if initial is not None:
    initial_arr = lax.convert_element_type(initial, lax.asarray(a).dtype)
    if initial_arr.shape != ():
      raise ValueError("initial value must be a scalar. "
                       f"Got array of shape {initial_arr.shape}")
    result = op(initial_arr, result)
  if keepdims:
    result = lax.expand_dims(result, pos_dims)
  return lax.convert_element_type(result, dtype or result_dtype)

def _canonicalize_axis_allow_named(x, rank):
  return maybe_named_axis(x, lambda i: canonicalize_axis(i, rank), lambda name: name)

def _reduction_dims(a: ArrayLike, axis: Axis):
  if axis is None:
    return (tuple(range(np.ndim(a))),) * 2
  if not isinstance(axis, (np.ndarray, tuple, list)):
    axes = (axis,)
  else:
    axes = axis
  canon_axis = tuple(_canonicalize_axis_allow_named(x, np.ndim(a))
                     for x in axes)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a Python scalar or 0-d array: initial=float(v) or jnp.asarray(v).reshape(())
2. Compute per-axis initials as separate reduction calls if per-axis values are needed

Example fix

// before
jnp.max(x, initial=jnp.array([0.0]))
// after
jnp.max(x, initial=0.0)  # or jnp.asarray(0.0)

import jax.numpy as jnp
initial = jnp.asarray(initial)
if initial.ndim != 0:
    initial = initial.reshape(())  # or take .item()
jnp.max(x, initial=initial)

When it happens

Trigger: jnp.max(x, initial=jnp.array([0, 0])) or passing a (1,)-shaped array (which is not treated as a scalar in JAX, unlike some numpy cases); passing a per-axis vector of initials.

Common situations: Reusing a broadcastable numpy pattern where initial had shape (1,); building initial from config values that end up as arrays (jnp.asarray of a list).

Related errors

• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac
• Incorrect output shape for return value #{i}: Expected: {out
• Length of sharding.spec ({len(out_s.spec)}) must be equal to

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