linalg: expm(A) by scaling and squaring

[axiom-of-choice]

Fixes #133.

Stacked on #170 (solve, #132), since both touch linalg.py. Merge that first and this becomes a two-file diff; happy to rebase onto main instead if you would rather take them in the other order.

Approach

Exactly what the issue describes: the host picks s from the 1-norm so A / 2^s has norm at most 1/2, which is where a low-order Taylor sum is accurate, then the result is squared back s times. The Taylor sum and the squarings are all _ane_gemm, the same building block matrix_power (#103) uses, so the whole thing is O(order + s) on-engine gemms.

The term is built by recurrence, term = term @ As / k, so each order costs one gemm rather than recomputing a power, and 1/k! folds into a host-side scalar so it needs no extra dispatch.

Verification

Apple M2 Pro (Mac14,12 Mac mini), macOS 26.5.2, against scipy.linalg.expm.

Modest norm, well conditioned:

n	relerr
4	4.8e-4
6	7.8e-4
8	7.6e-4

The scaling half is the part worth testing separately, since a Taylor sum alone would diverge there:

1-norm	s chosen	relerr
0.64	1	5.2e-4
2.80	3	2.6e-3
6.54	4	2.2e-3

Check	Result
tests/test_linalg.py	74 passed (57 on main, 9 from #170, 8 here)
tests/run_corpus.py	GATE: GREEN, 90/90
ruff / pyright / .githooks/pre-commit	clean

Two of the tests avoid a reference solver entirely, which matters because scipy is not in the dev extra (the scipy-based ones use importorskip): expm(0) must be the identity, and a diagonal argument must equal the elementwise exp, which pins the Taylor sum on its own.

Judgement calls

order=8 as the default. With the norm scaled to 1/2, the term at k=8 is already below fp16 resolution, so more orders buy nothing and cost gemms. It is a parameter rather than a constant in case you want it tunable, and order < 1 rejects.

The threshold is 1/2 rather than a Pade-style table. scipy switches order and threshold together from a lookup table tuned for fp64; at fp16 the accuracy is dominated by the squaring error, not the truncation, so a single conservative threshold is the honest choice. The docstring says so, and says this is for modest norms on well-conditioned input, the same caveat matrix_power carries.

I used the 1-norm computed host-side rather than calling norm(A, 1) (#118), since that dispatches a graph to get a scalar the host then needs immediately; happy to switch if you would rather have the one code path.

[sbryngelson]

Reviewed and it's a clean approve on the expm code -- the scaling picks s from the 1-norm so the scaled matrix lands at <= 1/2 (verified against your table: 6.54 -> s=4 -> 0.41), the term = term @ As / k recurrence builds exp(As) at one gemm per order, and the s squarings recover exp(A); zero matrix falls through to the identity. The two reference-free tests (identity, diagonal closed form) are exactly the right pins given scipy isn't in the dev extra. Nice.

#170 just merged, so please take the rebase-onto-main option you offered -- git rebase --onto main feat-132-linalg-solve feat-133-linalg-expm should drop the now-redundant solve commits and leave the expm-only two-file diff. Once it shows just expm I'll merge (closes #133). Thanks for the clean, well-verified series.

[axiom-of-choice]

Reviewed and it's a clean approve on the expm code -- the scaling picks s from the 1-norm so the scaled matrix lands at <= 1/2 (verified against your table: 6.54 -> s=4 -> 0.41), the term = term @ As / k recurrence builds exp(As) at one gemm per order, and the s squarings recover exp(A); zero matrix falls through to the identity. The two reference-free tests (identity, diagonal closed form) are exactly the right pins given scipy isn't in the dev extra. Nice.

#170 just merged, so please take the rebase-onto-main option you offered -- git rebase --onto main feat-132-linalg-solve feat-133-linalg-expm should drop the now-redundant solve commits and leave the expm-only two-file diff. Once it shows just expm I'll merge (closes #133). Thanks for the clean, well-verified series.

No Worries. Just rebased and pushed