jax-ml/jax · error · ValueError

{start=} must satisfy {-a_ndim}<=start<={a_ndim}

Error message

{start=} must satisfy {-a_ndim}<=start<={a_ndim}

What it means

Raised by jnp.rollaxis when the start parameter is outside [-ndim, ndim] of the input array; start must be a valid position to move the axis to, after canonicalization of axis itself.

Source

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

    Roll axis 1 to the end of the array:

    >>> jnp.rollaxis(a, 1, a.ndim).shape
    (2, 4, 5, 3)

    Equivalent of these two with :func:`~jax.numpy.moveaxis`

    >>> jnp.moveaxis(a, 2, 0).shape
    (4, 2, 3, 5)
    >>> jnp.moveaxis(a, 1, -1).shape
    (2, 4, 5, 3)
  """
  a = util.ensure_arraylike("rollaxis", a)
  start = core.concrete_or_error(operator.index, start, "'start' argument of jnp.rollaxis()")
  a_ndim = np.ndim(a)
  axis = _canonicalize_axis(axis, a_ndim)
  if not (-a_ndim <= start <= a_ndim):
    raise ValueError(f"{start=} must satisfy {-a_ndim}<=start<={a_ndim}")
  if start < 0:
    start += a_ndim
  if start > axis:
    start -= 1
  return moveaxis(a, axis, start)


@export
@api.jit(static_argnames=('axis', 'bitorder'))
def packbits(a: ArrayLike, axis: int | None = None, bitorder: str = "big") -> Array:
  """Pack array of bits into a uint8 array.

  JAX implementation of :func:`numpy.packbits`

  Args:
    a: N-dimensional array of bits to pack.
    axis: optional axis along which to pack bits. If not specified, ``a`` will
      be flattened.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clamp or validate start within [-a.ndim, a.ndim] before calling
2. Compute start adaptively from a.ndim rather than hardcoding
3. Prefer jnp.moveaxis(a, source, destination) which has clearer semantics and its own validation

Example fix

// before
jnp.rollaxis(a, axis=1, start=4)  # a.ndim == 3
// after
jnp.moveaxis(a, source=1, destination=2)  # explicit, validated destination

start = int(start)
n = jnp.ndim(a)
if not (-n <= start <= n): raise ValueError(f'start {start} out of range for ndim {n}')

def valid_rollaxis_start(a, start):
    n = jnp.ndim(a)
    return -n <= int(start) <= n

When it happens

Trigger: jnp.rollaxis(a, axis=0, start=5) on a 3-D array; start derived from user config or loop arithmetic exceeding the rank; start passed as a traced value is also rejected earlier by concrete_or_error.

Common situations: Hardcoded start values breaking after arrays are reshaped to lower rank; looping axis reordering with unclamped counters; mixing up rollaxis(start) vs moveaxis destination conventions.

Related errors

• `compute_on`'s compute_type argument must be a string.
• Argument '{x}' of type '{typ}' is not a valid JAX type
• Argument '{arg}' of type {type(arg)} is not a valid JAX type
• Expected kind to be a dtype, string, or tuple; got {kind=}
• at least one array or dtype is required

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