jax-ml/jax · error · ValueError

input must be 1d or non-empty square 2d array.

Error message

input must be 1d or non-empty square 2d array.

What it means

Thrown by jnp.poly when the input is neither a 1-D sequence of polynomial roots nor a non-empty square 2-D matrix. A square 2-D matrix is converted to eigenvalues (companion-matrix style, like numpy's poly), so any other shape (e.g. rectangular 2-D, empty 2-D, 3-D) is invalid.

Source

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

    >>> x = jnp.array([[2, 1, 5],
    ...                [3, 4, 7],
    ...                [1, 3, 5]])
    >>> jnp.round(jnp.poly(x))
    Array([  1.+0.j, -11.-0.j,   9.+0.j, -15.+0.j], dtype=complex64)
  """
  seq_of_zeros = ensure_arraylike('poly', seq_of_zeros)
  seq_of_zeros, = promote_dtypes_inexact(seq_of_zeros)
  seq_of_zeros_arr = atleast_1d(seq_of_zeros)
  del seq_of_zeros

  sh = seq_of_zeros_arr.shape
  if len(sh) == 2 and sh[0] == sh[1] and sh[0] != 0:
    # import at runtime to avoid circular import
    from jax._src.numpy import linalg
    seq_of_zeros_arr = linalg.eigvals(seq_of_zeros_arr)

  if seq_of_zeros_arr.ndim != 1:
    raise ValueError("input must be 1d or non-empty square 2d array.")

  dt = seq_of_zeros_arr.dtype
  if len(seq_of_zeros_arr) == 0:
    return ones((), dtype=dt)

  a = ones((1,), dtype=dt)
  for k in range(len(seq_of_zeros_arr)):
    a = convolve(a, array([1, -seq_of_zeros_arr[k]], dtype=dt), mode='full')

  return a


@export
@api.jit(static_argnames=['unroll'])
def polyval(p: ArrayLike, x: ArrayLike, *, unroll: int = 16) -> Array:
  r"""Evaluates the polynomial at specific values.

  JAX implementations of :func:`numpy.polyval`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a 1-D array of roots: jnp.poly(roots_1d)
2. If you meant to evaluate a polynomial from coefficients, use jnp.polyval(coeffs, x) instead
3. For matrices, ensure the array is square and non-empty (eigenvalues are computed)

Example fix

// before
jnp.poly(coeffs_2d_rectangular)
// after
jnp.poly(jnp.ravel(roots))          # 1-D roots
# or, if input is really coefficients:
jnp.polyval(coeffs, x)

def as_roots(a):
    import jax.numpy as jnp
    a = jnp.asarray(a)
    assert a.ndim == 1 or (a.ndim == 2 and a.shape[0] == a.shape[1] and a.shape[0] > 0), \
        f'poly expects 1-D roots or square matrix, got {a.shape}'
    return a

def is_valid_poly_input(a) -> bool:
    import jax.numpy as jnp
    a = jnp.asarray(a)
    return a.ndim == 1 or (a.ndim == 2 and a.shape[0] == a.shape[1] > 0)

When it happens

Trigger: Calling jnp.poly on a 2-D array that is rectangular (rows != cols) or 0x0, or on a 3-D array; passing a scalar-shaped array also fails the ndim==1 check.

Common situations: Passing a coefficient array by mistake (poly expects roots, not coefficients — use polyval for coefficients); passing a batch of root vectors as a 2-D array expecting vectorized behavior; empty input edge case.

Related errors

• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32
• The names should be exclusive and should not intersect in `n
• The return value of the policies should be a boolean. Got: {

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