jax-ml/jax · error · ValueError

jnp.linalg.cond: for {p=}, array must be square; got {arr.sh

Error message

jnp.linalg.cond: for {p=}, array must be square; got {arr.shape=}

What it means

Condition numbers for p norms other than 2/-2 (e.g. p=1 or p=jnp.inf) are computed as norm(x) * norm(inv(x)), which requires x to be square (invertible in the first place). jnp.linalg.cond therefore rejects non-square matrices when such a p is requested.

Source

Thrown at jax/_src/numpy/linalg.py:2302

    >>> x = jnp.array([[1, 2],
    ...                [0, 0]])
    >>> jnp.linalg.cond(x)
    Array(inf, dtype=float32)
  """
  arr = ensure_arraylike("cond", x)
  if arr.ndim < 2:
    raise ValueError(f"jnp.linalg.cond: input array must be at least 2D; got {arr.shape=}")
  if arr.shape[-1] == 0 or arr.shape[-2] == 0:
    raise ValueError(f"jnp.linalg.cond: input array must not be empty; got {arr.shape=}")
  if p is None or p == 2:
    s = svdvals(x)
    return s[..., 0] / s[..., -1]
  elif p == -2:
    s = svdvals(x)
    r = s[..., -1] / s[..., 0]
  else:
    if arr.shape[-2] != arr.shape[-1]:
      raise ValueError(f"jnp.linalg.cond: for {p=}, array must be square; got {arr.shape=}")
    r = norm(x, ord=p, axis=(-2, -1)) * norm(inv(x), ord=p, axis=(-2, -1))
  # Convert NaNs to infs where original array has no NaNs.
  return jnp.where(ufuncs.isnan(r) & ~ufuncs.isnan(x).any(axis=(-2, -1)), np.inf, r)


@export
def trace(x: ArrayLike, /, *,
          offset: int = 0, dtype: DTypeLike | None = None) -> Array:
  """Compute the trace of a matrix.

  JAX implementation of :func:`numpy.linalg.trace`.

  Args:
    x: array of shape ``(..., M, N)`` and whose innermost two
      dimensions form MxN matrices for which to take the trace.
    offset: positive or negative offset from the main diagonal
      (default: 0).
    dtype: data type of the returned array (default: ``None``). If ``None``,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use p=None (default, 2-norm via SVD), which supports rectangular matrices (ratio of extreme singular values).
2. Square the matrix first if a square normal-equations view is acceptable (A^T A).
3. Pass the original square operator if a rectangular one was produced by mistake.

Example fix

// before
c = jnp.linalg.cond(A, ord=1)  # A: (m, n), m != n
// after
c = jnp.linalg.cond(A)  # 2-norm condition via SVD, works for rectangular

if p not in (None, 2, -2) and a.shape[-2] != a.shape[-1]:
    p = None  # fall back to SVD-based 2-norm cond
c = jnp.linalg.cond(a, p)

def cond_p_needs_square(p) -> bool:
    return p is not None and p not in (2, -2)

When it happens

Trigger: jnp.linalg.cond(x, ord=1) or ord=jnp.inf with x of shape (m, n), m != n; rectangular inputs from least-squares problems.

Common situations: Analyzing conditioning of a design matrix in a regression pipeline with ord=1/inf; defaulting p to 1 in code reused from square-matrix contexts.

Related errors

• jnp.linalg.cond: input array must be at least 2D; got {arr.s
• matrix_transpose requires at least 2 dimensions; got {ndim=}
• After moving axes to end, leading shape of a must match shap
• Input arrays must have prod(a.shape[:b.ndim]) == prod(a.shap
• multi_dot: input arrays must all be two-dimensional, except

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