jax-ml/jax · error · ValueError

hessenberg requires the last dimension of a to be constant,

Error message

hessenberg requires the last dimension of a to be constant, got a.shape of {a.shape}.

What it means

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.

Source

Thrown at jax/_src/lax/linalg.py:1397

def _hessenberg_shape_rule(shape, **_):
  if shape[0] != shape[-1]:
    raise ValueError(
        "Argument to Hessenberg reduction must have shape [..., n, n], "
        f"got shape {shape}"
    )
  return shape, shape[:-2] + (shape[-1] - 1,)


def _hessenberg_dtype_rule(dtype, **_):
  return dtype, dtype


def _hessenberg_cpu_lowering(ctx, a):
  a_aval, = ctx.avals_in
  batch_dims = a_aval.shape[:-2]
  n = a_aval.shape[-1]
  if not core.is_constant_dim(n):
    raise ValueError("hessenberg requires the last dimension of a to be "
                     f"constant, got a.shape of {a.shape}.")
  target_name = lapack.prepare_lapack_call("gehrd_ffi", a_aval.dtype)
  avals_out = [*ctx.avals_out, ShapedArray(batch_dims, np.int32)]
  rule = _linalg_ffi_lowering(target_name, avals_out=avals_out,
                              operand_output_aliases={0: 0})
  a, taus, info = rule(ctx, a, low=np.int32(1), high=np.int32(n))
  ok = mlir.compare_hlo(
      info, mlir.full_like_aval(ctx, 0, ShapedArray(batch_dims, np.dtype(np.int32))),
      "EQ", "SIGNED")
  return [
      _replace_not_ok_with_nan(ctx, batch_dims, ok, a, ctx.avals_out[0]),
      _replace_not_ok_with_nan(ctx, batch_dims, ok, taus, ctx.avals_out[1]),
  ]


hessenberg_p = linalg_primitive(
    _hessenberg_dtype_rule, (_float | _complex,), (2,), _hessenberg_shape_rule,
    "hessenberg", multiple_results=True)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the matrix is padded to a fixed size so the last dim is static under jit
2. Call hessenberg outside jit / on concrete shapes (eager path)
3. On GPU the FFI lowering may tolerate dynamic dims — switch backend, or restructure to avoid hessenberg under dynamic tracing

Example fix

// before
f = jax.jit(lambda a: jax.lax.linalg.hessenberg(a))  # traced with dynamic n
// after
f = jax.jit(lambda a: jax.lax.linalg.hessenberg(a), static_argnums=())  # ensure static shape
# or pad to fixed size:
a_fixed = jnp.pad(a, ((0, N - a.shape[-2]), (0, N - a.shape[-1])))

n = a.shape[-1]
assert isinstance(n, int) or core.is_constant_dim(n), 'pad to static size before jit'

When it happens

Trigger: 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).

Common situations: 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.

Related errors

• ragged_dot_general requires the group count (last dimension
• Formatting arguments to checkify.check need to be PyTrees of
• Value of type {type(self)} is not convertible to float.
• Value of type {type(self)} is not convertible to complex.
• {self.__class__.__name__} has no attribute {name}

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