From 11e966eb6c7ed38cdf70b547db8d74a758eb64fa Mon Sep 17 00:00:00 2001 From: Jonathan Brodrick Date: Fri, 19 Jun 2026 12:12:13 +0100 Subject: [PATCH 1/4] Pure JAX implementation of ormqr for faster QR solves --- desc/_qr_multiply.py | 109 +++++++++++++++++++++++++++++++++++++++++++ desc/backend.py | 13 +----- 2 files changed, 111 insertions(+), 11 deletions(-) create mode 100644 desc/_qr_multiply.py diff --git a/desc/_qr_multiply.py b/desc/_qr_multiply.py new file mode 100644 index 0000000000..e0174bfd75 --- /dev/null +++ b/desc/_qr_multiply.py @@ -0,0 +1,109 @@ +"""Pure-JAX ``qr_multiply`` for JAX < 0.10.0. + +Implements :func:`jax.scipy.linalg.qr_multiply` without the ``ormqr`` primitive +(which requires a jaxlib rebuild), so it runs on any installed jaxlib. The +Householder reflectors ``Q`` are applied to ``c`` via the blocked CWY/UT +transform of Puglisi (1992) and Joffrain & Low (2006) without ever forming +``Q`` -- on large tall systems this beats forming ``Q`` (``orgqr``) and, on GPU, +cuSOLVER's ``ormqr``. + +Ported from https://github.com/jax-ml/jax/pull/36575. The ``DotAlgorithmPreset`` +used there only mattered for float32 throughput, which DESC does not use, so +plain matmuls are used here. Block sizes are the upstream A100-tuned values. +""" + +from functools import partial + +import jax +import jax.numpy as jnp +from jax import vmap +from jax.scipy.linalg import solve_triangular + + +def _householder_multiply(a, taus, c, *, transpose=False): + """Apply the reflectors in ``a``/``taus`` to ``c`` from the left, one block. + + Forms ``Q = I - V T^{-1} V^H`` via the identity ``T^{-1} + T^{-H} = V^H V``, + recovering the triangular ``T^{-1}`` by correcting the diagonal. + """ + m, k = a.shape + # V: unit lower-trapezoidal (reflectors below the diagonal, unit diagonal). + V = jnp.where(jnp.tril(jnp.ones((m, k), bool), -1), a, jnp.eye(m, k, dtype=a.dtype)) + diag_correction = (1 / taus) if transpose else (1 / taus).conj() + diag_correction = jnp.expand_dims(diag_correction, -1) * jnp.eye(k, dtype=a.dtype) + Vh = V.conj().swapaxes(-1, -2) + # solve_triangular reads only the relevant triangle, so passing the full + # Gram matrix V^H V (minus the diagonal correction) recovers T^{-1}. + T_inv = Vh @ V - diag_correction + z = solve_triangular(T_inv, Vh @ c, lower=transpose) + with jax.default_matmul_precision("highest"): + return c - V @ z + + +def _blocked_householder_multiply(a, taus, c, *, left, transpose): + """Apply ``Q`` (or ``Q^H``) to ``c`` in blocks (block sizes tuned on A100).""" + if not left: # c @ Q == (Q^H @ c^H)^H + ct = c.conj().swapaxes(-1, -2) + out = _blocked_householder_multiply( + a, taus, ct, left=True, transpose=not transpose + ) + return out.conj().swapaxes(-1, -2) + + if a.ndim > 2: # batch dims via vmap, keep the core logic 2-D + fn = partial(_blocked_householder_multiply, left=True, transpose=transpose) + return vmap(fn)(a, taus, c) + + m = a.shape[0] + k = taus.shape[0] + if k == 0: # no reflectors -> Q is the identity + return c + # Balances the per-block V^H V cost against the number of sequential blocks. + esize = a.dtype.itemsize + hi_limit = 4096 * max(1, 8 // esize) + mid_limit = 4096 * esize + nb = min(k, 256 if m <= hi_limit else 128 if m <= mid_limit else 64) + + blocks = range(0, k, nb) + for j0 in blocks if transpose else reversed(blocks): + c = c.at[j0:, :].set( + _householder_multiply( + a[j0:, j0 : j0 + nb], taus[j0 : j0 + nb], c[j0:, :], transpose=transpose + ) + ) + return c + + +def qr_multiply(a, c, mode="right"): + """Pure-JAX drop-in for ``jax.scipy.linalg.qr_multiply``; returns ``(CQ, R)``. + + ``a = Q @ R`` is the (economic) QR factorization. For ``mode="right"`` + returns ``c @ Q`` (with 1-D ``c`` treated as a length-``M`` row vector); for + ``mode="left"`` returns ``Q @ c``. ``R`` has shape ``(min(M, N), N)``. + """ + m, n = a.shape + k = min(m, n) + # mode="raw" returns the packed reflectors (transposed) plus tau factors, + # via the existing geqrf primitive -- no new primitive / jaxlib rebuild. + h, taus = jnp.linalg.qr(a, mode="raw") + packed = h.swapaxes(-1, -2) # (M, N): lower triangle = reflectors, upper = R + R = jnp.triu(packed)[:k, :] + # When m <= n geqrf's last reflector (row m-1) is the identity (tau == 0). + # Drop it statically -- the economic Q is unchanged and we avoid 1/0. + n_refl = k - 1 if m <= n else k + V = packed[:, :n_refl] + taus = taus[:n_refl] + c1d = c.ndim == 1 + + if mode == "right": + cq = _blocked_householder_multiply( + V, taus, c[None, :] if c1d else c, left=False, transpose=False + ) + cq = cq[:, :k] # economic Q has min(M, N) columns + return (cq[0] if c1d else cq), R + + C = c[:, None] if c1d else c + pad = jnp.zeros((m - k, C.shape[1]), C.dtype) + cq = _blocked_householder_multiply( + V, taus, jnp.vstack([C, pad]), left=True, transpose=False + ) + return (cq[:, 0] if c1d else cq), R diff --git a/desc/backend.py b/desc/backend.py index e893d540f5..2de3fe0d32 100644 --- a/desc/backend.py +++ b/desc/backend.py @@ -92,17 +92,8 @@ def _diag_to_full(d, e): if Version(jax.__version__) >= Version("0.10.0"): from jax.scipy.linalg import qr_multiply else: - - def qr_multiply(a, c, mode="right"): - """Fallback for ``jax.scipy.linalg.qr_multiply`` (added in JAX 0.10.0).""" - Q, R = qr(a, mode="economic") - if mode == "right": - # 1-D c (all DESC uses) matches the old Q.T @ c; c @ Q keeps - # higher-dim c consistent with qr_multiply rather than silently wrong - cq = Q.T @ c if c.ndim == 1 else c @ Q - else: - cq = Q @ c - return cq, R + # pure-JAX implementation; needs no ormqr primitive / jaxlib rebuild + from desc._qr_multiply import qr_multiply from jax.scipy.special import gammaln from jax.tree_util import ( From de5b730ab011e94af6d9ac3daaf9e7cc2e0db1f7 Mon Sep 17 00:00:00 2001 From: Jonathan Brodrick Date: Fri, 19 Jun 2026 12:47:32 +0100 Subject: [PATCH 2/4] update changelog --- CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2449a4d868..a77af8eb73 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,8 @@ Changelog Performance Improvements -- Speeds up the ``"qr"`` trust-region subproblem and Newton-step solves in the least-squares optimizers by reusing the Jacobian QR factorization across the Levenberg-Marquardt parameter sweep. On ``jax >= 0.10.0`` this uses ``qr_multiply`` to additionally avoid forming ``Q`` explicitly; on older versions a fallback preserves the same results. +- Speeds up the ``"qr"`` trust-region subproblem and Newton-step solves in the least-squares optimizers by reusing the Jacobian QR factorization across the Levenberg-Marquardt parameter sweep, and by using ``qr_multiply`` to apply ``Q`` without forming it explicitly. +- Adds a pure-JAX ``qr_multiply`` fallback (a blocked Householder / compact-WY implementation) for ``jax < 0.10.0``, so the above ``Q``-avoidance speedup is available on the currently pinned JAX with no jaxlib rebuild (~1.5-2x faster than forming ``Q`` on CPU for tall Jacobians, larger on GPU). v0.17.2 From 7b41861d27bf452ba70c653c804e4faa45d32237 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Fri, 19 Jun 2026 19:02:11 +0300 Subject: [PATCH 3/4] update the tr method to ue svd for the singular matrix --- tests/test_optimizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 233e47cda7..fad69672fb 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -1079,7 +1079,7 @@ def test_constrained_AL_lsq(): ctol=ctol, x_scale="auto", copy=True, - options={}, + options={"tr_method": "svd"}, ) V2 = eq2.compute("V")["V"] AR2 = eq2.compute("R0/a")["R0/a"] From 9a8b511b57067b7fdc5d0ea2e7676b871c099269 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Fri, 19 Jun 2026 19:08:04 +0300 Subject: [PATCH 4/4] jit qr_multiply to reduce overhead, update imports --- desc/_qr_multiply.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/desc/_qr_multiply.py b/desc/_qr_multiply.py index e0174bfd75..3b6a3b5fdc 100644 --- a/desc/_qr_multiply.py +++ b/desc/_qr_multiply.py @@ -14,10 +14,7 @@ from functools import partial -import jax -import jax.numpy as jnp -from jax import vmap -from jax.scipy.linalg import solve_triangular +from desc.backend import jax, jit, jnp, solve_triangular, vmap def _householder_multiply(a, taus, c, *, transpose=False): @@ -73,6 +70,7 @@ def _blocked_householder_multiply(a, taus, c, *, left, transpose): return c +@partial(jit, static_argnames="mode") def qr_multiply(a, c, mode="right"): """Pure-JAX drop-in for ``jax.scipy.linalg.qr_multiply``; returns ``(CQ, R)``.