Repository: GPflow/GPflow

Issue:
Handling unconnected gradients in Scipy optimizer
# Feature request

Our Scipy optimizer wrapper should explicitly handle unconnected gradients.

## Motivation

Currently, the GPflow Scipy optimizer wrapper (`gpflow.optimizers.Scipy`) simply errors (with one of the typically unhelpful TensorFlow error messages) when one of the passed-in `variables` is not connected to the loss to be minimized. This can easily happen in a valid case, for example when using an SVGP model with a `Constant` kernel (which might be used as a very simple baseline) but having inducing locations not explicitly set to non-trainable. This could also happen erroneously in a bug case, e.g. when calling `Scipy().minimize(model1.training_loss, model2.trainable_variables)`.

## Proposal

**Describe the solution you would like**
An explicit error message when the variables are not connected to the loss (and hence the `grads` list contains `None` entries).

**What alternatives have you considered?**
We could also simply change the call for the gradient computation to use `unconnected_gradients=tf.UnconnectedGradients.ZERO`. This might silently hide bug cases (see above), so perhaps this option could be an extra flag to the minimize() call?



Discussion:
Hey,
If this issue is still active, I would be happy to work on it. 
One question: is checking whether one entry of `gradient` is a zero tensor with `tf.math.count_nonzero` appropriate? 
Hi @antonykamp ! thanks for your interest. Yes, nothing has changed about this issue - if you try to optimize a loss that's not connected to the variables you pass in, you still get an unhelpful error message.
Regarding your question, can you expand a bit more on how you'd want to use count_nonzero?
> Can you expand a bit more on how you'd want to use count_nonzero?

Of course! When using `unconnected_gradients=tf.UnconnectedGradients.ZERO` and the loss is not connected to the variables, the associated item in the gradient equals a zero tensor. GPflow could test for this by counting the nonzeros of each item. Should the situation arise that every value equals zero, GPflow warns the user, [eg](https://github.com/antonykamp/GPflow/blob/599bc49ed902c2296a550c8f0409dcc9bdcef900/gpflow/optimizers/scipy.py#L181-L182).
Ah! You mean to give a warning, yet continue running?
Will the if check you added in your link work correctly within tf.function()? even if it does, it might make the call a bit slower. How about moving the check to https://github.com/antonykamp/GPflow/blob/599bc49ed902c2296a550c8f0409dcc9bdcef900/gpflow/optimizers/scipy.py#L119 (where we've got to be in numpy-land anyways), and maybe having a flag so it only actually checks on the first call ?
> Ah! You mean to give a warning, yet continue running?

As you pointed out, sometimes its desired behavior eg. in bug cases. Am I right? :)

Moving the check was a great idea ofc :D I would add the new flag as a simple attribute of `Scipy`-objects (and `uncon_gradients_handle` stays?)
Would you like to open a PR so we can discuss more concretely on the code? :)

Resolving patch:
diff --git a/gpflow/optimizers/scipy.py b/gpflow/optimizers/scipy.py
--- a/gpflow/optimizers/scipy.py
+++ b/gpflow/optimizers/scipy.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
 from typing import Callable, Iterable, List, Optional, Sequence, Tuple, TypeVar, Union
 
 import numpy as np
@@ -36,6 +37,7 @@ def minimize(
         method: Optional[str] = "L-BFGS-B",
         step_callback: Optional[StepCallback] = None,
         compile: bool = True,
+        allow_unused_variables: bool = False,
         **scipy_kwargs,
     ) -> OptimizeResult:
         """
@@ -59,7 +61,8 @@ def minimize(
             compile: If True, wraps the evaluation function (the passed `closure`
                 as well as its gradient computation) inside a `tf.function()`,
                 which will improve optimization speed in most cases.
-
+            allow_unused_variables: Whether to allow variables that are not
+                actually used in the closure.
             scipy_kwargs: Arguments passed through to `scipy.optimize.minimize`
                 Note that Scipy's minimize() takes a `callback` argument, but
                 you probably want to use our wrapper and pass in `step_callback`.
@@ -79,7 +82,9 @@ def minimize(
             )  # pragma: no cover
         initial_params = self.initial_parameters(variables)
 
-        func = self.eval_func(closure, variables, compile=compile)
+        func = self.eval_func(
+            closure, variables, compile=compile, allow_unused_variables=allow_unused_variables
+        )
         if step_callback is not None:
             if "callback" in scipy_kwargs:
                 raise ValueError("Callback passed both via `step_callback` and `callback`")
@@ -97,13 +102,32 @@ def initial_parameters(cls, variables: Sequence[tf.Variable]) -> tf.Tensor:
 
     @classmethod
     def eval_func(
-        cls, closure: LossClosure, variables: Sequence[tf.Variable], compile: bool = True
+        cls,
+        closure: LossClosure,
+        variables: Sequence[tf.Variable],
+        compile: bool = True,
+        allow_unused_variables: bool = False,
     ) -> Callable[[np.ndarray], Tuple[np.ndarray, np.ndarray]]:
+        first_call = True
+
         def _tf_eval(x: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
+            nonlocal first_call
+
             values = cls.unpack_tensors(variables, x)
             cls.assign_tensors(variables, values)
 
-            loss, grads = _compute_loss_and_gradients(closure, variables)
+            if first_call:
+                # Only check for unconnected gradients on the first function evaluation.
+                loss, grads = _compute_loss_and_gradients(
+                    closure, variables, tf.UnconnectedGradients.NONE
+                )
+                grads = cls._filter_unused_variables(variables, grads, allow_unused_variables)
+                first_call = False
+            else:
+                loss, grads = _compute_loss_and_gradients(
+                    closure, variables, tf.UnconnectedGradients.ZERO
+                )
+
             return loss, cls.pack_tensors(grads)
 
         if compile:
@@ -115,6 +139,32 @@ def _eval(x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
 
         return _eval
 
+    @staticmethod
+    def _filter_unused_variables(
+        variables: Sequence[tf.Variable], grads: Sequence[tf.Tensor], allow_unused_variables: bool
+    ) -> Sequence[tf.Tensor]:
+        filtered_grads = []
+        unused_variables = []
+        for i, grad in enumerate(grads):
+            if grad is None:
+                variable = variables[i]
+                filtered_grads.append(tf.zeros_like(variable))
+                unused_variables.append(variable.name)
+            else:
+                filtered_grads.append(grad)
+
+        if unused_variables:
+            msg = (
+                "Some variables does not have a gradient, and appear unused in / not connected to"
+                f" the loss closure: {unused_variables}."
+            )
+            if allow_unused_variables:
+                warnings.warn(msg)
+            else:
+                raise ValueError(msg)
+
+        return filtered_grads
+
     @classmethod
     def callback_func(
         cls, variables: Sequence[tf.Variable], step_callback: StepCallback
@@ -165,10 +215,12 @@ def assign_tensors(to_tensors: Sequence[tf.Variable], values: Sequence[tf.Tensor
 
 
 def _compute_loss_and_gradients(
-    loss_closure: LossClosure, variables: Sequence[tf.Variable]
+    loss_closure: LossClosure,
+    variables: Sequence[tf.Variable],
+    unconnected_gradients: tf.UnconnectedGradients,
 ) -> Tuple[tf.Tensor, Sequence[tf.Tensor]]:
     with tf.GradientTape(watch_accessed_variables=False) as tape:
         tape.watch(variables)
         loss = loss_closure()
-    grads = tape.gradient(loss, variables)
+    grads = tape.gradient(loss, variables, unconnected_gradients=unconnected_gradients)
     return loss, grads


Test patch:
diff --git a/tests/gpflow/optimizers/test_scipy.py b/tests/gpflow/optimizers/test_scipy.py
--- a/tests/gpflow/optimizers/test_scipy.py
+++ b/tests/gpflow/optimizers/test_scipy.py
@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
+
 import numpy as np
 import pytest
 import tensorflow as tf
 from numpy.testing import assert_allclose
 
 import gpflow
-from gpflow.config import default_jitter
+from gpflow.config import default_float, default_jitter
 from gpflow.mean_functions import Constant
 
 rng = np.random.RandomState(0)
@@ -70,3 +72,83 @@ def get_values(model):
     # due to the changes introduced by PR #1213, which removed some implicit casts
     # to float32. With TF 2.4 / TFP 0.12, we had to further increase atol from 1e-14.
     np.testing.assert_allclose(get_values(m1), get_values(m2), rtol=1e-14, atol=2e-14)
+
+
+@pytest.mark.parametrize("compile", [True, False])
+def test_scipy__optimal(compile: bool) -> None:
+    target1 = [0.2, 0.8]
+    target2 = [0.6]
+    v1 = tf.Variable([0.5, 0.5], dtype=default_float())
+    v2 = tf.Variable([0.5], dtype=default_float())
+    compilation_count = 0
+
+    def f() -> tf.Tensor:
+        nonlocal compilation_count
+        compilation_count += 1
+        return tf.reduce_sum((target1 - v1) ** 2) + tf.reduce_sum((target2 - v2) ** 2)
+
+    opt = gpflow.optimizers.Scipy()
+    result = opt.minimize(f, [v1, v2], compile=compile)
+
+    if compile:
+        assert 1 == compilation_count
+    else:
+        assert 1 < compilation_count
+    assert result.success
+    np.testing.assert_allclose(target1 + target2, result.x)
+    np.testing.assert_allclose(target1, v1)
+    np.testing.assert_allclose(target2, v2)
+
+
+@pytest.mark.parametrize("compile", [True, False])
+def test_scipy__partially_disconnected_variable(compile: bool) -> None:
+    target1 = 0.2
+    target2 = 0.6
+    v1 = tf.Variable([0.5, 0.5], dtype=default_float())
+    v2 = tf.Variable(0.5, dtype=default_float())
+
+    def f() -> tf.Tensor:
+        # v1[1] not used.
+        v10 = v1[0]
+        return (target1 - v10) ** 2 + (target2 - v2) ** 2
+
+    opt = gpflow.optimizers.Scipy()
+    result = opt.minimize(f, [v1, v2], compile=compile)
+
+    assert result.success
+    np.testing.assert_allclose([target1, 0.5, target2], result.x)
+    np.testing.assert_allclose([target1, 0.5], v1)
+    np.testing.assert_allclose(target2, v2)
+
+
+@pytest.mark.parametrize("compile", [True, False])
+@pytest.mark.parametrize("allow_unused_variables", [True, False])
+def test_scipy__disconnected_variable(compile: bool, allow_unused_variables: bool) -> None:
+    target1 = [0.2, 0.8]
+    v1 = tf.Variable([0.5, 0.5], dtype=default_float(), name="v1")
+    v2 = tf.Variable([0.5], dtype=default_float(), name="v2")
+
+    def f() -> tf.Tensor:
+        # v2 not used.
+        return tf.reduce_sum((target1 - v1) ** 2)
+
+    opt = gpflow.optimizers.Scipy()
+
+    if allow_unused_variables:
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            result = opt.minimize(
+                f, [v1, v2], compile=compile, allow_unused_variables=allow_unused_variables
+            )
+
+        (warning,) = w
+        msg = warning.message.args[0]
+        assert v2.name in msg
+
+        assert result.success
+        np.testing.assert_allclose(target1 + [0.5], result.x)
+        np.testing.assert_allclose(target1, v1)
+        np.testing.assert_allclose([0.5], v2)
+    else:
+        with pytest.raises(ValueError, match=v2.name):
+            opt.minimize(f, [v1, v2], allow_unused_variables=allow_unused_variables)