test: pin Fitness's NaN guard contract under jax.grad

[Jammy2211]

Part of PyAutoFit#1391. Upstream PR: PyAutoLabs/PyAutoFit#1392 (must merge first — library-first gate).

Pins what Fitness's NaN/inf resample guard actually promises: it protects the value, never the gradient.

Why

Fitness.call guards with xp.where(xp.isnan(ll), resample_figure_of_merit, ll). Searches reading only the figure of merit are protected. Gradient consumers are not: jax.grad differentiates both branches of a where and multiplies the unselected one by zero, so 0 * NaN = NaN. Diagnosed in autolens_workspace_developer#104; it retroactively explains why #101's trajectories died with a non-finite objective rather than a NaN one, and why optax.apply_if_finite latched at the cliff.

These assertions exist so that limitation cannot silently regress into a false promise — and so the disproof of the obvious-looking fix is recorded where the next person will find it.

Scripts Changed

Script	Change
scripts/jax_assertions/fitness_nan_gradient_contract.py	new — six assertions pinning the guard's real contract.

Asserts:

• assert_guard_repairs_the_value — the guard delivers what it advertises; no NaN value reaches the search.
• assert_guard_does_not_repair_the_gradient — a guarded point still yields a NaN gradient. Deliberately pins broken-looking behaviour: if it ever fails, the guard's reach changed and fitness.py's docstring must be revisited.
• assert_gradient_is_fine_when_the_masked_derivative_is_finite — the trap is narrower than "never protects gradients": log(x<0) is NaN in value but its derivative 1/x is finite, so 0 * -1 = 0 and nothing leaks.
• assert_output_side_double_where_does_not_fix_it — the disproof, so nobody re-prescribes it.
• assert_input_side_safe_x_does_fix_it — what a real fix looks like, and why it cannot live in Fitness (it needs to know which input is invalid).
• assert_valid_points_are_unaffected — none of the above perturbs a well-behaved model.

Placement

JAX-only, so it belongs here rather than test_autofit/ — the library's unit tests are numpy-only, which is exactly why jax_assertions/ exists. Mirrors the sibling fitness_dispatch.py (same af.ex.Gaussian + Fitness construction, module-level assert_* functions, __main__ runner).

Testing

python scripts/jax_assertions/fitness_nan_gradient_contract.py → all 6 assertions pass. CPU, seconds, no A100.

Gate

• Tests — n/a (workspace_test has no unit-test dir); the script is the test and passes.
• Smoke — n/a: additive jax_assertions/ script, no existing script touched.
• Heart — RED, acknowledged contemporaneously by the human on 5 pre-existing unrelated reasons; RED verbatim: PyAutoLens: 1 uncommitted source change(s).