jejjohnson / jejjohnson/pipekit
feat(pipekit-train): NumPyro adapter gaps blocking unsupervised Bayesian fits — kernel kwargs, unsupervised NumpyroTask, full-posterior predictive
- 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
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