{"record":{"id":"bb784f5af283b2e2","repo":"jax-ml/jax","slug":"hessenberg-requires-the-last-dimension-of-a-to-be","errorCode":null,"errorMessage":"hessenberg requires the last dimension of a to be constant, got a.shape of {a.shape}.","messagePattern":"hessenberg requires the last dimension of a to be constant, got a\\.shape of (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/linalg.py","lineNumber":1397,"sourceCode":"def _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)\n  avals_out = [*ctx.avals_out, ShapedArray(batch_dims, np.int32)]\n  rule = _linalg_ffi_lowering(target_name, avals_out=avals_out,\n                              operand_output_aliases={0: 0})\n  a, taus, info = rule(ctx, a, low=np.int32(1), high=np.int32(n))\n  ok = mlir.compare_hlo(\n      info, mlir.full_like_aval(ctx, 0, ShapedArray(batch_dims, np.dtype(np.int32))),\n      \"EQ\", \"SIGNED\")\n  return [\n      _replace_not_ok_with_nan(ctx, batch_dims, ok, a, ctx.avals_out[0]),\n      _replace_not_ok_with_nan(ctx, batch_dims, ok, taus, ctx.avals_out[1]),\n  ]\n\n\nhessenberg_p = linalg_primitive(\n    _hessenberg_dtype_rule, (_float | _complex,), (2,), _hessenberg_shape_rule,\n    \"hessenberg\", multiple_results=True)","sourceCodeStart":1379,"sourceCodeEnd":1415,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/linalg.py#L1379-L1415","documentation":"jax/_src/lax/linalg.py:1397 in _hessenberg_cpu_lowering. The CPU LAPACK gehrd call needs the matrix dimension n as a concrete integer (low/high bounds). If n is polymorphic (a DimExpr from dynamic shapes / jax.jit with symbolic dimensions), JAX cannot pass it and raises ValueError.","triggerScenarios":"Calling jax.lax.linalg.hessenberg on CPU inside a jit-compiled function whose trailing matrix dimension is symbolic/dynamic, e.g. exported with dynamic shapes or when using jax.export with dimension variables; n = a_aval.shape[-1] fails core.is_constant_dim(n).","commonSituations":"JAX export/Import serving pipelines with dynamic batch-less matrix size; using shape polymorphism (jax.experimental.jax2vec or export dimension variables); code that worked eagerly or with static shapes breaks under dynamic-shape tracing.","solutions":["Ensure the matrix is padded to a fixed size so the last dim is static under jit","Call hessenberg outside jit / on concrete shapes (eager path)","On GPU the FFI lowering may tolerate dynamic dims — switch backend, or restructure to avoid hessenberg under dynamic tracing"],"exampleFix":"// before\nf = jax.jit(lambda a: jax.lax.linalg.hessenberg(a))  # traced with dynamic n\n// after\nf = jax.jit(lambda a: jax.lax.linalg.hessenberg(a), static_argnums=())  # ensure static shape\n# or pad to fixed size:\na_fixed = jnp.pad(a, ((0, N - a.shape[-2]), (0, N - a.shape[-1])))","handlingStrategy":"validation","validationCode":"n = a.shape[-1]\nassert isinstance(n, int) or core.is_constant_dim(n), 'pad to static size before jit'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pad matrices to fixed size under jit","Avoid hessenberg under dynamic-shape export on CPU"],"tags":["jax","hessenberg","dynamic-shapes","jit","cpu"],"backgroundTag":"dynamic-shape-not-supported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}