[Feature Request] Extend pseudo-point support beyond Cholesky variational distributions

[DiogoRibeiro7]

🚀 Feature Request

Suggested labels: enhancement, variational, refactor

Motivation

Is your feature request related to a problem? Please describe.

Pseudo-point support in variational strategies is currently restricted to CholeskyVariationalDistribution.

Both gpytorch/variational/variational_strategy.py and gpytorch/variational/unwhitened_variational_strategy.py raise:

NotImplementedError(
    "Only CholeskyVariationalDistribution has pseudo-point support currently, ..."
)

That leaves other variational distribution implementations without access to the same pseudo-point / amortized exact GP path.

Pitch

Describe the solution you'd like

Extend pseudo-point support so it is not hard-coded to CholeskyVariationalDistribution.

A natural direction would be to move the required mean/covariance access behind a method on _VariationalDistribution, as already suggested by the inline TODO in VariationalStrategy.pseudo_points.

This would make the pseudo-point computation less distribution-specific and make it easier to support additional variational distributions consistently.

Describe alternatives you've considered

The current alternative is to keep pseudo-point support limited to Cholesky-based variational distributions, but that preserves an avoidable internal limitation and pushes users into a specific distribution choice.

Are you willing to open a pull request? (We LOVE contributions!!!)

Potentially, yes.

Additional context

Relevant code comments already point in this direction:

• gpytorch/variational/variational_strategy.py
• gpytorch/variational/unwhitened_variational_strategy.py

In particular, VariationalStrategy.pseudo_points contains:

• TODO: have var_mean, var_cov come from a method of _variational_distribution