Develop multi-output GP model with linear mean for each task and with learnable hyperparameters

Questions v5
6

Thank you @bwengals for your reply. I defined the following mean function

class Linear(pm.gp.mean.Mean):
def init(self, coeffs, intercept):
super().init()
self.b = intercept
self.A = coeffs

def __call__(self, X):
    return pt.squeeze(pt.dot(X, self.A) + self.b)

and then tried to used the function to define the model as the following code

with pm.Model() as model:

# Priors
ell = pm.Gamma("ell", alpha=3, beta=0.2, shape=2)
eta = pm.Gamma("eta", alpha=3, beta=0.2, shape=2)
kernels = [pm.gp.cov.ExpQuad, pm.gp.cov.Matern52]
sigma = pm.HalfNormal("sigma", sigma=3)
beta_m = pm.Gamma("beta_m", alpha=3, beta=0.2, shape=2)
b_m = pm.Gamma("b_m", alpha=3, beta=0.2, shape=2)

# Define a list of covariance functions
cov_list = [
    eta[idx] ** 2 * kernel(input_dim=2, ls=ell[idx], active_dims=[0])
    for idx, kernel in enumerate(kernels)
]

# Get the LCM kernel
cov_lcm = get_lcm(input_dim=2, active_dims=[1], num_outputs=n_outputs, kernels=cov_list)

# define linear mean
mean = Linear(coeffs=beta_m, intercept=b_m)

# Define a Multi-output GP
mogp = pm.gp.Marginal(mean_func=mean, cov_func=cov_lcm)
y_ = mogp.marginal_likelihood("f", x, y, sigma=sigma)

But I got the following error:

ValueError Traceback (most recent call last)
File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\tensor\elemwise.py:441, in Elemwise.get_output_info(self, dim_shuffle, inputs)
439 try:
440 out_shapes = [
→ 441 [
442 get_most_specialized_shape(shape)
443 for shape in zip([inp.type.shape for inp in inputs])
444 ]
445 ] * shadow.nout
446 except ValueError:

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\tensor\elemwise.py:442, in (.0)
439 try:
440 out_shapes = [
441 [
→ 442 get_most_specialized_shape(shape)
443 for shape in zip(*[inp.type.shape for inp in inputs])
444 ]
445 ] * shadow.nout
446 except ValueError:

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\tensor\elemwise.py:434, in Elemwise.get_output_info..get_most_specialized_shape(shapes)
433 if len(shapes) > 1:
→ 434 raise ValueError
435 return tuple(shapes)[0]

ValueError:

During handling of the above exception, another exception occurred:

ValueError Traceback (most recent call last)
Cell In[64], line 25
23 # Define a Multi-output GP
24 mogp = pm.gp.Marginal(mean_func=mean, cov_func=cov_lcm)
—> 25 y_ = mogp.marginal_likelihood(“f”, x, y, sigma=sigma)

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pymc\gp\gp.py:473, in Marginal.marginal_likelihood(self, name, X, y, sigma, noise, jitter, is_observed, **kwargs)
470 sigma = _handle_sigma_noise_parameters(sigma=sigma, noise=noise)
472 noise_func = sigma if isinstance(sigma, Covariance) else pm.gp.cov.WhiteNoise(sigma)
→ 473 mu, cov = self._build_marginal_likelihood(X=X, noise_func=noise_func, jitter=jitter)
474 self.X = X
475 self.y = y

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pymc\gp\gp.py:429, in Marginal._build_marginal_likelihood(self, X, noise_func, jitter)
428 def _build_marginal_likelihood(self, X, noise_func, jitter):
→ 429 mu = self.mean_func(X)
430 Kxx = self.cov_func(X)
431 Knx = noise_func(X)

Cell In[63], line 9, in Linear.call(self, X)
8 def call(self, X):
----> 9 return pt.squeeze(pt.dot(X, self.A) + self.b)

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\tensor\var.py:106, in _tensor_py_operators.add(self, other)
104 def add(self, other):
105 try:
→ 106 return at.math.add(self, other)
107 # We should catch the minimum number of exception here.
108 # Otherwise this will convert error when PyTensor flags
109 # compute_test_value is used
110 # Evidently, we need to catch NotImplementedError
111 # TypeError from as_tensor_variable are caught in Elemwise.make_node
112 # Otherwise TensorVariable * SparseVariable won’t work!
113 except (NotImplementedError, TypeError):
114 # We must return NotImplemented and not an
115 # NotImplementedError or raise an NotImplementedError.
116 # That way python will give a good error message like this
117 # TypeError: unsupported operand type(s) for +: 118 # 'TensorVariable' and 'TensorVariable'

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\graph\op.py:295, in Op.call(self, *inputs, **kwargs)
253 r""“Construct an Apply node using :meth:Op.make_node and return its outputs.
254
255 This method is just a wrapper around :meth:Op.make_node.
(…)
292
293 “””
294 return_list = kwargs.pop(“return_list”, False)
→ 295 node = self.make_node(*inputs, **kwargs)
297 if config.compute_test_value != “off”:
298 compute_test_value(node)

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\tensor\elemwise.py:485, in Elemwise.make_node(self, *inputs)
479 “”"
480 If the inputs have different number of dimensions, their shape
481 is left-completed to the greatest number of dimensions with 1s
482 using DimShuffle.
483 “”"
484 inputs = [as_tensor_variable(i) for i in inputs]
→ 485 out_dtypes, out_shapes, inputs = self.get_output_info(DimShuffle, *inputs)
486 outputs = [
487 TensorType(dtype=dtype, shape=shape)()
488 for dtype, shape in zip(out_dtypes, out_shapes)
489 ]
490 return Apply(self, inputs, outputs)

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\tensor\elemwise.py:447, in Elemwise.get_output_info(self, dim_shuffle, inputs)
440 out_shapes = [
441 [
442 get_most_specialized_shape(shape)
443 for shape in zip([inp.type.shape for inp in inputs])
444 ]
445 ] * shadow.nout
446 except ValueError:
→ 447 raise ValueError(
448 f"Incompatible Elemwise input shapes {[inp.type.shape for inp in inputs]}"
449 )
451 # inplace_pattern maps output idx → input idx
452 inplace_pattern = self.inplace_pattern

ValueError: Incompatible Elemwise input shapes [(22,), (2,)]

I appreciate any advice to get rid of this error.

Thank you