jax-ml/jax · error · NotImplementedError

method='qdwh' only supports mxn matrices where m < n where s

Error message

method='qdwh' only supports mxn matrices where m < n where side='right' and m >= n side='left', got {a.shape} with {side=}

What it means

The QDWH-based polar decomposition in jax.scipy.linalg.polar only handles matrices where the shape/side combination lets it reduce to the wide m<n case with side='right' (transposing otherwise). If you pass side='right' with m>=n, or side='left' with m<n, the unsupported combination raises NotImplementedError. This is a limitation of the qdwh implementation path, not of polar decomposition in general.

Source

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

  eigenvectors = _compute_eigenvectors(alpha, beta, mid, key)
  return mid, eigenvectors.T

@jit(static_argnames=('side', 'method'))
@config.default_matmul_precision("float32")
def _polar_2d(a: Array, side: str, method: str, eps: float | None,
              max_iterations: int | None) -> tuple[Array, Array]:
  m, n = a.shape
  if method == "qdwh":
    # TODO(phawkins): return info also if the user opts in?
    if m >= n and side == "right":
      unitary, posdef, _, _ = qdwh.qdwh(a, is_hermitian=False, eps=eps)
    elif m < n and side == "left":
      a = a.T.conj()
      unitary, posdef, _, _ = qdwh.qdwh(a, is_hermitian=False, eps=eps)
      posdef = posdef.T.conj()
      unitary = unitary.T.conj()
    else:
      raise NotImplementedError("method='qdwh' only supports mxn matrices "
                                "where m < n where side='right' and m >= n "
                                f"side='left', got {a.shape} with {side=}")
  elif method == "svd":
    u_svd, s_svd, vh_svd = lax_linalg.svd(a, full_matrices=False)
    s_svd = s_svd.astype(u_svd.dtype)
    unitary = u_svd @ vh_svd
    if side == "right":
      posdef = (vh_svd.T.conj() * s_svd[None, :]) @ vh_svd
    else:
      posdef = (u_svd * s_svd[None, :]) @ (u_svd.T.conj())
  else:
    raise ValueError(f"Unknown polar decomposition method {method}.")
  return unitary, posdef


@jit(static_argnames=('side', 'method'))
def polar(a: ArrayLike, side: str = 'right', *, method: str = 'qdwh', eps: float | None = None,
          max_iterations: int | None = None) -> tuple[Array, Array]:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Switch to the SVD path: polar(a, side=..., method='svd') — always supported
2. Swap the side: use side='left' when m>=n or side='right' when m<n
3. Transpose the input manually and transpose the results back
4. Update JAX; newer versions may normalize shapes internally

Example fix

# before
U, H = jax.scipy.linalg.polar(a)  # a shape (100, 10) -> raises
# after
U, H = jax.scipy.linalg.polar(a, side='left')
# or
U, H = jax.scipy.linalg.polar(a, method='svd')

import jax.numpy as jnp

def polar_safe(a, side='right', method='qdwh'):
    a = jnp.asarray(a)
    m, n = a.shape[-2:]
    if method == 'qdwh':
        if side == 'right' and m >= n: side = 'left'
        elif side == 'left' and m < n: side = 'right'
    return jax.scipy.linalg.polar(a, side=side, method=method)

try:
    U, H = polar(a, side=side)
except NotImplementedError:
    U, H = polar(a, side=side, method='svd')

When it happens

Trigger: polar(a, side='right', method='qdwh') with a of shape (m, n) where m >= n; or polar(a, side='left', method='qdwh') with m < n.

Common situations: Using method='qdwh' (the default) on tall matrices with the default side='right', e.g. in orthogonalization/Procrustes pipelines ported from NumPy where side was never considered.

Related errors

• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac
• Incorrect output shape for return value #{i}: Expected: {out
• Length of sharding.spec ({len(out_s.spec)}) must be equal to

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