jejjohnson / jejjohnson/pipekit

feat(pipekit-train): NumPyro adapter gaps blocking unsupervised Bayesian fits — kernel kwargs, unsupervised NumpyroTask, full-posterior predictive

Open
#51 0 comments 0 reactions 0 assignees View on GitHub
area:code priority:p2 type:feature
Dominant language
Python
Stars
0
Forks
0
PR merge metrics
No merged PRs in 30d

Description

## Problem / Request

A reuse audit of xtremax (JAX/NumPyro extreme-value library) against pipekit-train found that EVT workflows *cannot* adopt `TrainingLoop(backend="numpyro-mcmc")` today, for three concrete adapter reasons. These generalize to any unsupervised Bayesian fitting (density estimation, EVT, unsupervised latent models), so fixing them widens pipekit-train's NumPyro seam beyond supervised regression.

1. **No kernel configuration**: `NumpyroTask` exposes only `num_warmup/num_samples/num_chains`; `adapters/numpyro_mcmc.py:56-62` builds `NUTS(task.model)` bare. Real EVT fits need `NUTS(target_accept_prob=0.995, init_strategy=init_to_median, max_tree_depth=12)` — currently unreachable.
2. **Supervised-shaped task**: `adapters/bayes.py::materialize` stacks `(x, y)` pairs and calls `model(x, y)`. Unsupervised fits are `model(data, priors...)` with no `(x, y)` split — today one must fabricate a dummy-x dataset.
3. **Point-estimate predictive**: `NumpyroPredictiveOp` returns the posterior-predictive **mean**; EVT users (and Bayesian users generally) need the full posterior / posterior-predictive draws for credible intervals on derived quantities (e.g. return levels).

## User Story

> As a pipekit-train user fitting an unsupervised NumPyro model, I want to configure the MCMC kernel, pass unpaired data, and retrieve full posterior draws, so that the adapter covers Bayesian workflows beyond supervised regression.

## Proposed API

```python
@dataclass
class NumpyroTask:
model: Callable
kernel_kwargs: dict[str, Any] = field(default_factory=dict) # → NUTS(model, **kernel_kwargs)
kernel_factory: Callable[..., MCMCKernel] | None = None # full override
supervised: bool = True # False → model(batch, **model_kwargs)
...

class NumpyroPredictiveOp(...):
def posterior_samples(self) -> dict[str, Array]: ...
def predictive(self, key, *args, num_samples=None) -> dict[str, Array]: ... # full draws
# existing mean-predict stays as the Operator __call__ surface
```

## References & Existing Code

- `packages/pipekit-train/src/pipekit_train/adapters/numpyro_mcmc.py:56-62`
- `packages/pipekit-train/src/pipekit_train/adapters/bayes.py` (`materialize`, `NumpyroTask`, `NumpyroPredictiveOp.serialize_weights`)
- Motivating downstream: jejjohnson/xtremax GEV/GPD block-maxima + POT notebooks (raw `MCMC(NUTS(model, target_accept_prob=0.995, ...))` today)
- Same kernel-kwargs consideration applies to `numpyro_svi.py` (optimizer/guide kwargs) and `blackjax.py` for symmetry

## Implementation Steps

- [ ] Add `kernel_kwargs`/`kernel_factory` to `NumpyroTask`; thread through `numpyro_mcmc.py` (and the SVI analogue)
- [ ] Add the unsupervised path in `bayes.materialize` (no (x, y) stacking; pass the dataset batch directly)
- [ ] Expose full posterior/predictive draws on `NumpyroPredictiveOp`; keep weight-blob serialization intact
- [ ] Round-trip test: unsupervised model fit → registry store → load → `posterior_samples()`

## Definition of Done

- [ ] An unsupervised NumPyro model fits through `TrainingLoop(backend="numpyro-mcmc")` with a tuned NUTS kernel, and full posterior draws are retrievable after a registry round-trip
- [ ] Workspace pre-commit checklist green (pytest, ruff, ty)

## Testing

- [ ] Unit test: kernel kwargs reach NUTS (spy/introspection)
- [ ] Integration test: unsupervised toy model (e.g. Normal location-scale) end-to-end

## Documentation

- [ ] Adapter docs: supervised vs unsupervised task shapes; kernel configuration example

## Relationships

- Parent (theme epic): #
- Related: jejjohnson/xtremax#42 (ecosystem-alignment epic tracking the downstream side)

Contributor guide

Open the contributing guide

Research direction

Start with adapters/bayes.py, especially materialize, NumpyroTask, and NumpyroPredictiveOp.serialize_weights, then read adapters/numpyro_mcmc.py:56-62 and the numpyro_svi.py analogue. Add tests for kernel-kwargs forwarding and an unsupervised toy model with a registry round-trip. Done means tuned unsupervised MCMC works, posterior draws reload successfully, and pytest, ruff, and ty pass.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.