Batched tridiagonal solve

This specialized solver handles

\[ \ell_jx_{j-1}+d_jx_j+u_jx_{j+1}=b_j, \qquad j=0,\ldots,n-1, \]

with the system dimension on the leading axis. Every trailing axis is an independent right-hand side or coefficient batch.

x = sx.tridiagonal_solve(lower, diag, upper, rhs, method="auto")

For a periodic line, retain the two corner coefficients in lower[0] and upper[-1] and use cyclic_tridiagonal_solve. It applies an exact Sherman–Morrison correction through one tridiagonal call with two right-hand sides, so the same fused accelerator backend remains available. Its derivation is given in the cyclic-systems section below.

If rhs.shape == (n, n_columns, n_fields), one call solves n_columns * n_fields systems. lower, diag, and upper may carry matching trailing batch axes and are broadcast across extra right-hand-side axes.

Thomas derivation

Forward elimination defines modified coefficients

\[ c'_0=\frac{u_0}{d_0},\qquad d'_0=\frac{b_0}{d_0}, \]

\[ q_j=d_j-\ell_jc'_{j-1},\qquad c'_j=\frac{u_j}{q_j},\qquad d'_j=\frac{b_j-\ell_jd'_{j-1}}{q_j}. \]

Back substitution gives

\[ x_{n-1}=d'_{n-1},\qquad x_j=d'_j-c'_jx_{j+1}. \]

The work is \(O(n)\) per system and storage is \(O(n)\); no pivoting occurs [GVL13, Tho49].

Backend selection

method	Implementation	Intended use
"thomas"	two jax.lax.scan sweeps	CPU, deterministic arithmetic, small systems
"lax"	jax.lax.linalg.tridiagonal_solve	fused accelerator/vendor path
"auto"	platform-dependent selection	default portable behavior

auto selects Thomas when lowering for CPU and the fused operation otherwise. Systems with fewer than three rows use Thomas because accelerator kernels may require a larger minimum dimension.

For real bands and a complex right-hand side, SOLVAX solves the real and imaginary parts independently. This keeps real band storage and the fast fused accelerator kernel; packing them as a multiple-right-hand-side field is avoided because supported GPU backends can lower that layout poorly. Genuinely complex bands use portable Thomas arithmetic for every method because the fused JAX primitive is real-only on supported JAX releases.

The Thomas arithmetic order is fixed and therefore reproducible on a fixed backend. The fused and Thomas paths should agree to floating-point accuracy, not necessarily bit for bit.

The fused primitive is wrapped in an implicit linear solve, supplying forward- and reverse-mode derivatives even on supported JAX versions whose primitive does not define its own differentiation rule.

Boundary entries

For tridiagonal_solve, lower[0] and upper[-1] do not correspond to matrix entries. Set them to zero for clarity, although the algorithm ignores their out-of-domain role. cyclic_tridiagonal_solve reuses those two slots to carry the periodic corner coupling (see below).

Cyclic (periodic) systems

A periodic line couples the two endpoints,

\[ d_0x_0+u_0x_1+\beta x_{n-1}=b_0,\qquad \alpha x_0+\ell_{n-1}x_{n-2}+d_{n-1}x_{n-1}=b_{n-1}, \]

with the corner entries stored as \(\beta=\) lower[0] and \(\alpha=\) upper[-1]. The matrix is tridiagonal apart from those two corners, so it is a rank-one update of an ordinary tridiagonal matrix [GVL13, PTVF07]. Pick any \(\gamma\neq0\) and write

\[ A=T+uv^{T},\qquad u=\gamma e_0+\alpha e_{n-1},\qquad v=e_0+\tfrac{\beta}{\gamma}e_{n-1}, \]

where \(T\) equals \(A\) with the corners removed and the endpoints shifted, \(T_{00}=d_0-\gamma\) and \(T_{n-1,n-1}=d_{n-1}-\alpha\beta/\gamma\). SOLVAX takes \(\gamma=-d_0\) (falling back to \(\gamma=-1\) when \(d_0\) underflows) so the first pivot stays well scaled. The Sherman–Morrison identity then gives the solution from two ordinary tridiagonal solves \(Ty=b\) and \(Tz=u\),

\[ x=y-\frac{v^{T}y}{1+v^{T}z}\,z . \]

Both right-hand sides are stacked into a single tridiagonal_solve call, so a periodic line costs one tridiagonal solve with two columns and inherits the same reproducible-Thomas / fused-cuSPARSE method selection. The construction is exact (not iterative) and fully differentiable.

x = sx.cyclic_tridiagonal_solve(lower, diag, upper, rhs, method="auto")

Trailing axes of rhs are extra right-hand sides solved together, exactly as for tridiagonal_solve. Use this instead of periodic banded LU when the band is a single sub- and superdiagonal; for wider periodic bands use solvax.banded.lu_solve_banded_periodic().

Example: many field-line columns

n, n_modes, n_fields = 256, 64, 3
lower = jnp.full((n, n_modes), -1.0)
diag = jnp.full((n, n_modes), 4.0)
upper = jnp.full((n, n_modes), -1.0)
rhs = jnp.ones((n, n_modes, n_fields))

x = sx.tridiagonal_solve(lower, diag, upper, rhs)

Use as a line preconditioner

For a two-dimensional anisotropic PDE, reshape a residual so the strongly coupled direction is leading, solve all transverse lines together, and reshape back. Combine orthogonal line solves with solvax.precond.line_smoother().

Comparison with alternatives

• Compared with general banded LU, Thomas has less storage and lower constants but handles only one sub- and superdiagonal.

• Compared with block Thomas, it treats trailing axes as independent systems, not dense within-point coupling.

• Compared with parallel cyclic reduction, Thomas is sequential along the system axis; the fused vendor path is important on accelerators.

• Compared with dense solve, it preserves \(O(n)\) rather than \(O(n^3)\) work.

Stability and diagnostics

Thomas elimination is safe for standard strictly diagonally dominant and many SPD tridiagonal operators. A zero or tiny modified pivot \(q_j\) causes numerical failure. There is no pivot count, so validate difficult coefficient regimes by checking the residual or using a pivoted banded/native solver.

The function returns only x; it does not return convergence diagnostics because it is a direct algorithm.

Differentiation

Both backends are JAX-traceable. Gradients can flow through coefficients and right-hand sides:

loss = lambda d: jnp.sum(
    sx.tridiagonal_solve(lower, d, upper, rhs, method="thomas") ** 2
)
gradient = jax.grad(loss)(diag)

API summary

• solvax.tridiagonal.tridiagonal_solve()

• solvax.tridiagonal.cyclic_tridiagonal_solve()

Runnable counterparts: examples/14_tridiagonal_solve.py and examples/21_cyclic_tridiagonal.py.