From e7de318875907db26f30580609fdc1817ee2450c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 17 Jul 2026 13:43:14 +0100 Subject: [PATCH] =?UTF-8?q?docs:=20pin=20Fitness's=20NaN=20guard=20contrac?= =?UTF-8?q?t=20=E2=80=94=20value-only,=20never=20gradient?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The xp.where NaN/inf guard in Fitness.call repairs the VALUE but not the gradient: under jax.grad, reverse-mode differentiates the masked branch and computes 0 * NaN = NaN whenever the likelihood's derivative is also non-finite. Searches reading only the figure of merit are protected; gradient consumers are not. A guard here cannot fix this — by the time call() receives log_likelihood the non-finite derivative is already on the tape, so an output-side double-where is measurably ineffective. Gradient-safety belongs at the site that creates the NaN. Docstring + comment only; no behaviour change. Diagnosed in autolens_workspace_developer#104; pinned by autofit_workspace_test/scripts/jax_assertions/fitness_nan_gradient_contract.py. Co-Authored-By: Claude Fable 5 --- autofit/non_linear/fitness.py | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/autofit/non_linear/fitness.py b/autofit/non_linear/fitness.py index ce49b81a7..8962ebdb8 100644 --- a/autofit/non_linear/fitness.py +++ b/autofit/non_linear/fitness.py @@ -207,6 +207,27 @@ def call(self, parameters): A private method that calls the fitness function with the given parameters and additional keyword arguments. This method is intended for internal use only. + The NaN/inf guard below protects the **value only, never the gradient**. + + A model whose likelihood is NaN or inf is mapped to `resample_figure_of_merit`, so searches that read only + the figure of merit (nested samplers, MCMC) get the resample sentinel and never select the point. Gradient + consumers get no such protection: under `jax.grad`, reverse-mode differentiates *both* branches of an + `xp.where` and multiplies the unselected one by zero, so if the likelihood's derivative is also non-finite + the guard yields `0 * NaN = NaN` and the returned gradient is NaN even though the value looks handled. + + This bites only when the masked branch's *derivative* is non-finite — not merely its value. `sqrt(x)` at + x < 0 and `cholesky(A)` for non-positive-definite `A` are NaN in both value and derivative, so they trigger + it; `log(x)` at x < 0 is NaN in value but its derivative `1/x` stays finite, so it does not. + + A guard here **cannot** repair this. By the time this method receives `log_likelihood` the non-finite + derivative is already recorded on the autodiff tape, and no transformation of the output can remove it — + an output-side "double-where" does not work. Gradient-safety must be established at the site that creates + the NaN, by never *evaluating* the offending operation at the invalid input. + + See autolens_workspace_developer#104, where this was diagnosed, and + `autofit_workspace_test/scripts/jax_assertions/fitness_nan_gradient_contract.py`, which pins the behaviour + described here. + Parameters ---------- parameters @@ -235,7 +256,9 @@ def call(self, parameters): except exc.FitException: return self.resample_figure_of_merit - # Penalize NaNs in the log-likelihood + # Penalize NaNs in the log-likelihood. Value-only: under jax.grad these `where`s still differentiate the + # masked branch, so a non-finite derivative propagates as `0 * NaN = NaN`. See the contract in the + # docstring above -- gradient-safety belongs at the site that creates the NaN, not here. log_likelihood = self._xp.where(self._xp.isnan(log_likelihood), self.resample_figure_of_merit, log_likelihood) log_likelihood = self._xp.where(self._xp.isinf(log_likelihood), self.resample_figure_of_merit, log_likelihood)