Repository: GPflow/GPflow

Issue:
Set a prior over a linear mean function
<!-- Lines like this are comments and will be invisible -->

Bug: I set a Gaussian prior over a linear mean function and when a I see the monitor, gpflow only assigns a prior over the b parameter, and no for the A parameter.

## To reproduce

**Minimal, reproducible example**
I just put the code for that part of the model, for simplicity:
```python
import gpflow
import numpy as np
import tensorflow_probability as tfp
import tensorflow as tf

#trivial example, from [here](https://gpflow.github.io/GPflow/2.5.2/notebooks/advanced/mcmc.html) 
rng = np.random.RandomState(14)
X = np.linspace(-3, 3, 20).reshape(1,20)
y = rng.exponential(np.sin(X) ** 2)

A = gpflow.Parameter(np.zeros(X.shape[1]),dtype=np.float64)
A.prior=tfp.distributions.Normal((np.zeros(X.shape[1])),np.ones(X.shape[1]))
b = gpflow.Parameter(np.zeros((1,)),dtype=np.float64)
b.prior=tfp.distributions.Normal(np.zeros((1,)),np.ones((1,)))
mf = gpflow.functions.Linear(A,b)

k = gpflow.kernels.SquaredExponential()
k.lengthscales.prior = tfp.distributions.Normal(np.zeros(X.shape[1]),np.ones(X.shape[1]))
k.variance.prior = tfp.distributions.Normal(np.zeros(1),np.ones(1))

lik = gpflow.likelihoods.Gaussian()
lik.variance.prior = tfp.distributions.Normal(np.zeros(1),np.ones(1))

model = gpflow.models.GPMC((X, y), kernel=k, mean_function=mf, likelihood=lik)
model
```
So, there is no warning, but, as can be seen from the monitor, gpflow doesn't put the prior over A:

![imagen](https://github.com/GPflow/GPflow/assets/112404457/1fb5d981-0e00-4e07-8a5a-e68c8735fec1)


System information

* GPflow version: 2.9.0
* GPflow installed from: "pip install gpflow"
* TensorFlow version: 2.10.1
* Python version: 3.10.9 | packaged by conda-forge | (main, Jan 11 2023, 15:15:40) [MSC v.1916 64 bit (AMD64)]
* Operating system: win11/Ubuntu 18.046 LTS (I probed it in both os)

I guess the setting is correct, but I don't understand why gpflow set the prior over b and not over A.





Discussion:
I fixed the bug creating a new class for the mean function like this:
```python
##can aggregate this 'from' in the first part
from gpflow.mean_functions import MeanFunction
class Linear(MeanFunction):
    def __init__(self, A=None, b=None):
        MeanFunction.__init__(self)
        self.A = gpflow.Parameter(A, dtype=gpflow.default_float()) if A is not None else None
        self.b = gpflow.Parameter(b, dtype=gpflow.default_float()) if b is not None else None

    def __call__(self, X):
        if self.A is None:
            return self.b * tf.ones((X.shape[0], 1), dtype=gpflow.default_float())
        else:
            return tf.matmul(X, self.A) + self.b
``` 

So, to set the priors over A and b, you can do as shown in the first piece of code, then design the model, e.g., as:

```python
gpflow.models.GPMC((X, y), kernel=k, mean_function=Linear(A,b), likelihood=lik)
``` 
Hi @santiromoli,
I stumbled over the same problem. Short question, did you try this fix also with a `gpflow.models.GPR()` model? Because when I do this I get an error.
Hi @JSchmiegel! mmm no, I just re-defined the linear mean func. I made this because I thought the error was in `gpflow.functions.Linear`, but maybe it is in the GPMC module. 
Or, do you mean you get an error while running my script? Maybe I got a mistake, because I changed the real code in order to avoid uploading X,y. Maybe it's something related to GPMC, which is a particular module of gpflow to make Hamiltonian Monte Carlo.


Hi @santiromoli,
first of all, thank you for your quick reply. I also believe that it is rather a problem with the mean function, because I see exactly the error pattern described in this issue by you also with `gpflow.model.GPR`. Your code also seems to be correct because when I use it to create a corresponding model, I get the desired output. 
![image](https://github.com/GPflow/GPflow/assets/51242445/9ffd8375-c2e0-46ed-a17d-822f4a9bc9fb)


The only problem is when I try to optimize this model. Then I get the following error. Don't you have this with GPMC?
Code:
```python
opt = gpflow.optimizers.Scipy()
print(opt.minimize(model.training_loss, model.trainable_variables, method='SLSQP'))
```
Error:
```text
{
	"name": "ValueError",
	"message": "in user code:

    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\optimizers\\scipy.py\", line 185, in _tf_eval  *
        loss, grads = _compute_loss_and_gradients(
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\optimizers\\scipy.py\", line 322, in _compute_loss_and_gradients  *
        loss = loss_closure()
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\integration\\tf.py\", line 57, in wrapped_method  *
        return wrapped_function(self, *args, **kwargs)
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\decorator.py\", line 121, in wrapped_function  *
        if not get_enable_check_shapes():
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\models\\training_mixins.py\", line 64, in training_loss  *
        return self._training_loss()  # type: ignore[attr-defined]
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\integration\\tf.py\", line 57, in wrapped_method  *
        return wrapped_function(self, *args, **kwargs)
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\decorator.py\", line 121, in wrapped_function  *
        if not get_enable_check_shapes():
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\models\\gpmc.py\", line 81, in _training_loss  *
        return -self.log_posterior_density()
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\integration\\tf.py\", line 57, in wrapped_method  *
        return wrapped_function(self, *args, **kwargs)
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\decorator.py\", line 121, in wrapped_function  *
        if not get_enable_check_shapes():
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\models\\gpmc.py\", line 76, in log_posterior_density  *
        return self.log_likelihood() + self.log_prior_density()
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\integration\\tf.py\", line 57, in wrapped_method  *
        return wrapped_function(self, *args, **kwargs)
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\decorator.py\", line 121, in wrapped_function  *
        if not get_enable_check_shapes():
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\models\\gpmc.py\", line 104, in log_likelihood  *
        F = tf.linalg.matmul(L, self.V) + self.mean_function(X_data)
    File \"C:\\Users\\Jonas\\AppData\\Local\\Temp\\ipykernel_14220\\4104737182.py\", line 12, in __call__  *
        return tf.matmul(X, self.A) + self.b

    ValueError: Shape must be rank 2 but is rank 1 for '{{node MatMul_2}} = MatMul[T=DT_DOUBLE, transpose_a=false, transpose_b=false](781, MatMul_2/identity_CONSTRUCTED_AT_top_level/forward/ReadVariableOp)' with input shapes: [1,20], [20].
",
	"stack": "---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
j:\\100Anwendung\\dataanalysis\\900ArtificialAnalysis\\PriorError.ipynb Cell 5 line 2
      <a href='vscode-notebook-cell:/j%3A/100Anwendung/dataanalysis/900ArtificialAnalysis/PriorError.ipynb#W6sZmlsZQ%3D%3D?line=0'>1</a> opt = gpflow.optimizers.Scipy()
----> <a href='vscode-notebook-cell:/j%3A/100Anwendung/dataanalysis/900ArtificialAnalysis/PriorError.ipynb#W6sZmlsZQ%3D%3D?line=1'>2</a> print(opt.minimize(model.training_loss, model.trainable_variables, method='SLSQP'))

File j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\optimizers\\scipy.py:152, in Scipy.minimize(self, closure, variables, method, step_callback, compile, allow_unused_variables, tf_fun_args, track_loss_history, **scipy_kwargs)
    149     callback = self.loss_history_callback_func(func, history, scipy_kwargs.get(\"callback\"))
    150     scipy_kwargs[\"callback\"] = callback
--> 152 opt_result = scipy.optimize.minimize(
    153     func, initial_params, jac=True, method=method, **scipy_kwargs
    154 )
    156 if track_loss_history:
    157     opt_result[\"loss_history\"] = history

File j:\\100Anwendung\\_venv_\\lib\\site-packages\\scipy\\optimize\\_minimize.py:705, in minimize(fun, x0, args, method, jac, hess, hessp, bounds, constraints, tol, callback, options)
    702     res = _minimize_cobyla(fun, x0, args, constraints, callback=callback,
    703                             **options)
    704 elif meth == 'slsqp':
--> 705     res = _minimize_slsqp(fun, x0, args, jac, bounds,
    706                           constraints, callback=callback, **options)
    707 elif meth == 'trust-constr':
    708     res = _minimize_trustregion_constr(fun, x0, args, jac, hess, hessp,
    709                                        bounds, constraints,
    710                                        callback=callback, **options)

File j:\\100Anwendung\\_venv_\\lib\\site-packages\\scipy\\optimize\\_slsqp_py.py:374, in _minimize_slsqp(func, x0, args, jac, bounds, constraints, maxiter, ftol, iprint, disp, eps, callback, finite_diff_rel_step, **unknown_options)
    371     xu[infbnd[:, 1]] = np.nan
    373 # ScalarFunction provides function and gradient evaluation
--> 374 sf = _prepare_scalar_function(func, x, jac=jac, args=args, epsilon=eps,
    375                               finite_diff_rel_step=finite_diff_rel_step,
    376                               bounds=new_bounds)
    377 # gh11403 SLSQP sometimes exceeds bounds by 1 or 2 ULP, make sure this
    378 # doesn't get sent to the func/grad evaluator.
    379 wrapped_fun = _clip_x_for_func(sf.fun, new_bounds)

File j:\\100Anwendung\\_venv_\\lib\\site-packages\\scipy\\optimize\\_optimize.py:332, in _prepare_scalar_function(fun, x0, jac, args, bounds, epsilon, finite_diff_rel_step, hess)
    328     bounds = (-np.inf, np.inf)
    330 # ScalarFunction caches. Reuse of fun(x) during grad
    331 # calculation reduces overall function evaluations.
--> 332 sf = ScalarFunction(fun, x0, args, grad, hess,
    333                     finite_diff_rel_step, bounds, epsilon=epsilon)
    335 return sf

File j:\\100Anwendung\\_venv_\\lib\\site-packages\\scipy\\optimize\\_differentiable_functions.py:158, in ScalarFunction.__init__(self, fun, x0, args, grad, hess, finite_diff_rel_step, finite_diff_bounds, epsilon)
    155     self.f = fun_wrapped(self.x)
    157 self._update_fun_impl = update_fun
--> 158 self._update_fun()
    160 # Gradient evaluation
    161 if callable(grad):

File j:\\100Anwendung\\_venv_\\lib\\site-packages\\scipy\\optimize\\_differentiable_functions.py:251, in ScalarFunction._update_fun(self)
    249 def _update_fun(self):
    250     if not self.f_updated:
--> 251         self._update_fun_impl()
    252         self.f_updated = True

File j:\\100Anwendung\\_venv_\\lib\\site-packages\\scipy\\optimize\\_differentiable_functions.py:155, in ScalarFunction.__init__.<locals>.update_fun()
    154 def update_fun():
--> 155     self.f = fun_wrapped(self.x)

File j:\\100Anwendung\\_venv_\\lib\\site-packages\\scipy\\optimize\\_differentiable_functions.py:137, in ScalarFunction.__init__.<locals>.fun_wrapped(x)
    133 self.nfev += 1
    134 # Send a copy because the user may overwrite it.
    135 # Overwriting results in undefined behaviour because
    136 # fun(self.x) will change self.x, with the two no longer linked.
--> 137 fx = fun(np.copy(x), *args)
    138 # Make sure the function returns a true scalar
    139 if not np.isscalar(fx):

File j:\\100Anwendung\\_venv_\\lib\\site-packages\\scipy\\optimize\\_optimize.py:76, in MemoizeJac.__call__(self, x, *args)
     74 def __call__(self, x, *args):
     75     \"\"\" returns the function value \"\"\"
---> 76     self._compute_if_needed(x, *args)
     77     return self._value

File j:\\100Anwendung\\_venv_\\lib\\site-packages\\scipy\\optimize\\_optimize.py:70, in MemoizeJac._compute_if_needed(self, x, *args)
     68 if not np.all(x == self.x) or self._value is None or self.jac is None:
     69     self.x = np.asarray(x).copy()
---> 70     fg = self.fun(x, *args)
     71     self.jac = fg[1]
     72     self._value = fg[0]

File j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\optimizers\\scipy.py:217, in Scipy.eval_func.<locals>._eval(x)
    216 def _eval(x: AnyNDArray) -> Tuple[AnyNDArray, AnyNDArray]:
--> 217     loss, grad = _tf_eval(tf.convert_to_tensor(x))
    218     return loss.numpy().astype(np.float64), grad.numpy().astype(np.float64)

File j:\\100Anwendung\\_venv_\\lib\\site-packages\\tensorflow\\python\\util\\traceback_utils.py:153, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    151 except Exception as e:
    152   filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153   raise e.with_traceback(filtered_tb) from None
    154 finally:
    155   del filtered_tb

File ~\\AppData\\Local\\Temp\\__autograph_generated_fileevsyg_0f.py:37, in outer_factory.<locals>.inner_factory.<locals>.tf___tf_eval(x)
     35 grads = ag__.Undefined('grads')
     36 loss = ag__.Undefined('loss')
---> 37 ag__.if_stmt(ag__.ld(first_call), if_body, else_body, get_state, set_state, ('grads', 'loss', 'first_call'), 2)
     38 try:
     39     do_return = True

File ~\\AppData\\Local\\Temp\\__autograph_generated_fileevsyg_0f.py:28, in outer_factory.<locals>.inner_factory.<locals>.tf___tf_eval.<locals>.if_body()
     26 def if_body():
     27     nonlocal grads, loss, first_call
---> 28     (loss, grads) = ag__.converted_call(ag__.ld(_compute_loss_and_gradients), (ag__.ld(closure), ag__.ld(variables), ag__.ld(tf).UnconnectedGradients.NONE), None, fscope)
     29     grads = ag__.converted_call(ag__.ld(self)._filter_unused_variables, (ag__.ld(variables), ag__.ld(grads), ag__.ld(allow_unused_variables)), None, fscope)
     30     first_call = False

File ~\\AppData\\Local\\Temp\\__autograph_generated_filei5upes9h.py:12, in outer_factory.<locals>.inner_factory.<locals>.tf___compute_loss_and_gradients(loss_closure, variables, unconnected_gradients)
     10 with ag__.ld(tf).GradientTape(watch_accessed_variables=False) as tape:
     11     ag__.converted_call(ag__.ld(tape).watch, (ag__.ld(variables),), None, fscope)
---> 12     loss = ag__.converted_call(ag__.ld(loss_closure), (), None, fscope)
     13 grads = ag__.converted_call(ag__.ld(tape).gradient, (ag__.ld(loss), ag__.ld(variables)), dict(unconnected_gradients=ag__.ld(unconnected_gradients)), fscope)
     14 try:

File ~\\AppData\\Local\\Temp\\__autograph_generated_file36u7ipsh.py:13, in outer_factory.<locals>.inner_factory.<locals>.tf__wrapped_method(self, *args, **kwargs)
     11 try:
     12     do_return = True
---> 13     retval_ = ag__.converted_call(ag__.ld(wrapped_function), (ag__.ld(self),) + tuple(ag__.ld(args)), dict(**ag__.ld(kwargs)), fscope)
     14 except:
     15     do_return = False

File ~\\AppData\\Local\\Temp\\__autograph_generated_file3aobj0zg.py:40, in outer_factory.<locals>.inner_factory.<locals>.tf__wrapped_function(*args, **kwargs)
     38     nonlocal retval_, do_return
     39     pass
---> 40 ag__.if_stmt(ag__.not_(ag__.converted_call(ag__.ld(get_enable_check_shapes), (), None, fscope)), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
     42 def get_state_8():
     43     return (arg_map, ag__.ldu(lambda : arg_map[RESULT_TOKEN], 'arg_map[RESULT_TOKEN]'), checker, do_return, retval_)

File ~\\AppData\\Local\\Temp\\__autograph_generated_filefvg6ssp8.py:15, in outer_factory.<locals>.inner_factory.<locals>.tf__training_loss(self)
     13 try:
     14     do_return = True
---> 15     retval_ = ag__.converted_call(ag__.ld(self)._training_loss, (), None, fscope)
     16 except:
     17     do_return = False

File ~\\AppData\\Local\\Temp\\__autograph_generated_file36u7ipsh.py:13, in outer_factory.<locals>.inner_factory.<locals>.tf__wrapped_method(self, *args, **kwargs)
     11 try:
     12     do_return = True
---> 13     retval_ = ag__.converted_call(ag__.ld(wrapped_function), (ag__.ld(self),) + tuple(ag__.ld(args)), dict(**ag__.ld(kwargs)), fscope)
     14 except:
     15     do_return = False

File ~\\AppData\\Local\\Temp\\__autograph_generated_file3aobj0zg.py:40, in outer_factory.<locals>.inner_factory.<locals>.tf__wrapped_function(*args, **kwargs)
     38     nonlocal retval_, do_return
     39     pass
---> 40 ag__.if_stmt(ag__.not_(ag__.converted_call(ag__.ld(get_enable_check_shapes), (), None, fscope)), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
     42 def get_state_8():
     43     return (arg_map, ag__.ldu(lambda : arg_map[RESULT_TOKEN], 'arg_map[RESULT_TOKEN]'), checker, do_return, retval_)

File ~\\AppData\\Local\\Temp\\__autograph_generated_fileqercka_v.py:12, in outer_factory.<locals>.inner_factory.<locals>.tf___training_loss(self)
     10 try:
     11     do_return = True
---> 12     retval_ = -ag__.converted_call(ag__.ld(self).log_posterior_density, (), None, fscope)
     13 except:
     14     do_return = False

File ~\\AppData\\Local\\Temp\\__autograph_generated_file36u7ipsh.py:13, in outer_factory.<locals>.inner_factory.<locals>.tf__wrapped_method(self, *args, **kwargs)
     11 try:
     12     do_return = True
---> 13     retval_ = ag__.converted_call(ag__.ld(wrapped_function), (ag__.ld(self),) + tuple(ag__.ld(args)), dict(**ag__.ld(kwargs)), fscope)
     14 except:
     15     do_return = False

File ~\\AppData\\Local\\Temp\\__autograph_generated_file3aobj0zg.py:40, in outer_factory.<locals>.inner_factory.<locals>.tf__wrapped_function(*args, **kwargs)
     38     nonlocal retval_, do_return
     39     pass
---> 40 ag__.if_stmt(ag__.not_(ag__.converted_call(ag__.ld(get_enable_check_shapes), (), None, fscope)), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
     42 def get_state_8():
     43     return (arg_map, ag__.ldu(lambda : arg_map[RESULT_TOKEN], 'arg_map[RESULT_TOKEN]'), checker, do_return, retval_)

File ~\\AppData\\Local\\Temp\\__autograph_generated_fileja4noeft.py:12, in outer_factory.<locals>.inner_factory.<locals>.tf__log_posterior_density(self)
     10 try:
     11     do_return = True
---> 12     retval_ = ag__.converted_call(ag__.ld(self).log_likelihood, (), None, fscope) + ag__.converted_call(ag__.ld(self).log_prior_density, (), None, fscope)
     13 except:
     14     do_return = False

File ~\\AppData\\Local\\Temp\\__autograph_generated_file36u7ipsh.py:13, in outer_factory.<locals>.inner_factory.<locals>.tf__wrapped_method(self, *args, **kwargs)
     11 try:
     12     do_return = True
---> 13     retval_ = ag__.converted_call(ag__.ld(wrapped_function), (ag__.ld(self),) + tuple(ag__.ld(args)), dict(**ag__.ld(kwargs)), fscope)
     14 except:
     15     do_return = False

File ~\\AppData\\Local\\Temp\\__autograph_generated_file3aobj0zg.py:40, in outer_factory.<locals>.inner_factory.<locals>.tf__wrapped_function(*args, **kwargs)
     38     nonlocal retval_, do_return
     39     pass
---> 40 ag__.if_stmt(ag__.not_(ag__.converted_call(ag__.ld(get_enable_check_shapes), (), None, fscope)), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
     42 def get_state_8():
     43     return (arg_map, ag__.ldu(lambda : arg_map[RESULT_TOKEN], 'arg_map[RESULT_TOKEN]'), checker, do_return, retval_)

File ~\\AppData\\Local\\Temp\\__autograph_generated_fileodijem7d.py:20, in outer_factory.<locals>.inner_factory.<locals>.tf__log_likelihood(self)
     18 K = ag__.converted_call(ag__.ld(self).kernel, (ag__.ld(X_data),), None, fscope)
     19 L = ag__.converted_call(ag__.ld(tf).linalg.cholesky, (ag__.ld(K) + ag__.converted_call(ag__.ld(tf).eye, (ag__.converted_call(ag__.ld(tf).shape, (ag__.ld(X_data),), None, fscope)[0],), dict(dtype=ag__.converted_call(ag__.ld(default_float), (), None, fscope)), fscope) * ag__.converted_call(ag__.ld(default_jitter), (), None, fscope),), None, fscope)
---> 20 F = ag__.converted_call(ag__.ld(tf).linalg.matmul, (ag__.ld(L), ag__.ld(self).V), None, fscope) + ag__.converted_call(ag__.ld(self).mean_function, (ag__.ld(X_data),), None, fscope)
     21 try:
     22     do_return = True

File ~\\AppData\\Local\\Temp\\__autograph_generated_fileqt6_ogiz.py:35, in outer_factory.<locals>.inner_factory.<locals>.tf____call__(self, X)
     33         do_return = False
     34         raise
---> 35 ag__.if_stmt(ag__.ld(self).A is None, if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
     36 return fscope.ret(retval_, do_return)

File ~\\AppData\\Local\\Temp\\__autograph_generated_fileqt6_ogiz.py:31, in outer_factory.<locals>.inner_factory.<locals>.tf____call__.<locals>.else_body()
     29 try:
     30     do_return = True
---> 31     retval_ = ag__.converted_call(ag__.ld(tf).matmul, (ag__.ld(X), ag__.ld(self).A), None, fscope) + ag__.ld(self).b
     32 except:
     33     do_return = False

ValueError: in user code:

    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\optimizers\\scipy.py\", line 185, in _tf_eval  *
        loss, grads = _compute_loss_and_gradients(
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\optimizers\\scipy.py\", line 322, in _compute_loss_and_gradients  *
        loss = loss_closure()
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\integration\\tf.py\", line 57, in wrapped_method  *
        return wrapped_function(self, *args, **kwargs)
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\decorator.py\", line 121, in wrapped_function  *
        if not get_enable_check_shapes():
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\models\\training_mixins.py\", line 64, in training_loss  *
        return self._training_loss()  # type: ignore[attr-defined]
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\integration\\tf.py\", line 57, in wrapped_method  *
        return wrapped_function(self, *args, **kwargs)
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\decorator.py\", line 121, in wrapped_function  *
        if not get_enable_check_shapes():
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\models\\gpmc.py\", line 81, in _training_loss  *
        return -self.log_posterior_density()
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\integration\\tf.py\", line 57, in wrapped_method  *
        return wrapped_function(self, *args, **kwargs)
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\decorator.py\", line 121, in wrapped_function  *
        if not get_enable_check_shapes():
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\models\\gpmc.py\", line 76, in log_posterior_density  *
        return self.log_likelihood() + self.log_prior_density()
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\integration\\tf.py\", line 57, in wrapped_method  *
        return wrapped_function(self, *args, **kwargs)
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\check_shapes\\decorator.py\", line 121, in wrapped_function  *
        if not get_enable_check_shapes():
    File \"j:\\100Anwendung\\_venv_\\lib\\site-packages\\gpflow\\models\\gpmc.py\", line 104, in log_likelihood  *
        F = tf.linalg.matmul(L, self.V) + self.mean_function(X_data)
    File \"C:\\Users\\Jonas\\AppData\\Local\\Temp\\ipykernel_14220\\4104737182.py\", line 12, in __call__  *
        return tf.matmul(X, self.A) + self.b

    ValueError: Shape must be rank 2 but is rank 1 for '{{node MatMul_2}} = MatMul[T=DT_DOUBLE, transpose_a=false, transpose_b=false](781, MatMul_2/identity_CONSTRUCTED_AT_top_level/forward/ReadVariableOp)' with input shapes: [1,20], [20].
"
}
```
yes!! I'm sorry! Since I was only interested in the first part, I didn't construct well the arrays. That's why you get that error. Shortly, I didn't check the shape of every prior over mean/cov functions
This is actually a bug in https://github.com/GPflow/GPflow/blob/a3a896e00de24d0548ad0dc8afd914f0386c281f/gpflow/functions.py#L113

Even if you pass in a `gpflow.Parameter` instance as `A`, the `np.atleast_2d` turns this back into a numpy array, and throws away any other information about the parameter (such as the prior).

The correct solution would be to check whether `A` is already a parameter, and if so, check whether the shape has 2 dimensions and otherwise raise an error. Would one of you be willing to fix this and prepare a pull request?
Hi @st-- 
yes I am willing to fix this issue. I am eager to do so. 
I'm currently in the final phase of my Master's thesis, but I'd like to sit down and do this after I've submitted it (end of November).

Resolving patch:
diff --git a/gpflow/functions.py b/gpflow/functions.py
--- a/gpflow/functions.py
+++ b/gpflow/functions.py
@@ -110,7 +110,15 @@ def __init__(self, A: TensorType = None, b: TensorType = None) -> None:
         MeanFunction.__init__(self)
         A = np.ones((1, 1), dtype=default_float()) if A is None else A
         b = np.zeros(1, dtype=default_float()) if b is None else b
-        self.A = Parameter(np.atleast_2d(A))
+        if isinstance(A, Parameter):
+            if len(A._shape) >= 2:
+                self.A = A
+            else:
+                raise ValueError(
+                    "Error 'gpflow.funcitons.Linear()' mean function. A has not the correct shape (at least 2d)."
+                )
+        else:
+            self.A = Parameter(np.atleast_2d(A))
         self.b = Parameter(b)
 
     @inherit_check_shapes


Test patch:
diff --git a/tests/gpflow/test_functions.py b/tests/gpflow/test_functions.py
--- a/tests/gpflow/test_functions.py
+++ b/tests/gpflow/test_functions.py
@@ -16,6 +16,7 @@
 
 import numpy as np
 import pytest
+import tensorflow_probability as tfp
 from check_shapes import check_shapes
 from numpy.testing import assert_allclose
 
@@ -371,3 +372,25 @@ def test_models_with_mean_functions_changes(model_class: Type[Any]) -> None:
     assert np.all(var_zero.numpy() == var_non_zero.numpy())
     # predictive mean changes after modifying mean function
     assert not np.all(mu_zero.numpy() == mu_non_zero.numpy())
+
+
+class TestIssue2091EnsureParameterTypeForLinearMeanFunction:
+    """
+    See github issue #2091. These are tests to ensure that the prior is kept if A is already given as Parameter.
+    """
+
+    # test_parameter_with_correct_shape
+    def test_parameter_with_correct_shape(self) -> None:
+        "Check that Linear copies the prior from a correctly shaped A Parameter"
+        A = gpflow.Parameter(
+            np.ones((1, 1)), dtype=np.float64, prior=tfp.distributions.Normal(0.0, 1.0)
+        )
+        linear_function = Linear(A, 1)
+        assert linear_function.A.prior is not None
+
+    # test_parameter_with_incorrect_shape
+    def test_parameter_with_incorrect_shape(self) -> None:
+        "Check that Linear throws an error when the A Parameter is not correctly shaped"
+        A = gpflow.Parameter(np.zeros(20), dtype=np.float64)
+        with pytest.raises(ValueError):
+            Linear(A, 1)