[upstream gaussx] SumKronecker solve/logdet — refactor exact multi-output GP path to structured operators
- Dominant language
- Python
- Stars
- 1
- Forks
- 0
- Avg merge
- 19h 2m
- Merged PRs (30d)
- 16
Description
## Status
Tracking issue. The actual primitive lives upstream in `gaussx` — mirrored there as **jejjohnson/gaussx#201** (structural `solve` / `logdet` dispatch for `SumKronecker`). This issue is the pyrox-side tracker so the exact multi-output GP layer has somewhere to land the structured-operator refactor once gaussx ships the dispatch.
**Blocked by jejjohnson/gaussx#201.**
## User Story
> As a pyrox-gp user fitting an exact LMC / ICM multi-output GP, I want `MultiOutputGPPrior` to hand gaussx the *structured* covariance operator (`Σ_q B_q ⊗ K_q` plus a diagonal noise shift) instead of a materialized `(P·N, P·N)` matrix, so conditioning and marginal-likelihood evaluation scale as O(P³ + N³) instead of O(P³N³).
## Motivation
- **Current state** — `MultiOutputGPPrior._prior_operator` / `_noisy_operator` in `packages/pyrox-gp/src/pyrox_gp/_multi_output_models.py` call `kernel.full_covariance(X)` (dense) and add jitter/noise to the diagonal, because gaussx's `SumKronecker` has no `solve` / `logdet` dispatch — the structure-preserving operators the kernels already expose (`LMCKernel.cross_covariance_operator` returns a sum of tagged `Kronecker`s; `ICMKernel` returns a single `Kronecker`) would silently hit gaussx's dense fallback anyway.
- **Ready and waiting** — the kernel layer already produces the right operators (`full_covariance_operator` on `LMCKernel` / `ICMKernel`), so the refactor is confined to the model layer's operator assembly; the public `MultiOutputGPPrior` / `MultiOutputConditionedGP` / `mo_gp_factor` surface does not change.
- **Scope note** — the sparse path (`MultiOutputSparseGPPrior`) already exploits structure end-to-end via `BlockDiag` dispatch and is *not* affected by this issue.
## Refactor Sketch (after gaussx#201 lands)
```python
# packages/pyrox-gp/src/pyrox_gp/_multi_output_models.py
def _noisy_operator(self, noise_var):
# BEFORE (dense):
# K = self.kernel.full_covariance(self.X) # (PN, PN) materialized
# K = K.at[jnp.diag_indices_from(K)].add(...)
# return _psd_operator(K)
# AFTER (structured):
K_op = self.kernel.full_covariance_operator(self.X) # SumKronecker / Kronecker
noise = _flat_noise(noise_var, self.X.shape[0], self.num_outputs)
return SumOperator(K_op, lx.DiagonalLinearOperator(self.jitter + noise),
tags=lx.positive_semidefinite_tag)
```
with `gaussx.solve` / `logdet` then dispatching structurally inside `log_marginal_likelihood`, `build_prediction_cache`, `predict_variance`, and `MultivariateNormal`. For the scalar-noise ICM case this is the classical Kronecker-exact GP (Saatçi 2011); for per-output noise the diagonal shift is `noise ⊗ 1_N`, which fits gaussx#201's Case-2 whitening form.
## Implementation Steps
- [ ] Bump the gaussx pin once jejjohnson/gaussx#201 is released
- [ ] Switch `MultiOutputGPPrior._prior_operator` / `_noisy_operator` to structured operator assembly (keep the dense construction as the fallback for `OILMMKernel`, whose operator is a generic `SumOperator`)
- [ ] Verify `condition` / `predict` / `log_prob` / `mo_gp_factor` route through the structured dispatch (no `DenseFallbackWarning`)
- [ ] Add an equivalence test: structured path matches the current dense path on LMC and ICM fixtures (`packages/pyrox-gp/tests/gp/test_multi_output_models.py`)
- [ ] Benchmark note in the docstring: expected scaling for ICM + scalar noise
## Definition of Done
- [ ] Exact multi-output conditioning / MLL no longer materializes `(P·N, P·N)` for LMC / ICM kernels
- [ ] Public API unchanged; all existing tests in `test_multi_output_models.py` still pass
- [ ] Tests, lint, typecheck pass: `make test && make lint && make typecheck`
## Relationships
- Blocked by: jejjohnson/gaussx#201
- Related: #145 (existing upstream-gaussx tracker precedent)
Contributor guide
Research direction
Start with the blocker jejjohnson/gaussx#201, then inspect _prior_operator and _noisy_operator in packages/pyrox-gp/src/pyrox_gp/_multi_output_models.py. After the upstream dispatch is available, run packages/pyrox-gp/tests/gp/test_multi_output_models.py and verify structured LMC/ICM conditioning and MLL without DenseFallbackWarning; finish with make test && make lint && make typecheck.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100