`DualVariationalGaussian` can leave the PD cone / NaN under plain `fit()` despite `dual_elbo` docstring calling it valid

[thomaspinder]

Problem

DualVariationalGaussian.dual_matrix (gpjax/variational_families.py:572, stored unconstrained at :590) has no PD-cone protection on the plain fit() path, even though objectives.dual_elbo's own docstring (gpjax/objectives.py:381-386) explicitly states:

"Plain fit() on a DualVariationalGaussian with this objective remains valid — it is ordinary gradient descent in the dual coordinates."

The class's own docstring (variational_families.py:545-547) says PSD-ness "comes from the convex-combination structure of the natural-gradient update" — but that guarantee only holds under fit_natgrads, not under plain fit(). _working_matrices() (:656-659) calls raw jnp.linalg.cholesky on dual_matrix-derived quantities with no stabilisation.

This is structurally the same defect #666 flagged for the now-deleted NaturalVariationalGaussian/ExpectationVariationalGaussian (unconstrained PSD-required matrix, no protection against leaving the PD cone under a documented-valid optimiser path) — it resurfaced in the new DualVariationalGaussian family added by the same refactor that deleted those two classes.

Notably, fit_natgrads already has a bespoke natgrad_lr <= 1 guard specifically for DualVariationalGaussian (fit.py:777), showing the maintainers are aware step-size bounds matter for this family's PD invariant — but no equivalent guard exists on the plain-fit() path that dual_elbo's own docstring advertises as valid.

Repro

import optax as ox
import gpjax as gpx
from gpjax.objectives import dual_elbo

# ... construct `model` as a DualVariationalGaussian over some JointModel/inducing set ...

opt_model, history = gpx.fit(
    model=model,
    objective=lambda m, d: -dual_elbo(m, d),
    train_data=D,
    optim=ox.sgd(1.0),
    num_iters=50,
)

Empirically reproduced against the current v1.0 branch: ox.sgd(1.0) for 50 iterations produces loss = nan and min_eig(dual_matrix) = nan. ox.adam(5.0) for 100 iterations produces min_eig(dual_matrix) = -3.07e-4 (already PSD-violating, pre-NaN).

tests/test_fit.py:1182 (test_fit_on_dual_family_still_works) exercises plain fit() on DualVariationalGaussian but only at lr=1e-2/20 iters — too mild to expose the break.

Recommendation

Either:

1. Add a PD-cone guard/step-size bound to the plain fit() path when training a DualVariationalGaussian (mirroring the natgrad_lr <= 1 guard fit_natgrads already has), or
2. Correct dual_elbo's docstring to state that plain fit() is only safe below some step-size bound (and ideally enforce it), rather than calling it unconditionally "valid".

Context

Found via adversarial re-verification of issue #666 during a repo-wide audit of open issues against the v1.0 conditioning refactor (#666 itself is being closed/labeled v1.0-resolved: its two named classes were deleted in ada0f1d0 and their design successor, DualVariationalGaussian.prior_kl, already implements the efficient factorisation #666's sibling issue #710 asked for — but this new PD-cone gap in the same class was found in the process and doesn't map to any existing issue).

Evidence

• gpjax/variational_families.py:519 (class DualVariationalGaussian), :545-547 (docstring), :572/:590 (dual_matrix stored as unconstrained Real), :656-659 (_working_matrices, raw jnp.linalg.cholesky, no stabilisation)
• gpjax/objectives.py:381-386 (dual_elbo docstring sanctioning plain fit())
• gpjax/fit.py fit() body (no NaN/PD guard in the optax scan loop); contrast fit.py:777-779 (natgrad_lr > 1.0 rejection specific to DualVariationalGaussian in fit_natgrads)
• tests/test_fit.py:1182 (test_fit_on_dual_family_still_works, too mild a step size/iteration count to expose the break)