Repository: GPflow/GPflow

Issue:
Bug check_shapes with keras layers
Hi @jesnie,

I've been using the `check_shapes` decorator in eager the last couple of days and I really love it!

I've encountered a small inconvenience when using it in combination with Keras layers. The following snippets show the weird behaviour:

1. Custom keras layer without `check_shapes` decorator runs:
```python
from gpflow.experimental.check_shapes import check_shapes
import tensorflow as tf

class MyModel(tf.keras.Model):

    def call(self, x, y):
        return x + y

layer = MyModel()
x = tf.random.normal((32, 3))
y = tf.random.normal((32, 1))
layer(x, y)
```

2. Adding `check_shapes` decorator breaks the code:
```python
from gpflow.experimental.check_shapes import check_shapes
import tensorflow as tf

class MyModel(tf.keras.Model):

    @check_shapes(
        "x: [batch, input_dim]",
        "y: [batch, 1]",
        "return: [batch, input_dim]",
    )
    def call(self, x, y):
        return x + y

layer = MyModel()
x = tf.random.normal((32, 3))
y = tf.random.normal((32, 1))
layer(x, y)
```
as now the layer expects a `training` kwarg to be added. The following error is thrown:
```
[TypeError: call() got an unexpected keyword argument 'training']()
```

3. After adding a `training` kwarg the code runs again. However, this isn't necessary when only using "native" keras.
```python
from gpflow.experimental.check_shapes import check_shapes
import tensorflow as tf

class MyModel(tf.keras.Model):

    @check_shapes(
        "x: [batch, input_dim]",
        "y: [batch, 1]",
        "return: [batch, input_dim]",
    )
    def call(self, x, y, training = None):
        return x + y

layer = MyModel()
x = tf.random.normal((32, 3))
y = tf.random.normal((32, 1))
layer(x, y)
```

I'm using the following versions with python 3.8.13:
```
gpflow                  2.4.0
tensorflow              2.5.3
```

Many thanks,
Vincent


Discussion:

Resolving patch:
diff --git a/gpflow/experimental/check_shapes/decorator.py b/gpflow/experimental/check_shapes/decorator.py
--- a/gpflow/experimental/check_shapes/decorator.py
+++ b/gpflow/experimental/check_shapes/decorator.py
@@ -15,9 +15,11 @@
 Decorator for checking the shapes of function using tf Tensors.
 """
 import inspect
-from functools import wraps
+from functools import update_wrapper
 from typing import Any, Callable, Sequence, cast
 
+import tensorflow as tf
+
 from ..utils import experimental
 from .accessors import set_check_shapes
 from .argument_ref import RESULT_TOKEN
@@ -70,7 +72,6 @@ def _check_shapes(func: C) -> C:
         bound_error_context = FunctionDefinitionContext(func)
         signature = inspect.signature(func)
 
-        @wraps(func)
         def wrapped(*args: Any, **kwargs: Any) -> Any:
             if not get_enable_check_shapes():
                 return func(*args, **kwargs)
@@ -137,10 +138,12 @@ def _check_specs(specs: Sequence[ParsedArgumentSpec]) -> None:
 
             return result
 
+        # Make TensorFlow understand our decoration:
+        tf.compat.v1.flags.tf_decorator.make_decorator(func, wrapped)
+
+        update_wrapper(wrapped, func)
         set_check_shapes(wrapped, _check_shapes)
-        wrapped.__doc__ = parse_and_rewrite_docstring(
-            wrapped.__doc__, func_spec, bound_error_context
-        )
+        wrapped.__doc__ = parse_and_rewrite_docstring(func.__doc__, func_spec, bound_error_context)
         return cast(C, wrapped)
 
     return _check_shapes
diff --git a/gpflow/experimental/check_shapes/exceptions.py b/gpflow/experimental/check_shapes/exceptions.py
--- a/gpflow/experimental/check_shapes/exceptions.py
+++ b/gpflow/experimental/check_shapes/exceptions.py
@@ -39,6 +39,9 @@ def __init__(self, context: ErrorContext) -> None:
 
         self.context = context
 
+        # Prevent Keras from rewriting our exception:
+        self._keras_call_info_injected = True
+
 
 class VariableTypeError(CheckShapesError):
     """


Test patch:
diff --git a/tests/gpflow/experimental/check_shapes/test_integration.py b/tests/gpflow/experimental/check_shapes/test_integration.py
--- a/tests/gpflow/experimental/check_shapes/test_integration.py
+++ b/tests/gpflow/experimental/check_shapes/test_integration.py
@@ -23,8 +23,13 @@
 import tensorflow as tf
 
 from gpflow.base import AnyNDArray
-from gpflow.experimental.check_shapes import check_shapes
-from gpflow.experimental.check_shapes.config import disable_check_shapes
+from gpflow.experimental.check_shapes import ShapeMismatchError
+from gpflow.experimental.check_shapes import check_shape as cs
+from gpflow.experimental.check_shapes import (
+    check_shapes,
+    disable_check_shapes,
+    inherit_check_shapes,
+)
 
 
 def test_check_shapes__numpy() -> None:
@@ -51,6 +56,64 @@ def f(a: tf.Tensor, b: tf.Tensor) -> tf.Tensor:
     f(tf.zeros((2, 3)), tf.zeros((2, 4)))  # Don't crash...
 
 
+def test_check_shapes__tensorflow__keras() -> None:
+    # pylint: disable=arguments-differ,abstract-method,no-value-for-parameter,unexpected-keyword-arg
+
+    class SuperLayer(tf.keras.layers.Layer):
+        def __init__(self) -> None:
+            super().__init__()
+            self._b = tf.Variable(0.0)
+
+        @check_shapes(
+            "x: [batch, input_dim]",
+            "y: [batch, 1]",
+            "return: [batch, input_dim]",
+        )
+        def call(self, x: tf.Tensor, y: tf.Tensor) -> tf.Tensor:
+            return x + y + self._b
+
+    class SubLayer(SuperLayer):
+        @inherit_check_shapes
+        def call(self, x: tf.Tensor, y: tf.Tensor) -> tf.Tensor:
+            return x - y + self._b
+
+    class MyModel(tf.keras.Model):
+        def __init__(self, join: SuperLayer) -> None:
+            super().__init__()
+            self._join = join
+
+        @check_shapes(
+            "xy: [batch, input_dim_plus_one]",
+            "return: [batch, input_dim]",
+        )
+        def call(self, xy: tf.Tensor) -> tf.Tensor:
+            x = cs(xy[:, :-1], "[batch, input_dim]")
+            y = cs(xy[:, -1:], "[batch, 1]")
+            return self._join(x, y)
+
+    x = tf.ones((32, 3))
+    y = tf.zeros((32, 1))
+    xy = tf.concat([x, y], axis=1)
+    y_bad = tf.zeros((32, 2))
+    targets = tf.ones((32, 3))
+
+    def test_layer(join: SuperLayer) -> None:
+        join(x, y)
+
+        with pytest.raises(ShapeMismatchError):
+            join(x, y_bad)
+
+        model = MyModel(join)
+        model.compile(
+            optimizer=tf.keras.optimizers.SGD(learning_rate=0.25),
+            loss="mean_squared_error",
+        )
+        model.fit(x=xy, y=targets)
+
+    test_layer(SuperLayer())
+    test_layer(SubLayer())
+
+
 _F = Callable[[tf.Tensor], tf.Tensor]
 _Loss = Callable[[], tf.Tensor]