Multi-chain, warm-started, metric-seeded BlackJAXNUTS (#1521) - #1522
Merged
Conversation
Adds a `from_result` classmethod to the existing `InitializerParamStartPoints` initializer (no new `WarmStart` class): maps a previous `Result`'s max-log-likelihood or median-PDF point onto a (possibly different) target model *by prior path*, optionally generating `n_points` physical-space jittered starting vectors (one per chain) via a seeded `RandomState`. Never reads or writes prior objects beyond `.id`/`path_for_prior` -- scientific priors are left untouched, unlike `Result.model_centred*`. Also adds `Result.start_point` / `Result.start_point_from(...)` convenience wrappers next to `model_centred`, with a docstring contrast against the prior-rewriting `model_centred*` family. Addresses PyAutoFit#1521 (part A). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EDABYoH6giHXhFJUks8yd6
Removes the num_chains=1-only NotImplementedError. Chains are sampled in parallel via jax.vmap: warmup (blackjax.window_adaptation) and the NUTS step (blackjax.nuts.build_kernel(), vmapped over per-chain step_size / inverse_mass_matrix rather than the fixed-metric blackjax.nuts(...)) both batch over a leading chains axis, so search_internal["positions"] is now always (n_samples, n_chains, n_dim), infos carry a matching chain axis, and num_trajectory_expansions (tree depth) is recorded alongside acceptance rate / integration steps / divergences. New kwargs: num_chains (now usable), inverse_mass_matrix (None / "diagonal" / "dense" / an explicit seed array / a Result or Samples object -- seeding from a Result's covariance_matrix raises a clear error if it looks MLE-only), mass_matrix_shrinkage (blackjax's imm_shrinkage_to_previous), and share_adaptation (average the per-chain tuned metrics into one shared kernel). The resolution/shape helpers (resolve_inverse_mass_matrix, inverse_mass_matrix_kind_from, stack_initial_positions, split_chain_diagnostics) live in a new sampler-agnostic blackjax/chains.py module, pure numpy (no jax/blackjax import at module scope) so a future second BlackJAX sampler can reuse them. __identifier_fields__ gains "inverse_mass_matrix" (its string kind); apply_test_mode now also clamps num_chains to at most 2. samples_info_from gains blackjax.diagnostics-derived multi-chain diagnostics (ess_bulk/tail per param + mins, rank-normalised rhat per param + max, divergent_indices, tree_depth_histogram, warm_start_source, inverse_mass_matrix_kind) alongside the existing ess_per_param/ess_min. SamplesMCMC gains matching read-only properties (ess_bulk, ess_tail, rhat, n_divergent, tree_depths), documented as split-chain-only when num_chains == 1. Addresses PyAutoFit#1521 (parts B/C). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EDABYoH6giHXhFJUks8yd6
…inx baseline) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EDABYoH6giHXhFJUks8yd6
Jammy2211
pushed a commit
to trexfr-ops/PyAutoFit
that referenced
this pull request
Aug 27, 2026
Extends the linear diagonal AbstractScaler to a per-coordinate bijection (identity / log / logit), vectorised over (n_starts, n_params) and built once per model via from_model() so jax.jit traces a single jnp.where selection tree. BijectorNone is the no-op default (byte-identical: x/1.0 == x exactly); BijectorAuto picks log for LogUniform/LogGaussian priors; BijectorLogit is an explicitly secondary arm (the scaler's unit-cube objections about boundary optima still apply); BijectorPerPath lets a caller declare kinds by model path; BijectorDiagonal adapts an existing AbstractScaler into this framework without changing scaler= behaviour. log_det_jacobian is exposed for a future phi-space sampler (PyAutoFit PyAutoLabs#1521/PyAutoLabs#1522) and is never needed for MAP: composing the objective through a bijection relabels points without changing its value set, so no Jacobian may be added without moving the MAP. Compose AbstractClipper.project with a bijector, clipping in transformed space against bijector.bounds_forward(lower_inset, upper_inset) -- valid because every kind is monotone increasing. Fixes a real bug found while wiring this up: the existing physical-relative margin (margin * (upper - lower)) is wrong for a log-kind coordinate by orders of magnitude -- for LogUniform(1e-6, 1e6) the default margin produces a ~1.0 PHYSICAL inset, fencing off virtually the entire support. Log-kind coordinates now get a log-space margin (margin * log(upper/lower)) via ClipperPriorBox._inset_from_model(kinds=...); the no-bijector bounds_from_model path is unaffected. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EDABYoH6giHXhFJUks8yd6
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Phase 6 pre-req of the autolens_profiling inference programme (PROGRAMME §4 Phase 6; Gate C = batched-pipeline value). Closes #1521.
InitializerParamStartPoints.from_result(result, model=None, point="max_log_likelihood"|"median_pdf", n_points=1, jitter=0.0, seed=0)maps the previous result's vector onto the target model by prior path (sourced fromsamples.model, so ordering is always consistent).Result.start_point(property) andResult.start_point_from(...)are thin conveniences besidemodel_centred*— explicitly non-mutating, unlike prior passing.BlackJAXNUTS:num_chains(vmapped warmup +blackjax.nuts.build_kernel()vmapped over per-chain step size / metric),inverse_mass_matrix ∈ {None, "diagonal", "dense", ndarray, Result, Samples}(a Result/Samples seeds a dense metric fromsamples.covariance_matrix; MLE-only sources raise and ask for an explicit Laplace array),mass_matrix_shrinkage(→imm_shrinkage_to_previous),share_adaptation.inverse_mass_matrixjoins__identifier_fields__; test mode clamps chains to 2.diagnostics): rank-normalisedess_bulk/ess_tailper param + min, split-R̂ per param + max,n_divergent+ indices, tree-depth histogram (num_trajectory_expansions— the H6.2 lockstep measurement), mean acceptance,warm_start_source,inverse_mass_matrix_kind, all insamples_infowith read-onlySamplesMCMCproperties.blackjax/chains.pyfor later ChEES/MAMS arms.Tests
test_autofitfull suite: 2101 passed, 3 skipped. New:test_initializer_from_result.py(9),test_blackjax_chains.py(23, pure numpy), construction/identifier/shape-contract tests intest_blackjax_nuts.py,Result.start_pointtest. No real JAX sampling loops in the unit suite (repo convention); the end-to-end 2-chain fit goes to the workspace_test script.Follow-ups
Low-rank metric (
window_adaptation_low_rank) and the ChEES/MAMS arms are out of scope here.🤖 Generated with Claude Code
https://claude.ai/code/session_01EDABYoH6giHXhFJUks8yd6