Hi team. I'm looking to implement the GP clustering method described in Hierarchical Bayesian modelling of gene expression time series across irregularly sampled replicates and clusters using GPJax.

Looking at GPy's implementation used in the paper as well as @thomaspinder's implementation in in GPFlow, it seems that the notion of a hierarchical kernel fits naturally into the CombinationKernel base class. However in GPJax this would require passing additional arguments to the pairwise combination_fn indicating which replicate the data point belongs to. For instance, see equation (5) in Fast Nonparametric Clustering of Structured Time-Series

The piecewise function above could then be written using lax.cond (or lax.switch for deeper hierarchies).

Do you think this is a over-engineered solution to remain consistent with GPJax's vmap-over-pairwise-kernel-functions approach instead of defining the covariance matrix using something like

def Ktilde(t, td, n, nd):
    K = kf.cross_covariance(kf_params, t, td)
    mask = n == nd.T
    K += mask * kg.cross_covariance(kg_params, t, td)
    return K

It would be nice to integrate smoothly with GPJax to experiment with sparse methods etc.