v0.2.0

stagereg 0.2.0

This release adds MC-Dropout as a third posterior inference method, alongside BayesByBackprop and PreconditionedSGLD. It also lets you infer a posterior over only part of a network, and adds tooling for checkpoints and evaluation.

Highlights

• MCDropout and MCDropoutPosterior: fine-tunes the inferred subnetwork with dropout turned on, and draws one dropout mask set per key. It supports early stopping and checkpoints. The XSin posterior example now demonstrates it (#92, #119).
• Inferring part of a network: the new PosteriorStage.inferred_subnetwork filter picks the subnetwork to infer, and the rest of the network stays fixed at its trained-mean values. The default still infers the whole network (#105).
• Diagnostics from inference methods: InferenceMethod.diagnostics names extra update outputs, which are recorded as posterior/train/<key> epoch means (#106).
• expected_sample_crps: computes the exact expected CRPS without a grid, batched over scoring units.
• CachedCheckpointStore: caches every save in memory and passes only the saves you select to a persistent store. Ready-made rules for which saves to keep: finalized_only, every_nth_save(n) and always.
• MeanStage and GammaVarianceStage now save a finalized checkpoint even without an early stopper, so a later process can restore them.

⚠️ Breaking changes

• Keyed posterior functions: mean_function and log_likelihood now take a PRNG key, on PosteriorProblem and on both posterior classes. A deterministic network gives the same results as before (#92).
• InferenceMethod needs a diagnostics property: a custom method can declare () to keep its current behaviour (#106).
• Posterior problems and priors cover only the inferred subnetwork. load_posterior refuses a posterior whose parameter tree doesn't match that subnetwork (#105).
• Training-state fields cleaned up:
  • TrainingState.stage is replaced by the active_stage field.
  • Checkpoint.rng_state is removed; the key is in Checkpoint.state.rng_state.
  • outer_iteration, data_fingerprint and TrainingEvent.iteration are removed.

See CHANGELOG.md for the full list.