Trouble with Linear mean_function with heteroscedastic GP

[risaueno]

Bug / performance issue / build issue

I am trying to set the prior mean function to X=Y in the model specification in a heteroskedastic regression. However Linear mean_function where the value of A is 1. does not work, whereas 0 does.

To reproduce

This can be reproduced using this gpflow tutorial https://gpflow.github.io/GPflow/2.4.0/notebooks/advanced/heteroskedastic.html but adding the mean_function as follows:

mean_f = gpf.mean_functions.Linear(A=np.array([[1.]]))
model = gpf.models.SVGP(
    kernel=kernel,
    likelihood=likelihood,
    inducing_variable=inducing_variable,
    num_latent_gps=likelihood.latent_dim,
    mean_function=mean_f,
)

This works fine: mean_f = gpf.mean_functions.Linear(A=np.array([[0]]))

Stack trace, or error message

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
[/data/hpcdata/users/risno/code/climate-gpbc/bcml/experiments/1_3_gpqmn/develop/heteroscedastic_datapoints.py](https://vscode-remote+ssh-002dremote-002bhpc6.vscode-resource.vscode-cdn.net/data/hpcdata/users/risno/code/climate-gpbc/bcml/experiments/1_3_gpqmn/develop/heteroscedastic_datapoints.py) in line 62
     735 logf = []
     736 for epoch in range(1, epochs + 1):
---> 737     optimisation_step()
     739     # For every 'log_freq' epochs, print the epoch and plot the predictions against the data
     740     if epoch % log_freq == 0 and epoch > 0:

File [/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py:780](https://vscode-remote+ssh-002dremote-002bhpc6.vscode-resource.vscode-cdn.net/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py:780), in Function.__call__(self, *args, **kwds)
    778 else:
    779   compiler = "nonXla"
--> 780   result = self._call(*args, **kwds)
    782 new_tracing_count = self._get_tracing_count()
    783 without_tracing = (tracing_count == new_tracing_count)

File [/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py:807](https://vscode-remote+ssh-002dremote-002bhpc6.vscode-resource.vscode-cdn.net/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py:807), in Function._call(self, *args, **kwds)
    804   self._lock.release()
    805   # In this case we have created variables on the first call, so we run the
    806   # defunned version which is guaranteed to never create variables.
--> 807   return self._stateless_fn(*args, **kwds)  # pylint: disable=not-callable
    808 elif self._stateful_fn is not None:
    809   # Release the lock early so that multiple threads can perform the call
    810   # in parallel.
    811   self._lock.release()

File [/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/function.py:2829](https://vscode-remote+ssh-002dremote-002bhpc6.vscode-resource.vscode-cdn.net/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/function.py:2829), in Function.__call__(self, *args, **kwargs)
   2827 with self._lock:
   2828   graph_function, args, kwargs = self._maybe_define_function(args, kwargs)
-> 2829 return graph_function._filtered_call(args, kwargs)

File [/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/function.py:1843](https://vscode-remote+ssh-002dremote-002bhpc6.vscode-resource.vscode-cdn.net/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/function.py:1843), in ConcreteFunction._filtered_call(self, args, kwargs, cancellation_manager)
   1827 def _filtered_call(self, args, kwargs, cancellation_manager=None):
   1828   """Executes the function, filtering arguments from the Python function.
   1829 
   1830   Objects aside from Tensors, CompositeTensors, and Variables are ignored.
   (...)
   1841     `args` and `kwargs`.
   1842   """
-> 1843   return self._call_flat(
   1844       [t for t in nest.flatten((args, kwargs), expand_composites=True)
   1845        if isinstance(t, (ops.Tensor,
   1846                          resource_variable_ops.BaseResourceVariable))],
   1847       captured_inputs=self.captured_inputs,
   1848       cancellation_manager=cancellation_manager)

File [/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/function.py:1923](https://vscode-remote+ssh-002dremote-002bhpc6.vscode-resource.vscode-cdn.net/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/function.py:1923), in ConcreteFunction._call_flat(self, args, captured_inputs, cancellation_manager)
   1918 possible_gradient_type = (
   1919     pywrap_tfe.TFE_Py_TapeSetPossibleGradientTypes(args))
   1920 if (possible_gradient_type == _POSSIBLE_GRADIENT_TYPES_NONE
   1921     and executing_eagerly):
   1922   # No tape is watching; skip to running the function.
-> 1923   return self._build_call_outputs(self._inference_function.call(
   1924       ctx, args, cancellation_manager=cancellation_manager))
   1925 forward_backward = self._select_forward_and_backward_functions(
   1926     args,
   1927     possible_gradient_type,
   1928     executing_eagerly)
   1929 forward_function, args_with_tangents = forward_backward.forward()

File [/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/function.py:545](https://vscode-remote+ssh-002dremote-002bhpc6.vscode-resource.vscode-cdn.net/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/function.py:545), in _EagerDefinedFunction.call(self, ctx, args, cancellation_manager)
    543 with _InterpolateFunctionError(self):
    544   if cancellation_manager is None:
--> 545     outputs = execute.execute(
    546         str(self.signature.name),
    547         num_outputs=self._num_outputs,
    548         inputs=args,
    549         attrs=attrs,
    550         ctx=ctx)
    551   else:
    552     outputs = execute.execute_with_cancellation(
    553         str(self.signature.name),
    554         num_outputs=self._num_outputs,
   (...)
    557         ctx=ctx,
    558         cancellation_manager=cancellation_manager)

File [/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/execute.py:59](https://vscode-remote+ssh-002dremote-002bhpc6.vscode-resource.vscode-cdn.net/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/tensorflow/python/eager/execute.py:59), in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     57 try:
     58   ctx.ensure_initialized()
---> 59   tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     60                                       inputs, attrs, num_outputs)
     61 except core._NotOkStatusException as e:
     62   if name is not None:

InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  Cholesky decomposition was not successful. The input might not be valid.
	 [[node NaturalGradient/natural_gradient_steps/Cholesky (defined at [/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/gpflow/optimizers/natgrad.py:395](https://vscode-remote+ssh-002dremote-002bhpc6.vscode-resource.vscode-cdn.net/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/gpflow/optimizers/natgrad.py:395)) ]]
	 [[Func/StatefulPartitionedCall_1/input_control_node/_276/_47]]
  (1) Invalid argument:  Cholesky decomposition was not successful. The input might not be valid.
	 [[node NaturalGradient/natural_gradient_steps/Cholesky (defined at [/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/gpflow/optimizers/natgrad.py:395](https://vscode-remote+ssh-002dremote-002bhpc6.vscode-resource.vscode-cdn.net/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/gpflow/optimizers/natgrad.py:395)) ]]
0 successful operations.
0 derived errors ignored. [Op:__inference_optimisation_step_265645]

Errors may have originated from an input operation.
Input Source operations connected to node NaturalGradient/natural_gradient_steps/Cholesky:
 NaturalGradient/natural_gradient_steps/sub (defined at [/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/gpflow/optimizers/natgrad.py:394](https://vscode-remote+ssh-002dremote-002bhpc6.vscode-resource.vscode-cdn.net/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/gpflow/optimizers/natgrad.py:394))

Input Source operations connected to node NaturalGradient/natural_gradient_steps/Cholesky:
 NaturalGradient/natural_gradient_steps/sub (defined at [/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/gpflow/optimizers/natgrad.py:394](https://vscode-remote+ssh-002dremote-002bhpc6.vscode-resource.vscode-cdn.net/data/hpcdata/users/risno/miniconda3/envs/py38/lib/python3.8/site-packages/gpflow/optimizers/natgrad.py:394))

Function call stack:
optimisation_step -> optimisation_step

Expected behavior

System information

• GPflow version: 2.3.0
• GPflow installed from: pip install gpflow
• TensorFlow version: 2.3.0
• Python version: 3.8.13
• Operating system CentOS Linux release 7.9.2009 (Core)