[upstream gaussx] Diagonal-array overloads for cavity_distribution and newton_update
- Dominant language
- Python
- Stars
- 1
- Forks
- 0
- Avg merge
- 19h 2m
- Merged PRs (30d)
- 16
Description
## Status
Tracking issue. The actual feature lives upstream in `gaussx` — mirrored there as **jejjohnson/gaussx#198** — and this issue is the pyrox-side tracker so the deferred parts of #124 item 8 have somewhere to land once gaussx ships the diagonal overloads.
Partially unblocked by #143 (which migrated `damped_natural_update` since it already accepts `(N,)` arrays); the cavity and Newton sites still need the upstream issue resolved.
## User Story
> As a pyrox maintainer of the site-based EP / PL / Laplace inference paths, I want the 4 cavity and 4 Newton open-coded arithmetic blocks in `_inference_nongauss.py` / `_inference_nongauss_markov.py` replaced by `gaussx.cavity_distribution` / `gaussx.newton_update` one-liners, so the numerics have a single tested home upstream and pyrox stops carrying duplicate formula copies.
## Motivation
- **Deduplication** — ~30 LOC of hand-rolled cavity/Newton arithmetic across 8 call sites, all reimplementing formulas gaussx already owns in operator form.
- **Asymptotics** — the workaround (wrap `(N,)` in `lx.DiagonalLinearOperator`, call the operator path, extract `.diagonal()`) materialises `(N, N)` matrices: an O(N) → O(N²) inflation on the hot inner loop.
- **Continuity** — completes #124 item 8, started in #143.
## Use case / Mathematics
Site-based EP / PL / Laplace inference for GPs with scalar latents represents per-site precisions as **diagonal `(N,)` arrays** — never as full `(N, N)` matrices. Each iteration computes a cavity:
```python
cav_prec = 1.0 / q_var - site_nat2 # (N,)
cav_var = 1.0 / cav_prec # (N,)
cav_mean = cav_var * (q_mean / q_var - site_nat1)
```
and a Newton site:
```python
Lambda = jnp.maximum(-hessian_diag, precision_floor) # (N,)
nat1_target = grad + Lambda * f # (N,)
nat2_target = Lambda # (N,)
```
The current `gaussx.cavity_distribution(post_mean, post_cov, site_nat1, site_nat2, power=1.0)` takes `post_cov` as `lx.AbstractLinearOperator` and calls `.as_matrix()` / `.mv(...)`. Wrapping a `(N,)` diag in `lx.DiagonalLinearOperator` and extracting `.diagonal()` on the way out works, but materialises an `(N, N)` matrix and is materially slower than the elementwise version above.
`gaussx.newton_update(mean, jacobian, hessian)` takes `hessian: (N, N)`. For diagonal sites that's an `O(N)` → `O(N²)` inflation.
## Proposed API (upstream, for reference)
Multiple-dispatch (or a separate `_diag` variant):
```python
# When post_cov / hessian / site_nat2 is a (N,) array:
cavity_distribution(
post_mean: (N,),
post_var: (N,), # diag instead of operator
site_nat1: (N,),
site_nat2: (N,),
power: float = 1.0,
) -> tuple[Array, Array] # (cav_mean, cav_var), both (N,)
newton_update(
mean: (N,),
grad: (N,),
hess_diag: (N,), # diag instead of (N, N)
) -> tuple[Array, Array] # (nat1, nat2), both (N,)
```
Keep the operator-typed signatures for the full-cov path; add the diag overload for the elementwise path. The full spec (tests, convention notes, dispatch decision) lives in the gaussx mirror issue.
## Convention compatibility
pyrox's site-based code uses `nat2 = +Λ` (positive precision) consistently. gaussx's `cavity_distribution` (operator form) already follows this convention (`cav_prec = post_prec - power · site_nat2`). The diagonal overload should preserve it, so callers don't need sign flips.
Note that gaussx's `mean_cov_to_natural` / `natural_to_mean_cov` use the exponential-family convention `eta2 = -0.5 · Λ`. That mismatch is orthogonal to this issue (and is the reason pyrox's `_inference.py` doesn't currently use those helpers either) — flagging it so the diagonal overload's docstring can be explicit about which convention it uses.
## Example usage (pyrox side, after migration)
```python
import gaussx
# before: 3 lines of open-coded cavity arithmetic per site — after:
cav_mean, cav_var = gaussx.cavity_distribution(q_mean, q_var, site_nat1, site_nat2, power=power)
# before: 3 lines of open-coded Newton-site arithmetic — after:
nat1_t, nat2_t = gaussx.newton_update(f, grad, hess_diag)
```
## Tasks / sub-tasks
- [ ] Upstream: jejjohnson/gaussx#198 ships diagonal-array overloads for `cavity_distribution` and `newton_update` (`(N,)` in / `(N,)` out, `nat2 = +Λ` convention)
- [ ] pyrox migration PR: replace the 4 cavity sites in `src/pyrox/gp/_inference_nongauss.py` and `src/pyrox/gp/_inference_nongauss_markov.py` with `gaussx.cavity_distribution` calls
- [ ] pyrox migration PR: replace the 4 Newton sites with `gaussx.newton_update` calls; delete the open-coded arithmetic
- [ ] Bump the minimum gaussx version in `pyproject.toml`
- [ ] Confirm `tests/gp/test_inference_nongauss.py` and `tests/gp/test_inference_nongauss_markov.py` pass without tolerance regressions
## Definition of Done
- [ ] All 8 call sites are one-liners delegating to gaussx; the duplicated formulas are gone (~30 LOC removed)
- [ ] Existing EP / PL / Laplace tests pass without tolerance regressions
- [ ] No sign-convention shims needed at the call sites
## References
- Wilkinson, W. J., Särkkä, S. & Solin, A. (2023). *Bayes–Newton methods for approximate Bayesian inference with PSD guarantees*. JMLR.
- Upstream mirror: jejjohnson/gaussx#198
- History: #124 (item 8), #143 (`damped_natural_update` migration)
## Relationships
- Blocked by: jejjohnson/gaussx#198
- Related: #124, #143
Contributor guide
Assessment
This issue has not been assessed yet.