The Pyro module pyro.contrib.gp is pretty stable and there are already a couple of examples (deep kernel learning, gplvm) of doing SVI inference for sparse GP models in Pyro. I would like to extend the feature request for having an example of deep sparse GP in NumPyro, which isn't duplicated with the current examples on Pyro and illustrates how to use flax_module primitive. Please see the following references as a guideline for an implementation:

• GP conditional: This utility implements the core math of GP and is used across all GP models.
• RBF kernel: a popular GP kernel
• Variational Sparse Gaussian Process model: the most Pyro-friendly implementation of sparse GP.
• See flax_module and random_flax_module for how to use the neural network library Flax in numpyro.
• Some Pyro examples: deep kernel learning, gplvm.
• The deep GP tutorial that I wrote 2 years ago: when you have sparse GP implemented, extending it to deep gp or deep kernel gp is pretty straightforward.