Per-parameter bijector for gradient searches (identity / log / logit) — Phase 8B library half (autolens_profiling#162) - #1525
Merged
Conversation
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 #1521/#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
Adds bijector: Optional[AbstractBijector] = None alongside scaler, mutually exclusive (raises ValueError at construction if both are non-default). Wherever the step loop maps between physical and stepped coordinates for the scaler (resume load, fresh-start seeding, best/ lane-best capture, dead-lane redraw, the write-back), a parallel bijector.forward/inverse branch is added; search_internal["params"] stays PHYSICAL either way so resume is safe across a scaler/bijector change. search_internal also gains "bijector": bijector.kinds, and samples_info records the bijector's class name and (from search_internal, not live state) its resolved kinds -- mirroring how "scaler"/"clipper" are already reported. AbstractClipper.project's new `bijector=` param is threaded through the clip call and the dead-lane redraw. Everything is gated behind the same has_scaler/has_bijector-style short-circuit already used for the scaler/clipper/constraint checks, so the default (BijectorNone) step loop and its jax.jit trace are unchanged. Also adds two independently flag-gated, off-by-default diagnostics (neither touches the compiled step when unused): record_lane_nan_history (per-step per-lane value/grad-NaN bits, numpy.packbits'd to (n_steps, ceil(n_starts/8))) and trace_param_indices (a physical per-step trace for a chosen subset of parameters, (n_steps, n_starts, k)). Both resume via the same `.get`-default discipline as the existing lifetime counters. Exports the Bijector* classes from autofit/__init__.py beside the Scaler* ones. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EDABYoH6giHXhFJUks8yd6
test_bijector.py (numpy-only, mirroring test_scaler.py's discipline): round trip and monotonicity per kind, the log-kind closed-form log-det-Jacobian checked against a finite difference of inverse() directly, kind-selection/fallback coverage for BijectorAuto/ BijectorLogit/BijectorPerPath (including a malformed LogUniformPrior), BijectorDiagonal reducing exactly to ScalerPriorWidth's map, BijectorNone's byte-identical no-op, the objective-composed-through-the- bijector equivalence pin and its end-to-end "cannot move the MAP" companion (L-BFGS-B in raw theta space vs. unconstrained BFGS in phi space -- the raw run needs a box because (log(theta)-2)^2 is undefined for theta <= 0 and unconstrained steps stray there, which is itself the motivating case for the log bijector), and clip-commutes-with-forward. test_clipper.py: scale+bijector together raises, ClipperNone ignores a bijector, the LogUniform log-space inset fix pinned directly against the bug (a LogUniform(1e-6, 1e6) physical-relative inset lands near 1.0, fencing off virtually the whole support; the log-space inset stays within 0.1% of the true bound), project(..., bijector=...) matching a direct physical clip mapped through, and the no-bijector bounds_from_model path proven byte-unchanged. test_multi_start_gradient.py: bijector dict round-trip, the scaler+bijector construction-time raise, bijector not offered by LBFGS, and samples_info recording the bijector's class name and (read from search_internal, not live resolved state) its kinds. Full suite green: test_autofit/non_linear (673 passed, 2 skipped) and the whole test_autofit suite (2167 passed, 3 skipped). ruff clean on every touched file except autofit/__init__.py, where the new Bijector* re-export lines repeat that file's pre-existing F401/E402 baseline (258 -> 270 errors, +12 for the 12 new import lines) rather than adding a new category of lint issue. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EDABYoH6giHXhFJUks8yd6
Extends the existing estim_lr/d progress-line comment (already noting d is in SCALED units under a scaler, not physical ones) to cover the bijector case explicitly: under a mixed identity/log bijector, d is a single global scalar spanning coordinates that are not even the same KIND of unit across parameters. Prodigy's own estimation rule (one global d from whole-tree norms) is unmodified by either knob -- this is existing, documented Prodigy behaviour, not something this change touches or attempts to fix. 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
Library half of autolens_profiling #162 (W5 · Phase 8B: log-coordinate stepping). The profiling A/B driver follows once this merges (library-first gate).
autofit/non_linear/bijector.py—AbstractBijector(forward/inverse/log_det_jacobian/bounds_forward/kinds), per-coordinate kinds baked once infrom_modelso JAX traces onewherechain (double-where safe surrogate, no NaN grads from the unused branch).BijectorNone(default, byte-identical step),BijectorAuto(log forLogUniform/LogGaussian, identity elsewhere),BijectorPerPath,BijectorLogit(secondary arm; the scaler-era unit-cube objections stand and are documented),BijectorDiagonal(scaler-equivalent).log_det_jacobianexists for φ-space samplers (multi-chain NUTS, Multi-chain, warm-started, metric-seeded BlackJAXNUTS (#1521) #1522) and is never called by MultiStart.project(..., bijector=)), valid because every kind is monotone. Bug fix:ClipperPriorBox's relative insetmargin·(hi−lo)≈ 1.0 physical forLogUniform(1e-6, 1e6)— under a log coordinate that would fence off all λ < 1. Log-kind coordinates now inset in log space (margin·log(hi/lo)), pinned by a test.bijector=besidescaler=(both non-default →ValueError);search_internal["params"]stays physical so resume across bijector on/off is safe;samples_inforecords class + kinds. Two flag-gated, off-by-default per-step diagnostics (record_lane_nan_history,trace_param_indices) — no-ops on the compiled step when unused.Tests
Full suite: 2167 passed, 3 skipped. New
test_bijector.py(27),test_clipper.py+6,test_multi_start_gradient.py+6. Pre-merge JAX smoke:MultiStartProdigyon a 3-param model with aLogUniformamplitude —BijectorNoneandBijectorAutoreach the same MAP (centre 1.000, σ 0.800, amp 300.0).ruffclean on touched files (the 12 new__init__re-export lines match that file's existing F401/E402 baseline).Notes
Prodigy's global
dis unchanged (one scalar over mixed units, exactly as under a scaler) — documented, to be measured by the profiling A/B's MGE byte-identity control and F4.🤖 Generated with Claude Code
https://claude.ai/code/session_01EDABYoH6giHXhFJUks8yd6