{"record":{"id":"7766d20f9ff7fd9b","repo":"jax-ml/jax","slug":"last-2-dimensions-of-the-array-must-be-square","errorCode":null,"errorMessage":"Last 2 dimensions of the array must be square","messagePattern":"Last 2 dimensions of the array must be square","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/linalg.py","lineNumber":384,"sourceCode":"    ...   jnp.linalg.matrix_power(a, -2)\n    Array([[ 5.5 , -2.5 ],\n           [-3.75,  1.75]], dtype=float32)\n\n    Negative powers are equivalent to matmul of the inverse:\n\n    >>> inv_a = jnp.linalg.inv(a)\n    >>> with jnp.printoptions(precision=3):\n    ...   inv_a @ inv_a\n    Array([[ 5.5 , -2.5 ],\n           [-3.75,  1.75]], dtype=float32)\n  \"\"\"\n  arr = ensure_arraylike(\"jnp.linalg.matrix_power\", a)\n\n  if arr.ndim < 2:\n    raise TypeError(\"{}-dimensional array given. Array must be at least \"\n                    \"two-dimensional\".format(arr.ndim))\n  if arr.shape[-2] != arr.shape[-1]:\n    raise TypeError(\"Last 2 dimensions of the array must be square\")\n  try:\n    n = operator.index(n)\n  except TypeError as err:\n    raise TypeError(f\"exponent must be an integer, got {n}\") from err\n\n  if n == 0:\n    return jnp.broadcast_to(jnp.eye(arr.shape[-2], dtype=arr.dtype), arr.shape)\n  elif n < 0:\n    arr = inv(arr)\n    n = abs(n)\n\n  if n == 1:\n    return arr\n  elif n == 2:\n    return arr @ arr\n  elif n == 3:\n    return (arr @ arr) @ arr\n","sourceCodeStart":366,"sourceCodeEnd":402,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/linalg.py#L366-L402","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Fix the input so the last two dims match: shape (..., n, n)","Use @ for repeated multiplication or QR/SVD-based APIs if a rectangular matrix power was intended","Assert squareness before calling: assert a.shape[-2] == a.shape[-1]"],"exampleFix":"// before\njnp.linalg.matrix_power(jnp.ones((3, 4)), 2)  # TypeError\n// after\nA = jnp.ones((3, 3))\njnp.linalg.matrix_power(A, 2)","handlingStrategy":"validation","validationCode":"a = jnp.asarray(a)\nassert a.shape[-2] == a.shape[-1], 'matrix_power requires square matrices'\njnp.linalg.matrix_power(a, n)","typeGuard":"def is_square(x) -> bool:\n    x = jnp.asarray(x)\n    return x.ndim >= 2 and x.shape[-2] == x.shape[-1]","tryCatchPattern":null,"preventionTips":["Validate squareness before matrix_power","Watch shape drift after transposes/refactors","Rectangular matrices have no matrix power — fix the data or the operation"],"tags":["jax","linalg","matrix-power","square-matrix-required"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}