jax-ml/jax · error · ValueError

Incompatible shapes for Sylvester equation: A: {A.shape} B:

Error message

Incompatible shapes for Sylvester equation:
A: {A.shape}
B: {B.shape}
C: {C.shape}

What it means

For the Sylvester equation AX + XB = C, C's last two dims (m, n) dictate that A be (m, m) and B be (n, n). JAX validates these trailing shapes after dtype promotion and before vectorizing over batch dims.

Source

Thrown at jax/_src/scipy/linalg.py:3364

    the eigen decomposition method because you need to perform a Schur decomposition and then scan the entire solution matrix.
    Second, it requires more system memory compared to the eigen decomposition method.

    The eigen decomposition method is the fastest method to solve a sylvester equation. However, this speed brings with it a couple of drawbacks.
    First, A and B must be diagonalizable otherwise the eigenvectors will be linearly dependent and ill-conditioned leading to accuracy issues.
    Second, when the eigenvectors are not orthogonal roundoff errors are amplified.

    Additionally, for complex types as the size of the matrix increases the accuracy of the results degrades. Float64 types are most robust to degradation.

    The tol argument allows you to specify how ill-conditioned a matrix can be and still estimate a solution.
    For matrices that are ill-conditioned we recommend using float64 instead of the default float32 dtype. The solver
    can still return good estimates for ill-conditioned matrices depending on how close to zero the sums of the eigenvalues of A and B
    are.
  """
  A, B, C = promote_dtypes_inexact(jnp.asarray(A), jnp.asarray(B), jnp.asarray(C))

  m, n = C.shape[-2:]
  if A.shape[-2:] != (m, m) or B.shape[-2:] != (n, n):
    raise ValueError(f"Incompatible shapes for Sylvester equation:\nA: {A.shape}\nB: {B.shape}\nC: {C.shape}")

  return jnp_vectorize.vectorize(
      partial(_solve_sylvester_2d, method=method, tol=tol),
      signature="(m,m),(n,n),(m,n)->(m,n)")(A, B, C)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Check C.shape[-2:] == (m, n), A.shape[-2:] == (m, m), B.shape[-2:] == (n, n) and transpose/reshape inputs accordingly
  2. Verify you are not confusing the row/column dimensions of C with A vs B
  3. Add an assert before calling in test code

Example fix

# before
X = linalg.solve_sylvester(A, B, C.T)  # wrong orientation
# after
m, n = C.shape[-2:]
assert A.shape[-2:] == (m, m) and B.shape[-2:] == (n, n)
X = linalg.solve_sylvester(A, B, C)
Defensive patterns

Strategy: validation

Validate before calling

m, n = C.shape[-2:]
assert A.shape[-2:] == (m, m) and B.shape[-2:] == (n, n)

Type guard

def sylvester_shapes_ok(A, B, C) -> bool:
    m, n = C.shape[-2:]
    return A.shape[-2:] == (m, m) and B.shape[-2:] == (n, n)

Prevention

When it happens

Trigger: Passing A of shape (m, k), B of shape (k, n), or C with shape not matching A/B sizes, e.g. transposed C.

Common situations: Building equation data from training batches where matrices got flattened or transposed; assuming B is (m, m) like A.

Related errors


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