jejjohnson / jejjohnson/pyrox

[upstream gaussx] Diagonal-array overloads for cavity_distribution and newton_update

Open
#145 1 comment 0 reactions 0 assignees View on GitHub
area:gp area:integration enhancement type:feature upstream-gaussx
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.