jax-ml/jax · error · ValueError

The length of `a` along the last axis must be at least 2; go

Error message

The length of `a` along the last axis must be at least 2; got shape {a.shape}.

What it means

jax.scipy.linalg.companion returns the companion matrix of a polynomial with at least degree 1, so the coefficient array a needs length >= 2 along its last axis (leading coefficient first). After atleast_1d promotion, scalars and length-1 arrays raise ValueError.

Source

Thrown at jax/_src/scipy/linalg.py:2692

  Returns:
    A companion matrix of shape ``(..., N - 1, N - 1)``.

  Note:
    Unlike :func:`scipy.linalg.companion`, this function does not check at
    runtime that ``a[..., 0]`` is non-zero; if the leading coefficient is
    zero, the result will contain ``inf`` or ``nan`` entries.

  Examples:
    >>> jax.scipy.linalg.companion(jnp.array([1., -10., 31., -30.]))
    Array([[ 10., -31.,  30.],
           [  1.,   0.,   0.],
           [  0.,   1.,   0.]], dtype=float32)
  """
  a, = promote_args_inexact("companion", a)
  a = jnp.atleast_1d(a)
  if a.shape[-1] < 2:
    raise ValueError(
        "The length of `a` along the last axis must be at least 2; "
        f"got shape {a.shape}.")
  return _companion(a)

@partial(jnp_vectorize.vectorize, signature="(n)->(m,m)")
def _companion(a: Array) -> Array:
  first_row = -a[1:] / a[0]
  m = a.shape[0] - 1
  out = jnp.eye(m, m, k=-1, dtype=first_row.dtype)
  return out.at[0].set(first_row)


def fiedler(a: ArrayLike) -> Array:
  r"""Construct a symmetric Fiedler matrix.

  JAX implementation of :func:`scipy.linalg.fiedler`.

  The Fiedler matrix has entries :math:`F_{ij} = |a_i - a_j|` for

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Supply at least two coefficients, e.g. companion([1, 2, 3]) for x^2 + 2x + 3
2. Guard degree-0/constant polynomials in caller code (they have no companion matrix)
3. Check a.shape[-1] >= 2 before calling

Example fix

# before
C = companion(jnp.array([5.0]))
# after
C = companion(jnp.array([5.0, 1.0]))  # 5 + x

a = jnp.atleast_1d(a)
if a.shape[-1] < 2:
    raise ValueError('polynomial must have degree >= 1 (>= 2 coefficients)')
C = companion(a)

try:
    companion(a)
except ValueError as e:
    if 'at least 2' in str(e):
        a = jnp.concatenate([a, jnp.zeros_like(a)])  # only if sensible
        companion(a)
    else: raise

When it happens

Trigger: Calling companion(jnp.array([1.0])) or companion(2.0); batched input whose last axis was reduced to size 1.

Common situations: Passing a constant instead of polynomial coefficients; trimming roots/coefficient lists one element too many; building coefficients dynamically where a degree-0 case slips through.

Related errors

• numpy masked arrays are not supported as direct inputs to JA
• Invalid compute type {c_type}. Current supported values are
• numpy masked arrays are not supported as direct inputs to JA
• key cannot be empty
• Mesh axis names cannot be None. Got: {axis_names}

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