Test suite fails with jax >= 0.10: spurious FutureWarning from jax's internal `cho_solve` dispatch, escalated by `filterwarnings = ["error"]`

[s-sasaki-earthsea-wizard]

Summary

With jax/jaxlib 0.10.x (conda-forge currently serves 0.10.2), 44 tests fail on an unmodified checkout of main (c2f7c24). This is not a dolphin regression: jax 0.10.0 added a deprecation warning that jax.scipy.linalg.solve(assume_a="pos") trips internally, and dolphin's pytest config (filterwarnings = ["error"]) escalates it to an error at trace time.

Evidence

CI run of unmodified main (c2f7c24) on my fork, same test-build-push.yml workflow: https://github.com/s-sasaki-earthsea-wizard/dolphin/actions/runs/31772307761

• 44 failed / 702 passed / 19 skipped (test_phase_link_core, test_phase_link_compress, workflow tests)
• Every failure's traceback ends the same way: _core.py:521 (process_coherence_matrices) → crlb.py:187 (_crlb_from_x) → jax/_src/scipy/linalg.py:170: FutureWarning

FutureWarning: jax.scipy.linalg.cho_solve: batched 1D solves with b.ndim > 1 are deprecated, and in the future will be treated as a batched 2D solve. Use cho_solve(c_and_lower, b[..., None]).squeeze(-1) to avoid this warning.

Any CI run whose environment resolves jax >= 0.10.0 will hit this (reproduced with 0.10.1 and 0.10.2); the last unaffected jax release is 0.9.2.

Root cause

crlb.py calls jax.scipy.linalg.solve(..., assume_a="pos") on batched matrices (crlb.py:187 on every CRLB computation; also crlb.py:178 when aps_variance > 0 and crlb.py:241 in compute_crlb_jax). Inside jax 0.10.x, _solve implements assume_a="pos" by vmapping an internal cho_solve(factors, x) over the RHS columns (linalg.py#L1196); the per-example rhs then has x.ndim == factors.ndim - 1, which is exactly the batched-1D dispatch that jax 0.10.0 deprecated (linalg.py#L164-L177) — so jax's own internal call emits the warning. User code never touches cho_solve; results are numerically correct; only the warning is spurious.

Dolphin's two direct cho_solve call sites are not affected (verified on 0.10.2): phase_link/_core.py:472 has equal-ndim c/b and takes the matrix branch; timeseries.py:1496 is vmapped with a per-pixel 1D rhs.

Options

1. Targeted warning ignore (smallest diff, my suggestion): add to pyproject.toml's filterwarnings, following the existing h5py precedent:

   # jax >= 0.10 solve(assume_a="pos") trips its own cho_solve deprecation
   # warning internally; remove once fixed upstream in jax:
   "ignore:jax.scipy.linalg.cho_solve. batched 1D solves.*:FutureWarning",

   Keeps CI on latest jax, no behavior change, self-documenting removal condition.

2. Avoid the warning at the source: replace the assume_a="pos" calls in crlb.py with explicit cho_factor + cho_solve on the matrix RHS (the non-deprecated branch). Numerically identical; would forgo custom_linear_solve's factorization reuse under differentiation, which the CRLB path doesn't appear to use.

3. Pin jax<0.10: least attractive — that walks back three minor releases and the warning is still present on jax main, so the pin would linger.

Happy to open the small PR for option 1 (or 2), and/or to file this upstream with jax as a minimal-repro issue — your call.

---

Disclosure: AI tools assisted with research, implementation, and drafting. I verified the technical claims and take responsibility for the proposal.

[s-sasaki-earthsea-wizard]

Upstream follow-up: this is a jax bug, and a fix is already in flight.

I filed a minimal repro as jax-ml/jax#39986, and a jax maintainer opened jax-ml/jax#39996 the same day. The fix is the one-liner the analysis above predicted — in _solve's custom_linear_solve callback, where jax passes a batched-1D RHS into its own cho_solve:

-      solve=lambda _, x: cho_solve(factors, x),
+      solve=lambda _, x: cho_solve(factors, x[..., None]).squeeze(-1),

So nothing needs to change in dolphin's own code — once a jax release containing the fix is out, this is just a pyproject.toml matter.

That PR is still open and unreleased (latest is 0.11.0), so CI stays red on any environment resolving jax >= 0.10 until then. If you'd rather not wait, the option-1 filterwarnings ignore works as a bridge and can be dropped as soon as CI resolves a fixed jax — happy to open that PR. Otherwise, leaving this issue open until the release lands is fine by me.

[s-sasaki-earthsea-wizard]

Update: jax-ml/jax#39996 was merged and jax-ml/jax#39986 is now closed.

Note the fix is only on jax main for now: every released version from 0.10.0 through 0.11.0 still carries the bug, so CI here will stay red until the next jax release reaches conda-forge. No dolphin-side changes are needed. If the wait is a problem, I'm still happy to open the small filterwarnings bridge PR mentioned above as a stopgap.

[scottstanie]

i fixed it with scottstanie#9 to avoid the deprecation warning on 0.10 releases

[s-sasaki-earthsea-wizard]

Thanks for the fix! For the record, that's option 2 from the writeup (explicit cho_factor + cho_solve on the matrix RHS, i.e. the non-deprecated branch) — and it looks like the right call over waiting on the environment:

• The upstream fix did get released: jax-ml/jax#39996 landed in jax 0.11.1 (2026-08-17), verified at linalg.py#L1196.
• But conda-forge still tops out at 0.10.2 (no 0.11.x there yet), and every released 0.10.0–0.11.0 carries the warning permanently — your patch makes dolphin quiet on all of them without waiting for environments to catch up.

The one theoretical cost I noted for option 2 (losing custom_linear_solve's factorization reuse under differentiation) doesn't apply here, since nothing differentiates through the CRLB path.

Happy to close this whenever the change lands on main here.