`GaussianDistribution.kl_divergence` raises for `BlockDiag` and `Kronecker` covariances

[thomaspinder]

Problem

GaussianDistribution.kl_divergence fails outright for the two custom operators in gpjax/linalg/custom_operators.py:

from gpjax.linalg import BlockDiag, Kronecker
A = lx.MatrixLinearOperator(jnp.array([[2., 0.3], [0.3, 1.5]]))
B = lx.MatrixLinearOperator(jnp.array([[3., 0.1], [0.1, 2.0]]))
q = GaussianDistribution(loc=jnp.zeros(4), scale=BlockDiag([A, B]))
q.kl_divergence(q)

NotImplementedError: `lineax.has_unit_diagonal` has not been implemented for
  <class 'gpjax.linalg.custom_operators.BlockDiag'>

Both BlockDiag and Kronecker are affected, and both fail identically.

Root cause

_kl_divergence computes the Mahalanobis term with lx.linear_solve(sqrt_p, diff, solver=lx.Triangular()), and lx.Triangular() queries lx.has_unit_diagonal on the operator. custom_operators.py registers seven lineax predicates for both operators but not that one:

predicate	BlockDiag	Kronecker
is_symmetric	registered	registered
is_diagonal	registered	registered
is_tridiagonal	registered	registered
is_lower_triangular	registered	registered
is_upper_triangular	registered	registered
is_positive_semidefinite	registered	registered
is_negative_semidefinite	registered	registered
has_unit_diagonal	missing	missing

So this is a one-line-per-operator omission rather than a design problem. Everything else in the KL path already handles these operators: cholesky_factor has BlockDiag/Kronecker implementations, as does logdet.

Recommendation

Register has_unit_diagonal for both, returning False (neither operator guarantees a unit diagonal in general; BlockDiag could recurse over blocks if a tighter answer is wanted). Add a KL round-trip test covering both operators so the gap can't reopen.

Scope note

Worth auditing whether any other lineax predicate is queried elsewhere in the codebase for these operators — the audit above only covers the eight predicates I checked, and the same omission class could bite in another solver path.

_{Found while verifying #664 (see #707). Pre-existing on main; not introduced by that work.}