{"record":{"id":"b06cefc435f5ba43","repo":"jax-ml/jax","slug":"argument-to-hessenberg-reduction-must-have-shape","errorCode":null,"errorMessage":"Argument to Hessenberg reduction must have shape [..., n, n], got shape {shape}","messagePattern":"Argument to Hessenberg reduction must have shape \\[\\.\\.\\., n, n\\], got shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/linalg.py","lineNumber":1381,"sourceCode":"                precision=lax.Precision.HIGHEST)\n  vdag_adot_v = dot(dot(_H(v), a_dot), v)\n  dv = dot(v, Fmat * vdag_adot_v)\n  dw = _extract_diagonal(vdag_adot_v.real)\n  return (v, w_real), (dv, dw)\n\n\neigh_p = linalg_primitive(\n    _eigh_dtype_rule, (_float | _complex,), (2,), _eigh_shape_rule, \"eigh\",\n    multiple_results=True)\nad.primitive_jvps[eigh_p] = _eigh_jvp_rule\nregister_cpu_gpu_lowering(eigh_p, _eigh_cpu_gpu_lowering)\n\n\n# Hessenberg reduction\n\ndef _hessenberg_shape_rule(shape, **_):\n  if shape[0] != shape[-1]:\n    raise ValueError(\n        \"Argument to Hessenberg reduction must have shape [..., n, n], \"\n        f\"got shape {shape}\"\n    )\n  return shape, shape[:-2] + (shape[-1] - 1,)\n\n\ndef _hessenberg_dtype_rule(dtype, **_):\n  return dtype, dtype\n\n\ndef _hessenberg_cpu_lowering(ctx, a):\n  a_aval, = ctx.avals_in\n  batch_dims = a_aval.shape[:-2]\n  n = a_aval.shape[-1]\n  if not core.is_constant_dim(n):\n    raise ValueError(\"hessenberg requires the last dimension of a to be \"\n                     f\"constant, got a.shape of {a.shape}.\")\n  target_name = lapack.prepare_lapack_call(\"gehrd_ffi\", a_aval.dtype)","sourceCodeStart":1363,"sourceCodeEnd":1399,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/linalg.py#L1363-L1399","documentation":"jax/_src/lax/linalg.py:1381 in _hessenberg_shape_rule. jax.lax.linalg.hessenberg reduces a square matrix to upper Hessenberg form; the shape rule requires the last two dims equal ([..., n, n]). A rectangular operand raises ValueError with the observed shape.","triggerScenarios":"Calling jax.lax.linalg.hessenberg on an array whose last two dimensions differ, e.g. (m, k) with m != k, or a batch of non-square matrices.","commonSituations":"Applying Hessenberg reduction as a preprocessing step for eigensolvers on data matrices instead of square operators; transposition/slicing bugs producing off-by-one shapes.","solutions":["Ensure the operand is square: assert a.shape[-2] == a.shape[-1]","If you meant to tridiagonalize/bidiagonalize a rectangular matrix, use qr or svd-based preprocessing instead","Fix upstream shape construction (e.g. use A @ A.T or correct slicing)"],"exampleFix":"// before\nh, q = jax.lax.linalg.hessenberg(A)  # A: (m, k), m != k\n// after\nassert A.shape[-2] == A.shape[-1]\nh, q = jax.lax.linalg.hessenberg(A)","handlingStrategy":"validation","validationCode":"assert a.ndim >= 2 and a.shape[-2] == a.shape[-1]","typeGuard":"def is_square(a: jax.Array) -> bool:\n    return a.ndim >= 2 and a.shape[-2] == a.shape[-1]","tryCatchPattern":null,"preventionTips":["Hessenberg only defined for square matrices"],"tags":["jax","linalg","hessenberg","shape-validation"],"backgroundTag":"matrix-not-square","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}