Repository navigation
feat: implement gradient/Hessian for TruncatedNormalMessage - #1238
Merged
Merged
Conversation
TruncatedNormalMessage._normal_gradient_hessian previously raised NotImplementedError, which silently broke EP + Laplace optimisation for any GaussianPrior with finite limits — the EP outer loop catches the exception per factor, so runs looked successful but truncated parameters never received their cavity update. Implement the gradient and Hessian w.r.t. x: inside the truncation support these match the untruncated Gaussian (Z and log sigma are constant in x); logl additionally subtracts log Z. Outside the support, return logl = -inf and grad = 0 so the optimiser sees a flat region rather than NaNs. Override on TruncatedNaturalNormal because its mean / sigma properties return *truncated* moments (via scipy.stats.truncnorm), not the underlying Gaussian's parameters that the gradient formula needs; the override reconstructs the underlying (mu, sigma) from the natural parameters. Closes #1237 Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
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
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
Implement
TruncatedNormalMessage._normal_gradient_hessianso EP + Laplace optimisation actually updates parameters with truncated Gaussian priors. The previousNotImplementedErrorstub was caught silently by the EP outer loop, leaving truncated-prior parameters glued to their priors across the entire run.Closes #1237
API Changes
TruncatedNormalMessage.logpdf_gradientand.logpdf_gradient_hessiannow return real values instead of raisingNotImplementedError.TruncatedNaturalNormalgets the same behaviour via an override that reconstructs the underlying Gaussian from its natural parameters. No public signatures changed — this is a behaviour fix that unblocks a previously unimplemented code path.See full details below.
Test Plan
test_autofit/messages/test_truncated_normal.py(12 tests, all passing)NormalMessageinside the support;loglshifted by exactly-log Zlogl = -inf,grad = 0xshapes verifiednumerical_logpdf_gradient/numerical_logpdf_gradient_hessianto 1e-2 toleranceTruncatedNaturalNormalreturns finite gradients matching the equivalentTruncatedNormalMessagetest_autofit/mapper/prior/test_truncated_gaussian.py,test_autofit/graphical/functionality/test_messages.pypytest test_autofitregression: 1245 passed, 0 failedz_projects/concr/scripts/toy/ep.pywithTruncatedGaussianPrioronnormalizationandsigma) now runs to completion with noNotImplementedError, and per-dataset truncated-prior posteriors visibly move off their priorsFull API Changes (for automation & release notes)
Changed Behaviour
TruncatedNormalMessage._normal_gradient_hessian(x)— previously raisedNotImplementedError; now returns(logl, grad_logl, hess_logl)with the same shape conventions asNormalMessage._normal_gradient_hessian. Inside[lower_limit, upper_limit]the gradient and Hessian match the untruncated Gaussian;loglis shifted by-log ZwhereZ = Φ((b-μ)/σ) - Φ((a-μ)/σ). Outside the support,logl = -infandgrad_logl = 0.TruncatedNormalMessage.logpdf_gradient(x)and.logpdf_gradient_hessian(x)— propagate the new behaviour (no longer raise).TruncatedNaturalNormal.logpdf_gradient(x)and.logpdf_gradient_hessian(x)— propagate the new behaviour. The class now overrides_normal_gradient_hessianto reconstruct the underlying Gaussian(μ, σ)from its natural parameters, because its inheritedmean/sigmaproperties return the truncated moments viascipy.stats.truncnorm.Added (internal)
TruncatedNormalMessage._normal_gradient_hessian_from(mean, sigma, x)— extracted helper that runs the body using explicit(μ, σ). Used by both the base class and theTruncatedNaturalNormaloverride. Not part of the public API.Migration
None — this is a pure bug fix. Existing code that called these methods previously hit
NotImplementedError; that code now simply works.🤖 Generated with Claude Code