diff --git a/scripts/jax_assertions/fitness_nan_gradient_contract.py b/scripts/jax_assertions/fitness_nan_gradient_contract.py new file mode 100644 index 0000000..c5a1297 --- /dev/null +++ b/scripts/jax_assertions/fitness_nan_gradient_contract.py @@ -0,0 +1,152 @@ +""" +JAX Assertions — Fitness's NaN guard contract +============================================= + +Pins what ``Fitness``'s NaN/inf resample guard (``autofit/non_linear/fitness.py``, +in ``call``) actually promises: it protects the **value**, never the **gradient**. + +Why this exists (PyAutoFit#1391, diagnosed in autolens_workspace_developer#104): +the guard maps a NaN/inf likelihood to ``resample_figure_of_merit``, so searches +that read only the figure of merit never select the point. Gradient consumers get +nothing: ``jax.grad`` differentiates *both* branches of an ``xp.where`` and +multiplies the unselected one by zero, so a non-finite derivative propagates as +``0 * NaN = NaN``. The value looks handled; the gradient is poison. + +These assertions exist to stop that limitation silently regressing into a false +promise — and to record the disproof of the fix that looks obvious but does not +work, so nobody re-prescribes it. An output-side "double-where" **cannot** help: +by the time ``Fitness.call`` sees ``log_likelihood`` the non-finite derivative is +already on the tape. Only refusing to *evaluate* the offending op at the invalid +input works, and that decision lives at the NaN's source (for the pixelized +likelihood, autoarray's Cholesky), not in ``Fitness``. + +These are jax-only assertions, so they live here rather than in ``test_autofit/`` +(the library's unit tests are numpy-only). + +Run from the workspace root: + + python scripts/jax_assertions/fitness_nan_gradient_contract.py +""" + +import jax +import jax.numpy as jnp + +jax.config.update("jax_enable_x64", True) + +RESAMPLE = -jnp.inf + + +def _nan_derivative_likelihood(x): + """NaN in value AND derivative for x < 0 -- the shape of a real failure. + + ``sqrt`` stands in for the site that actually bites in production: a Cholesky + of a non-positive-definite curvature/regularization matrix, which JAX returns + as NaN rather than raising (autolens_workspace_developer#104). + """ + return jnp.sqrt(x) + + +def _finite_derivative_likelihood(x): + """NaN in value but FINITE derivative (d/dx log = 1/x) for x < 0.""" + return jnp.log(x) + + +def _guarded(likelihood): + """The guard as `Fitness.call` applies it (value path only).""" + + def call(x): + ll = likelihood(x) + ll = jnp.where(jnp.isnan(ll), RESAMPLE, ll) + ll = jnp.where(jnp.isinf(ll), RESAMPLE, ll) + return ll * -2.0 # convert_to_chi_squared + + return call + + +def assert_guard_repairs_the_value(): + """The guard delivers what it advertises: no NaN reaches the search.""" + call = _guarded(_nan_derivative_likelihood) + fom = call(-1.0) + assert not jnp.isnan(fom), "the guard must never let a NaN value through" + assert jnp.isinf(fom), "an invalid model's figure of merit is the resample sentinel" + + +def assert_guard_does_not_repair_the_gradient(): + """The contract's limit: a guarded point still yields a NaN gradient. + + This assertion is deliberately pinning *broken-looking* behaviour. If it ever + fails, the guard's reach has changed and the docstring in `Fitness.call` (plus + every caller relying on "resampled points are safe") must be revisited. + """ + call = _guarded(_nan_derivative_likelihood) + grad = jax.grad(call)(-1.0) + assert jnp.isnan(grad), ( + "expected the where-guard to leak a NaN gradient (0 * NaN); if this now " + "passes, the value-only contract has changed -- update fitness.py's docstring" + ) + + +def assert_gradient_is_fine_when_the_masked_derivative_is_finite(): + """The trap is narrower than 'the guard never protects gradients'. + + It fires only when the masked branch's DERIVATIVE is non-finite. `log(x<0)` is + NaN in value but its derivative 1/x is finite, so `0 * -1 = 0` and nothing + leaks. Recording this stops the limitation being overstated. + """ + call = _guarded(_finite_derivative_likelihood) + grad = jax.grad(call)(-1.0) + assert jnp.isfinite(grad), ( + "a finite masked derivative must not trigger the 0 * NaN trap" + ) + + +def assert_output_side_double_where_does_not_fix_it(): + """The disproof. Sanitising the OUTPUT is too late -- the tape already has NaN.""" + + def call(x): + ll_raw = _nan_derivative_likelihood(x) + bad = jnp.isnan(ll_raw) | jnp.isinf(ll_raw) + safe = jnp.where(bad, 0.0, ll_raw) # sanitise after the likelihood ran + return jnp.where(bad, RESAMPLE, safe) * -2.0 + + grad = jax.grad(call)(-1.0) + assert jnp.isnan(grad), ( + "an output-side double-where must NOT be presented as the remedy: the " + "non-finite derivative is already on the tape by then" + ) + + +def assert_input_side_safe_x_does_fix_it(): + """What a real fix looks like -- and why it cannot live in `Fitness`. + + The likelihood is never *evaluated* at the invalid input. That requires knowing + which input is invalid, which is knowable only at the NaN's source, not at the + Fitness boundary where the likelihood is an opaque already-computed scalar. + """ + + def call(x): + bad = x < 0 + safe_x = jnp.where(bad, 1.0, x) # sanitise the INPUT + ll = _nan_derivative_likelihood(safe_x) + return jnp.where(bad, RESAMPLE, ll) * -2.0 + + grad = jax.grad(call)(-1.0) + assert jnp.isfinite(grad), "input-side safe-x must yield a finite gradient" + assert grad == 0.0, "an invalid point carries no gradient information" + + +def assert_valid_points_are_unaffected(): + """None of the above perturbs a well-behaved model.""" + call = _guarded(_nan_derivative_likelihood) + assert jnp.isfinite(call(4.0)) + assert jnp.isfinite(jax.grad(call)(4.0)) + + +if __name__ == "__main__": + assert_guard_repairs_the_value() + assert_guard_does_not_repair_the_gradient() + assert_gradient_is_fine_when_the_masked_derivative_is_finite() + assert_output_side_double_where_does_not_fix_it() + assert_input_side_safe_x_does_fix_it() + assert_valid_points_are_unaffected() + print("fitness_nan_gradient_contract: all assertions passed")