Mean function and kernel with common/shared hyperparameters

[theo-brown]

I would like to construct a GP that has a hyperparameter that appears in both the mean function and the kernel, such as

$$ \begin{align} \mu(x) &= ax + b. \\ k(x, x') &= \exp\left(-a^2 (x - x')^2\right). \end{align} $$

I can implement this easy enough in vanilla JAX NumPy, as follows.

import jax.numpy as jnp
from jax import vmap 
import jax

def mean_function(x, a, b):
    return a * x + b

def kernel(x1, x2, a):
    return jnp.exp(-a**2 * jnp.sum((x1 - x2) ** 2, axis=-1))

def gram(kernel, x1, x2):
    return vmap(lambda x1: vmap(lambda x2: kernel(x1, x2))(x2))(x1)

def generate_samples(key, a, b, n_samples=5):
    x = jnp.linspace(0, 10, 100).reshape(-1, 1)
    K = gram(lambda x1, x2: kernel(x1, x2, a), x, x)
    mu = mean_function(x, a, b)

    L = jnp.linalg.cholesky(K + 1e-6 * jnp.eye(K.shape[0]))
    z = jax.random.normal(key, (n_samples, K.shape[0]))
    samples = jnp.dot(z, L.T) + mu.reshape(-1) 

    return x, samples

This produces something like the following plot.

However, I couldn't work out how I'd do the same thing within GPJax, as the gpx.gps.Prior takes in a mean_function and kernel that have already been constructed (and hence with an $a$ parameter each).

How can I make the $a$ parameter common between the mean and kernel function in GPJax? In particular, I'd like to learn $a$ using the MLL.

Thanks all for your amazing work on this library. Looking forward to hopefully contributing to it in future!

[theo-brown]

My first thought is to subclass gpx.gps.Prior.

I'm not sure whether there's a neater way of doing it, but here goes:

class MyPrior(gpx.gps.Prior):
    def __init__(
        self,
        a: float,
        b: float,
        jitter: float = 1e-6,
    ):
        self.a = parameters.Real(a)
        self.b = parameters.Real(b)
        self.jitter = jitter

    @property
    def kernel(self):       
        return MyKernel(self.b.value)
    
    @property
    def mean_function(self):
        return MyMean(self.a.value, self.b.value)

where MyKernel and MyMean are GPJax-ified versions of the functions above.

[thomaspinder]

Hey! Cool problem!

Have you tried directly passing the parameter yourself? My first thought is that it would work.

import gpjax as gpx
from gpjax.parameters import PositiveReal

shared_param = PositiveReal(1.)
kernel = MyKernel(shared_param)
meanf = MyMean(shared_param)
prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
...

Perhaps try that, and feel free to reply here if it does not work - I can try to help further.

[theo-brown]

I had tried this but it was erroring and I've now found the bit of code I was missing.

The custom kernel in the docs has the following:

def __init__(
        self,
        tau: float = 5.0,
        ...
    ):
        ....
        self.tau = PositiveReal(jnp.array(tau), tag="polar")

Previously I was missing the jnp.array(tau), and was trying to do something like self.tau = PositiveReal(tau), which throws an error if tau is already a gpx.Parameter.

[theo-brown]

However, I think this will still register the parameters as separate nnx.Variables / leaves of the tree. I'm not entirely sure of this, but my hunch is that this means that when choosing the parameters with MLL they will not be constrained to have the same value.

import gpjax as gpx 
import jax.numpy as jnp
from jax import config 
import flax.nnx as nnx 
import jax 
config.update("jax_enable_x64", True)

class MyKernel(gpx.kernels.AbstractKernel):
    def __init__(self, a: float | gpx.parameters.PositiveReal, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.a = gpx.parameters.PositiveReal(jnp.array(a))

    def __call__(self, x1: jax.Array, x2: jax.Array) -> jax.Array:
        return jnp.exp(-self.a**2 * jnp.sum((x1 - x2) ** 2, axis=-1))
    
class MyMeanFunction(gpx.mean_functions.AbstractMeanFunction):
    def __init__(self, a: float | gpx.parameters.PositiveReal, b: float | gpx.parameters.PositiveReal, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.a = gpx.parameters.PositiveReal(jnp.array(a))
        self.b = gpx.parameters.PositiveReal(jnp.array(b))
    
    def __call__(self, x: jax.Array) -> jax.Array:
        return self.a * x + self.b
        
a = gpx.parameters.PositiveReal(1.)
kernel = MyKernel(a=a)
meanf = MyMeanFunction(a=a, b=0.5)
prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
graphdef, state = nnx.split(prior)

Not entirely sure I'm parsing this correctly, but it looks like there are more than 2 leaves.
graphdef:

GraphDef(
  nodedef=NodeDef(
    type=Prior,
    index=0,
    attributes=('jitter', 'kernel', 'mean_function'),
    subgraphs={
      'kernel': NodeDef(
        type=MyKernel,
        index=1,
        attributes=('a', 'active_dims', 'compute_engine', 'n_dims'),
        subgraphs={
          'n_dims': NodeDef(
            type=PytreeType,
            index=-1,
            attributes=(),
            subgraphs={},
            static_fields={},
            leaves={},
            metadata=PyTreeDef(None)
          )
        },
        static_fields={
          'active_dims': slice(None, None, None),
          'compute_engine': <gpjax.kernels.computations.dense.DenseKernelComputation object at 0x7f3879e427a0>
        },
        leaves={
          'a': 2
        },
        metadata=<class '__main__.MyKernel'>
      ),
      'mean_function': NodeDef(
        type=MyMeanFunction,
        index=3,
        attributes=('a', 'b'),
        subgraphs={},
        static_fields={},
        leaves={
          'a': 4,
          'b': 5
        },
        metadata=<class '__main__.MyMeanFunction'>
      )
    },
    static_fields={
      'jitter': 1e-06
    },
    leaves={},
    metadata=<class 'gpjax.gps.Prior'>
  ),
  index_mapping=None
)

state:

State({
  'kernel': {
    'a': VariableState(
      type=PositiveReal,
      value=Array(1., dtype=float64),
      _tag='positive'
    )
  },
  'mean_function': {
    'a': VariableState(
      type=PositiveReal,
      value=Array(1., dtype=float64),
      _tag='positive'
    ),
    'b': VariableState(
      type=PositiveReal,
      value=Array(0.5, dtype=float64, weak_type=True),
      _tag='positive'
    )
  }
})

[theo-brown]

Looking at the source code, rather than the docs, the implemented kernels have an additional check in in their constructor:

        if isinstance(lengthscale, nnx.Variable):
            self.lengthscale = lengthscale
        else:
            self.lengthscale = PositiveReal(lengthscale)

which seems to get round this. If I implement this logic:

class MyKernel(gpx.kernels.AbstractKernel):
   def __init__(self, a: float | nnx.Variable, *args, **kwargs):
       super().__init__(*args, **kwargs)
       
       if isinstance(a, nnx.Variable):
           self.a = a 
       else:
           self.a = gpx.parameters.PositiveReal(jnp.array(a))

   def __call__(self, x1: jax.Array, x2: jax.Array) -> jax.Array:
       return jnp.exp(-self.a**2 * jnp.sum((x1 - x2) ** 2, axis=-1))
   
class MyMeanFunction(gpx.mean_functions.AbstractMeanFunction):
   def __init__(self, a: float | nnx.Variable, b: float | gpx.parameters.PositiveReal, *args, **kwargs):
       super().__init__(*args, **kwargs)
       
       if isinstance(a, nnx.Variable):
           self.a = a
       else:
           self.a = gpx.parameters.PositiveReal(jnp.array(a))
       if isinstance(b, nnx.Variable):
           self.b = b
       else:
           self.b = gpx.parameters.PositiveReal(jnp.array(b))
   
   def __call__(self, x: jax.Array) -> jax.Array:
       return self.a * x + self.b

then prior has graphdef

GraphDef(
 nodedef=NodeDef(
   type=Prior,
   index=0,
   attributes=('jitter', 'kernel', 'mean_function'),
   subgraphs={
     'kernel': NodeDef(
       type=MyKernel,
       index=1,
       attributes=('a', 'active_dims', 'compute_engine', 'n_dims'),
       subgraphs={
         'n_dims': NodeDef(
           type=PytreeType,
           index=-1,
           attributes=(),
           subgraphs={},
           static_fields={},
           leaves={},
           metadata=PyTreeDef(None)
         )
       },
       static_fields={
         'active_dims': slice(None, None, None),
         'compute_engine': <gpjax.kernels.computations.dense.DenseKernelComputation object at 0x7f3879e427a0>
       },
       leaves={
         'a': 2
       },
       metadata=<class '__main__.MyKernel'>
     ),
     'mean_function': NodeDef(
       type=MyMeanFunction,
       index=3,
       attributes=('a', 'b'),
       subgraphs={},
       static_fields={},
       leaves={
         'a': 2,
         'b': 4
       },
       metadata=<class '__main__.MyMeanFunction'>
     )
   },
   static_fields={
     'jitter': 1e-06
   },
   leaves={},
   metadata=<class 'gpjax.gps.Prior'>
 ),
 index_mapping=None
)

and state

State({
  'kernel': {
    'a': VariableState(
      type=PositiveReal,
      value=Array(1., dtype=float64, weak_type=True),
      _tag='positive'
    )
  },
  'mean_function': {
    'b': VariableState(
      type=PositiveReal,
      value=Array(0.5, dtype=float64, weak_type=True),
      _tag='positive'
    )
  }
})

which is the desired effect.

[thomaspinder]

Super! I'm glad this solved your issue. It's great that you're pushing the boundaries of GPJax :D

A question for you - what content could we have included in our docs to make this a simpler issue for you to solve?