jax-ml/jax · error · ValueError

jax.numpy.arange: arguments must be scalars; got {name}={val

Error message

jax.numpy.arange: arguments must be scalars; got {name}={val}

What it means

Unlike NumPy, jnp.arange requires start, stop, and step to be scalar values (ndim == 0) because it must compute the output length statically for shape purposes. Passing any array with ndim > 0 raises this ValueError naming the offending argument and value.

Source

Thrown at jax/_src/numpy/lax_numpy.py:5956

def _arange(start: ArrayLike | DimSize, stop: ArrayLike | DimSize | None = None,
            step: ArrayLike | None = None, dtype: DTypeLike | None = None,
            out_sharding: NamedSharding | None = None) -> Array:
  # Validate inputs
  if dtype is not None:
    dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "arange")
  util.check_arraylike_or_none("arange", start, stop, step)

  # Ensure start/stop/step are concrete
  start_name = "stop" if stop is None and step is None else "start"
  start = core.concrete_or_error(None, start, f"It arose in the jnp.arange argument '{start_name}'")
  stop = core.concrete_or_error(None, stop, "It arose in the jnp.arange argument 'stop'")
  step = core.concrete_or_error(None, step, "It arose in the jnp.arange argument 'step'")

  # Ensure start/stop/step are scalars
  for name, val in [(start_name, start), ("stop", stop), ("step", step)]:
    if val is not None and np.ndim(val) != 0:
      raise ValueError(f"jax.numpy.arange: arguments must be scalars; got {name}={val}")

  # Handle symbolic dimensions
  if any(core.is_symbolic_dim(v) for v in (start, stop, step)):
    if stop is None:
      start, stop = 0, start
    if step is None:
      step = 1
    return _arange_dynamic(start, stop, step, dtype or dtypes.default_int_dtype())

  if dtype is None:
    dtype = dtypes.result_type(start, *(x for x in [stop, step] if x is not None))
  dtype = dtypes.jax_dtype(dtype)

  if iscomplexobj(start) or iscomplexobj(stop) or iscomplexobj(step):
    raise ValueError(
        "Passing complex start/stop/step to jnp.arange is no longer supported"
        " starting in JAX v0.10.0.")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Extract a scalar dimension: jnp.arange(x.shape[0])
2. Use int(...) conversion for Python/numpy scalars: jnp.arange(int(stop))
3. For per-batch ranges, use vmap over a scalarized version or jnp.repeat-based construction

Example fix

// before
idx = jnp.arange(x.shape)        # tuple/array of dims
// after
idx = jnp.arange(x.shape[0])     # scalar bound

for name, v in [('start', start), ('stop', stop), ('step', step)]:
    if v is not None and np.ndim(v) != 0:
        raise ValueError(f'{name} must be scalar')

def is_scalar(v) -> bool:
    return np.ndim(v) == 0

When it happens

Trigger: jnp.arange(jnp.array([0, 10])), jnp.arange(start=arr, stop=arr+5), or passing a shape tuple, e.g. jnp.arange(x.shape). Under jit, non-concrete values instead fail earlier in core.concrete_or_error.

Common situations: Using jnp.arange(x.shape[0]) works (scalar), but jnp.arange(x.shape) with a multi-dim shape fails; passing batched/traced bounds inside jitted code; iterating with array-valued endpoints.

Related errors

• In arange with non-constant arguments all of start, stop, an
• arange has step == 0
• Invalid value "{default}" for JAX flag {name}
• new enum value must be in {enum_values}, got {new_val} of ty
• new enum value must be None or in {enum_values}, got {new_va

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