jax-ml/jax · error · ValueError

Input must be a rank-1 array.

Error message

Input must be a rank-1 array.

What it means

jnp.roots finds polynomial roots by building a companion matrix from the coefficient vector, which requires a rank-1 (1-D) coefficient array. After atleast_1d promotion, any input with ndim != 1 (a matrix, batched coefficients, or scalar plus a stray axis) is rejected.

Source

Thrown at jax/_src/numpy/polynomial.py:111

  Examples:
    >>> coeffs = jnp.array([0, 1, 2])

    The default behavior matches numpy and strips leading zeros:

    >>> jnp.roots(coeffs)
    Array([-2.+0.j], dtype=complex64)

    With ``strip_zeros=False``, extra roots are set to NaN:

    >>> jnp.roots(coeffs, strip_zeros=False)
    Array([-2. +0.j, nan+nanj], dtype=complex64)
  """
  p = ensure_arraylike("roots", p)
  p, = promote_dtypes_inexact(p)
  p_arr = atleast_1d(p)
  del p
  if p_arr.ndim != 1:
    raise ValueError("Input must be a rank-1 array.")
  if p_arr.size < 2:
    return array([], dtype=dtypes.to_complex_dtype(p_arr.dtype))
  num_leading_zeros = _where(all(p_arr == 0), len(p_arr), argmin(p_arr == 0))

  if strip_zeros:
    num_leading_zeros = core.concrete_or_error(int, num_leading_zeros,
      "The error occurred in the jnp.roots() function. To use this within a "
      "JIT-compiled context, pass strip_zeros=False, but be aware that leading zeros "
      "will result in some returned roots being set to NaN.")
    return _roots_no_zeros(p_arr[num_leading_zeros:])
  else:
    return _roots_with_zeros(p_arr, num_leading_zeros)


@export
@api.jit(static_argnames=('deg', 'rcond', 'full', 'cov'))
def polyfit(x: ArrayLike, y: ArrayLike, deg: int, rcond: float | None = None,
            full: bool = False, w: ArrayLike | None = None, cov: bool = False

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Flatten the coefficients: jnp.roots(p.ravel()) or squeeze extra axes.
2. Ensure you pass a plain 1-D coefficient vector, highest degree first.
3. For batches, vmap over individual coefficient vectors.

Example fix

// before
r = jnp.roots(p)  # p.shape == (3, 1)
// after
r = jnp.roots(p.ravel())

p = jnp.ravel(jnp.asarray(p))
assert p.ndim == 1, p.shape
r = jnp.roots(p)

def is_rank1(p) -> bool:
    return getattr(p, 'ndim', 0) == 1

When it happens

Trigger: jnp.roots(jnp.array([[1, 0, -1]])) (shape (1, 3)); passing batched coefficient arrays; passing a 2-D coefficient matrix from a fitting routine.

Common situations: Wrapping coefficients in an extra axis when loading from datasets; expecting vectorized/batched root-finding (jnp.roots does not batch); scalars becoming shape-(1,) is fine, but (n, 1) column vectors are not.

Related errors

• expected 1D vector for x
• expected 1D or 2D array for y
• expected a 1-d array for weights
• matrix_transpose requires at least 2 dimensions; got {ndim=}
• After moving axes to end, leading shape of a must match shap

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