jax-ml/jax · error · TypeError

Last 2 dimensions of the array must be square

Error message

Last 2 dimensions of the array must be square

What it means

jnp.linalg.matrix_power requires the last two dimensions of the input to be equal (a square matrix stack), since matrix powers are only defined for square matrices. If arr.shape[-2] != arr.shape[-1], TypeError('Last 2 dimensions of the array must be square') is raised.

Source

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

    ...   jnp.linalg.matrix_power(a, -2)
    Array([[ 5.5 , -2.5 ],
           [-3.75,  1.75]], dtype=float32)

    Negative powers are equivalent to matmul of the inverse:

    >>> inv_a = jnp.linalg.inv(a)
    >>> with jnp.printoptions(precision=3):
    ...   inv_a @ inv_a
    Array([[ 5.5 , -2.5 ],
           [-3.75,  1.75]], dtype=float32)
  """
  arr = ensure_arraylike("jnp.linalg.matrix_power", a)

  if arr.ndim < 2:
    raise TypeError("{}-dimensional array given. Array must be at least "
                    "two-dimensional".format(arr.ndim))
  if arr.shape[-2] != arr.shape[-1]:
    raise TypeError("Last 2 dimensions of the array must be square")
  try:
    n = operator.index(n)
  except TypeError as err:
    raise TypeError(f"exponent must be an integer, got {n}") from err

  if n == 0:
    return jnp.broadcast_to(jnp.eye(arr.shape[-2], dtype=arr.dtype), arr.shape)
  elif n < 0:
    arr = inv(arr)
    n = abs(n)

  if n == 1:
    return arr
  elif n == 2:
    return arr @ arr
  elif n == 3:
    return (arr @ arr) @ arr

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Fix the input so the last two dims match: shape (..., n, n)
2. Use @ for repeated multiplication or QR/SVD-based APIs if a rectangular matrix power was intended
3. Assert squareness before calling: assert a.shape[-2] == a.shape[-1]

Example fix

// before
jnp.linalg.matrix_power(jnp.ones((3, 4)), 2)  # TypeError
// after
A = jnp.ones((3, 3))
jnp.linalg.matrix_power(A, 2)

a = jnp.asarray(a)
assert a.shape[-2] == a.shape[-1], 'matrix_power requires square matrices'
jnp.linalg.matrix_power(a, n)

def is_square(x) -> bool:
    x = jnp.asarray(x)
    return x.ndim >= 2 and x.shape[-2] == x.shape[-1]

When it happens

Trigger: jnp.linalg.matrix_power(jnp.ones((3, 4)), 2) — a non-square (3, 4) matrix; also triggered in tests like testMatrixPowerBool with rectangular bool input.

Common situations: Feeding weight matrices whose shapes drifted during refactoring; applying matrix powers to attention/transition matrices after a transpose bug; rectangular data matrices mistaken for operators.

Related errors

• {}-dimensional array given. Array must be at least two-dimen
• Unsupported dtype: {dtype}
• Argument to symmetric eigendecomposition must have shape [..
• Argument to Hessenberg reduction must have shape [..., n, n]
• The first argument to householder_product must have at least

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