docs: pin Fitness's NaN guard contract — value-only, never gradient

[Jammy2211]

Overview

Fitness.call (autofit/non_linear/fitness.py:238-240) guards a bad likelihood with xp.where(xp.isnan(ll), resample_figure_of_merit, ll). It repairs the value — confirmed: autolens_workspace_developer#104's rejected draws report loss = inf, not nan, so the guard fires. But under jax.grad, reverse-mode AD differentiates the masked branch and multiplies by zero, so 0 * NaN = NaN — the gradient is NaN regardless of the guard. This retroactively explains two #100/#101 mysteries: trajectories died with a non-finite objective rather than a NaN one, and optax.apply_if_finite latched at the cliff instead of stepping past it.

Investigation established that Fitness.call is structurally incapable of fixing this, so this issue pins the guard's true contract and proves the limitation. It writes no code fix — the real fix belongs at the NaN source (the autoarray phase-2 issue).

Plan

• Document the guard's real contract on Fitness.call: it protects the value, never the gradient.
• Add a JAX assertion script pinning that contract — including the disproof of the remedy that looks obvious but fails.
• Correct the findings doc shipped in autolens_workspace_developer#105, which carries two claims this investigation disproved.
• Write no code fix; gradient-safety must be established at each NaN source.

Detailed implementation plan

What the investigation established (do not re-derive)

Proved on CPU in seconds (~20 lines of JAX — no A100 needed; an earlier draft wrongly specified #104's pixelized probe, which needs ~11 GiB):

1. The trap is narrower than "the guard never protects gradients." It fires only when the masked branch's derivative is non-finite:

primal at a bad point	value	derivative	trap fires?
log(x), x<0	NaN	1/x = finite	no — 0 * -1 = 0
sqrt(x), x<0	NaN	NaN	yes
cholesky(A), non-PD (#104's site)	NaN	NaN	yes

2. An output-side double-where does NOT fix it. Measured:

current (fitness.py:239-240)     x=-1.0  value=inf  grad=nan
output-side double-where         x=-1.0  value=inf  grad=nan   <-- still NaN
input-side safe-x                x=-1.0  value=inf  grad=0.0   <-- works

By the time Fitness.call receives log_likelihood, the NaN derivative is already on the tape. Sanitising the output cannot repair it — only avoiding the evaluation at the bad input works, and that decision lives deep inside the likelihood (autoarray's cholesky), not here.

3. Therefore the fix cannot live in Fitness.call. The guard advertises protection it cannot provide to gradient consumers.

Work Classification

Both → library first (PyAutoFit), workspaces follow.

Worktree root

~/Code/PyAutoLabs-wt/fitness-nan-guard-contract/

Affected Repositories

• PyAutoFit (primary) — docstring only, no behaviour change
• autofit_workspace_test — the JAX assertion script
• autolens_workspace_developer — correct the merged findings doc

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean
./autofit_workspace_test	main	clean
./autolens_workspace_developer	main	23 modified (pre-existing black reformats, unrelated)

Suggested branch: feature/fitness-nan-guard-contract

Implementation Steps

1. PyAutoFit — autofit/non_linear/fitness.py, docstring on call: the guard maps a NaN/inf likelihood to resample_figure_of_merit in the value only; under jax.grad the masked branch is still differentiated, so when the likelihood's derivative is also non-finite the gradient is NaN; gradient-safety must be established at the NaN source, because an output-side guard cannot repair a tape that already carries NaN. Cross-reference autolens_workspace_developer#104. No behaviour change.
2. autofit_workspace_test — new scripts/jax_assertions/fitness_nan_gradient_contract.py, mirroring the sibling fitness_dispatch.py (same af.ex.Gaussian + Fitness construction). On a NaN-derivative likelihood assert: the value is guarded; jax.grad is NaN (pins the limitation so it cannot silently regress into a false promise); the output-side double-where does not fix it; input-side safe-x does (grad 0.0). Unit tests stay numpy-only per repo policy — this belongs here, not test_autofit/.
3. autolens_workspace_developer — correct searches_minimal/pix_nonfinite_findings.md: the over-strong "does not protect gradient consumers at all" claim, and the double-where remedy line (disproved). Replace this bug's "needs an A100" note with the CPU repro.

Key Files

• autofit/non_linear/fitness.py:205-254 — Fitness.call; the guard at 238-240.
• autofit_workspace_test/scripts/jax_assertions/fitness_dispatch.py — the sibling to mirror (JAX-only assertions live here precisely because unit tests are numpy-only).
• autolens_workspace_developer/searches_minimal/pix_nonfinite_findings.md — the doc to correct.

Testing

pytest test_autofit/ (docstring-only — expect no delta); run the new assertion script directly. CPU only, no A100.

Out of scope

A search-level gradient sanitiser (where(isfinite(g), g, 0)). It fabricates a zero gradient at invalid points, and #101 already showed optax.apply_if_finite latches at the cliff. The principled response to an invalid point is restart/resample, which the multistart-resurrection prompt covers.

Original Prompt

Click to expand starting prompt

# Pin the real contract of Fitness's NaN resample guard (value, never gradient)

Type: bug
Target: autofit
Repos:
- @PyAutoFit
- @autofit_workspace_test
- @autolens_workspace_developer
Difficulty: small
Autonomy: supervised
Priority: normal
Status: draft

Found while localising the pixelized NaN (autolens_workspace_developer#104,
shipped via #105).

`Fitness.call` (`autofit/non_linear/fitness.py:238-240`) penalises a bad
likelihood with:

```python
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)

It repairs the value — confirmed empirically: #104's rejected draws report
loss = inf, not nan, so the guard fires. But under jax.grad, reverse-mode
AD differentiates the masked branch and multiplies by zero, so 0 * NaN = NaN:
the gradient is NaN regardless of the guard. This retroactively explains two
#100/#101 mysteries — trajectories died with a non-finite objective rather than a
NaN one, and optax.apply_if_finite LATCHED at the cliff instead of stepping
past it (unguarded runs died silently — #100's -39888).

What the investigation established (do not re-derive)

Three things were proved on CPU in seconds (~20 lines of JAX; no A100 needed —
an earlier draft of this prompt wrongly specified the #104 pixelized probe, which
needs ~11 GiB):

1. The trap is narrower than "the guard never protects gradients". It fires
   only when the masked branch's derivative is non-finite. sqrt(x) at x<0
   NaNs its derivative -> trap fires. log(x) at x<0 is NaN in value but its
   derivative 1/x stays finite -> 0 * -1 = 0, no trap. Cholesky on a non-PD
   matrix (Feature/eden #104's actual site) NaNs its derivative, so the trap fires there.

2. An output-side double-where DOES NOT FIX IT. Measured:

   current (fitness.py:239-240)     x=-1.0  value=inf  grad=nan
   output-side double-where         x=-1.0  value=inf  grad=nan   <-- still NaN
   input-side safe-x                x=-1.0  value=inf  grad=0.0   <-- works
   

   By the time Fitness.call receives log_likelihood, the NaN derivative is
   already on the tape. Sanitising the output cannot repair it. Only avoiding the
   evaluation of the likelihood at the bad input works — and that decision lives
   deep inside the likelihood (autoarray's cholesky), not in Fitness.call.

3. Therefore Fitness.call is structurally incapable of fixing this. The
   guard advertises protection it cannot provide to gradient consumers. The real
   fix belongs at each NaN source — for the pixelized case, that is
   draft/bug/autoarray/reg_matrix_logdet_nonfinite_fix.md (phase 2).

Task: contract + test, NOT a code fix

Human decision 2026-07-17: pin the true contract; write no code fix here.

1. PyAutoFit — docstring on Fitness.call stating the guard's real contract:
   value-only; gradient-safety must be established at the NaN source; a
   downstream guard cannot repair a NaN-carrying tape. Note the
   masked-derivative condition (point 1 above). No behaviour change. Cross-ref
   autolens_workspace_developer#104.

2. autofit_workspace_test — new
   scripts/jax_assertions/fitness_nan_gradient_contract.py, mirroring the
   sibling fitness_dispatch.py (same af.ex.Gaussian + Fitness setup).
   Assert on a NaN-derivative likelihood: the value is guarded; jax.grad
   is NaN (pins the limitation so it cannot silently regress into a false
   promise); the output-side double-where does not fix it (the disproof, so
   nobody re-prescribes it); input-side safe-x does. Unit tests stay
   numpy-only per repo policy — this belongs here, not in test_autofit/.

3. autolens_workspace_developer — correct
   searches_minimal/pix_nonfinite_findings.md, which shipped in Feature/break promises #105 carrying
   two errors from before this investigation: the over-strong "does not protect
   gradient consumers at all" claim, and the double-where remedy line (now
   disproved). Replace the "needs an A100" reproduction note for this specific
   bug with the CPU repro.

Explicitly out of scope: a search-level gradient sanitiser
(where(isfinite(g), g, 0)). It fabricates a zero gradient at invalid points, and
#101 already showed optax.apply_if_finite latches at the cliff; the principled
response to an invalid point is restart/resample, which
draft/feature/autofit/multistart_resurrection_and_contrib_rules.md covers.

</details>

[Jammy2211]

Library PR Created: #1392 (pending-release)

Docstring + comment only, no behaviour change. pytest test_autofit/ → 1499 passed / 1 skipped (no delta, as expected).

API Changes: none — so there is no downstream migration. The workspace work here is not migration but the deliverable itself (option (i), new coverage):

• autofit_workspace_test — scripts/jax_assertions/fitness_nan_gradient_contract.py, pinning the contract (value guarded; jax.grad NaN; output-side double-where disproved; input-side safe-x works). All 6 assertions pass on CPU in seconds.
• autolens_workspace_developer — correcting searches_minimal/pix_nonfinite_findings.md, which shipped in Feature/break promises #105 carrying the two claims this investigation disproved.

Both follow via /ship_workspace behind the library-first merge gate.

[Jammy2211]

Shipped — all three PRs merged in library-first order: PyAutoFit#1392 → autofit_workspace_test#53 → autolens_workspace_developer#106.

The task inverted under investigation. It was filed as "fix the where guard"; the investigation proved the guard cannot be fixed where it lives — by the time Fitness.call receives log_likelihood, the non-finite derivative is already on the autodiff tape, so an output-side double-where measurably does nothing:

current (fitness.py:239-240)     x=-1.0  value=inf  grad=nan
output-side double-where         x=-1.0  value=inf  grad=nan   <-- does NOT help
input-side safe-x                x=-1.0  value=inf  grad=0.0   <-- works

So the deliverable is a pinned contract, not a patch: a docstring stating the guard is value-only, never gradient, and six assertions in fitness_nan_gradient_contract.py. The load-bearing one deliberately pins broken-looking behaviour — if assert_guard_does_not_repair_the_gradient ever fails, the guard's reach changed and the docstring is lying.

Two claims from the original #104 write-up were disproved and corrected on main (#106): the "does not protect gradient consumers at all" overreach (the trap needs the masked branch's derivative to be non-finite — sqrt/cholesky yes, log no), and the prescribed double-where remedy. The reproduction is ~20 lines of CPU JAX, not an A100 run.

Gradient-safety belongs at each NaN source. For the pixelized case that's draft/bug/autoarray/reg_matrix_logdet_nonfinite_fix.md.

Completion record: PyAutoMind/complete/2026/07/fitness-nan-guard-contract.md. Closing.