jax-ml/jax · error · ValueError

In arange with non-constant arguments all of start, stop, an

Error message

In arange with non-constant arguments all of start, stop, and step must be either dimension expressions or integers: start={start}, stop={stop}, step={step}

What it means

When jnp.arange falls back to the dynamic (non-constant argument) path, all of start, stop, and step must be either static integers or symbolic dimension expressions (polynomials over dimension variables). Mixing a dimension expression with a non-dimension value (e.g. a traced float) raises this ValueError.

Source

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

                lax.mul(lax.convert_element_type(step, working_dtype),
                        lax.broadcasted_iota(working_dtype, (size,), 0,
                                             out_sharding=out_sharding))),
        dtype)
  elif start == 0:
    # arange(M) or arange(0, M)
    size = max(0, int(np.ceil(stop)))
    return lax.broadcasted_iota(dtype, (size,), 0, out_sharding=out_sharding)
  else:
    # arange(N, M)
    size = max(0, int(np.ceil(stop - start)))
    return lax.add(lax.convert_element_type(start, dtype),
                    lax.broadcasted_iota(dtype, (size,), 0, out_sharding=out_sharding))

def _arange_dynamic(
    start: DimSize, stop: DimSize, step: DimSize, dtype: DTypeLike) -> Array:
  # Here if at least one of start, stop, step are dynamic.
  if any(not core.is_dim(v) for v in (start, stop, step)):
    raise ValueError(
        "In arange with non-constant arguments all of start, stop, and step "
        f"must be either dimension expressions or integers: start={start}, "
        f"stop={stop}, step={step}")
  # Must resolve statically if step is {<0, ==0, >0}
  try:
    if step == 0:
      raise ValueError("arange has step == 0")
    step_gt_0 = (step > 0)
  except core.InconclusiveDimensionOperation as e:
    raise core.InconclusiveDimensionOperation(
        f"In arange with non-constant arguments the step ({step}) must " +
        f"be resolved statically if it is > 0 or < 0.\nDetails: {e}")
  gap = step if step_gt_0 else - step
  distance = (stop - start) if step_gt_0 else (start - stop)
  size = core.max_dim(0, distance + gap - 1) // gap
  return (array(start, dtype=dtype) +
          array(step, dtype=dtype) * lax.iota(dtype, size))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Keep all three arguments as integers/dimension expressions when any is symbolic: use step=1 (or int step)
2. Re-express float-step ranges via linspace on an integer range, or scale after: jnp.arange(n) * 0.1
3. Mark the bound static (e.g. pass a Python int) if the size is actually known

Example fix

// before
xs = jnp.arange(n, step=0.5)  # n symbolic dim
// after
xs = jnp.arange(n) * 0.5

from jax._src import core
if any(core.is_symbolic_dim(v) for v in (start, stop, step)):
    assert all(core.is_dim(v) for v in (start, stop, step)), 'mixed symbolic/non-symbolic arange args'

When it happens

Trigger: Calling jnp.arange with a symbolic dimension bound (from jax.export or shape polymorphism) mixed with a non-dim start/step, e.g. jnp.arange(n, step=0.5) where n is a symbolic dimension; or arange(dim, 2*dim, 0.1).

Common situations: Shape-polymorphic jax.jit/export code that iterates a range over a batch dimension with a float step; migrating dynamic-shape code where a tracer leaks into arange bounds.

Related errors

• jax.numpy.arange: arguments must be scalars; got {name}={val
• 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/07a968d6d41f1858. Report an issue: GitHub.