We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `gpflow/experimental/check_shapes/decorator.py`
Content:
```
1 # Copyright 2022 The GPflow Contributors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 # http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 """
15 Decorator for checking the shapes of function using tf Tensors.
16 """
17 import inspect
18 from functools import wraps
19 from typing import Any, Callable, Sequence, cast
20 
21 from ..utils import experimental
22 from .accessors import set_check_shapes
23 from .argument_ref import RESULT_TOKEN
24 from .base_types import C
25 from .checker import ShapeChecker
26 from .checker_context import set_shape_checker
27 from .config import get_enable_check_shapes
28 from .error_contexts import (
29     ConditionContext,
30     FunctionCallContext,
31     FunctionDefinitionContext,
32     NoteContext,
33     ParallelContext,
34     StackContext,
35 )
36 from .parser import parse_and_rewrite_docstring, parse_function_spec
37 from .specs import ParsedArgumentSpec
38 
39 
40 def null_check_shapes(func: C) -> C:
41     """
42     Annotates the given function so that it looks like it has shape checks, but without actually
43     checking anything.
44 
45     This is necessary not to break `@inherit_check_shapes` when shape checking is disabled.
46     """
47     set_check_shapes(func, null_check_shapes)
48     return func
49 
50 
51 @experimental
52 def check_shapes(*specs: str) -> Callable[[C], C]:
53     """
54     Decorator that checks the shapes of tensor arguments.
55 
56     :param specs: Specification of arguments to check. See: `Check specification`_.
57     """
58     if not get_enable_check_shapes():
59         return null_check_shapes
60 
61     unbound_error_context = FunctionCallContext(check_shapes)
62 
63     func_spec = parse_function_spec(specs, unbound_error_context)
64 
65     pre_specs = [spec for spec in func_spec.arguments if not spec.argument_ref.is_result]
66     post_specs = [spec for spec in func_spec.arguments if spec.argument_ref.is_result]
67     note_specs = func_spec.notes
68 
69     def _check_shapes(func: C) -> C:
70         bound_error_context = FunctionDefinitionContext(func)
71         signature = inspect.signature(func)
72 
73         @wraps(func)
74         def wrapped(*args: Any, **kwargs: Any) -> Any:
75             if not get_enable_check_shapes():
76                 return func(*args, **kwargs)
77 
78             try:
79                 bound_arguments = signature.bind(*args, **kwargs)
80             except TypeError as e:
81                 # TypeError is raised if *args and **kwargs don't actually match the arguments of
82                 # `func`. In that case we just call `func` normally, which will also result in an
83                 # error, but an error with the error message the user is used to.
84                 func(*args, **kwargs)
85                 raise AssertionError(
86                     "The above line should fail so this line should never be reached."
87                 ) from e
88             bound_arguments.apply_defaults()
89             arg_map = bound_arguments.arguments
90 
91             checker = ShapeChecker()
92             for note_spec in note_specs:
93                 checker.add_context(StackContext(bound_error_context, NoteContext(note_spec)))
94 
95             def _check_specs(specs: Sequence[ParsedArgumentSpec]) -> None:
96                 processed_specs = []
97 
98                 for arg_spec in specs:
99                     for arg_value, relative_arg_context in arg_spec.argument_ref.get(
100                         arg_map, bound_error_context
101                     ):
102                         arg_context = StackContext(bound_error_context, relative_arg_context)
103 
104                         if arg_spec.condition is not None:
105                             condition, condition_context = arg_spec.condition.get(
106                                 arg_map,
107                                 StackContext(arg_context, ConditionContext(arg_spec.condition)),
108                             )
109                             if not condition:
110                                 continue
111                             arg_context = StackContext(
112                                 bound_error_context,
113                                 ParallelContext(
114                                     (
115                                         StackContext(
116                                             relative_arg_context,
117                                             StackContext(
118                                                 ConditionContext(arg_spec.condition),
119                                                 condition_context,
120                                             ),
121                                         ),
122                                     )
123                                 ),
124                             )
125 
126                         processed_specs.append((arg_value, arg_spec.tensor, arg_context))
127 
128                 checker.check_shapes(processed_specs)
129 
130             _check_specs(pre_specs)
131 
132             with set_shape_checker(checker):
133                 result = func(*args, **kwargs)
134             arg_map[RESULT_TOKEN] = result
135 
136             _check_specs(post_specs)
137 
138             return result
139 
140         set_check_shapes(wrapped, _check_shapes)
141         wrapped.__doc__ = parse_and_rewrite_docstring(
142             wrapped.__doc__, func_spec, bound_error_context
143         )
144         return cast(C, wrapped)
145 
146     return _check_shapes
147
```

Path: `gpflow/experimental/check_shapes/exceptions.py`
Content:
```
1 # Copyright 2022 The GPflow Contributors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 # http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 # pylint: disable=broad-except
16 
17 """
18 Exceptions generated by `check_shapes`.
19 
20 These rely heavily on the infrastructure in `error_contexts.py`.
21 """
22 from .error_contexts import ErrorContext, MessageBuilder
23 
24 
25 class CheckShapesError(Exception):
26     """
27     Common super class for `check_shapes` errors.
28     """
29 
30     error_message: str
31 
32     def __init__(self, context: ErrorContext) -> None:
33         builder = MessageBuilder()
34         builder.add_line("")
35         builder.add_line(self.error_message)
36         with builder.indent() as b:
37             context.print(b)
38         super().__init__(builder.build())
39 
40         self.context = context
41 
42 
43 class VariableTypeError(CheckShapesError):
44     """
45     Error raised if a variable is used both as a rank-1 and a variable-rank variable.
46     """
47 
48     error_message = (
49         "Cannot use the same variable to bind both a single dimension"
50         " and a variable number of dimensions."
51     )
52 
53 
54 class SpecificationParseError(CheckShapesError):
55     """
56     Error raised if there was an error parsing the shape specification.
57     """
58 
59     error_message = "Unable to parse shape specification."
60 
61 
62 class DocstringParseError(CheckShapesError):
63     """
64     Error raised if there was an error parsing the shape specification.
65     """
66 
67     error_message = "Unable to parse docstring."
68 
69 
70 class ArgumentReferenceError(CheckShapesError):
71     """
72     Error raised if the argument to check the shape of could not be resolved.
73     """
74 
75     error_message = "Unable to resolve argument / missing argument."
76 
77 
78 class ShapeMismatchError(CheckShapesError):
79     """
80     Error raised if a function is called with tensors of the wrong shape.
81     """
82 
83     error_message = "Tensor shape mismatch."
84 
85 
86 class NoShapeError(CheckShapesError):
87     """
88     Error raised if we are trying to get the shape of an object that does not have a shape.
89     """
90 
91     error_message = "Unable to determine shape of object."
92
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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 @@
 
             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 @@
 
         self.context = context
 
+        # Prevent Keras from rewriting our exception:
+        self._keras_call_info_injected = True
+
 
 class VariableTypeError(CheckShapesError):
     """

{"golden_diff": "diff --git a/gpflow/experimental/check_shapes/decorator.py b/gpflow/experimental/check_shapes/decorator.py\n--- a/gpflow/experimental/check_shapes/decorator.py\n+++ b/gpflow/experimental/check_shapes/decorator.py\n@@ -15,9 +15,11 @@\n Decorator for checking the shapes of function using tf Tensors.\n \"\"\"\n import inspect\n-from functools import wraps\n+from functools import update_wrapper\n from typing import Any, Callable, Sequence, cast\n \n+import tensorflow as tf\n+\n from ..utils import experimental\n from .accessors import set_check_shapes\n from .argument_ref import RESULT_TOKEN\n@@ -70,7 +72,6 @@\n         bound_error_context = FunctionDefinitionContext(func)\n         signature = inspect.signature(func)\n \n-        @wraps(func)\n         def wrapped(*args: Any, **kwargs: Any) -> Any:\n             if not get_enable_check_shapes():\n                 return func(*args, **kwargs)\n@@ -137,10 +138,12 @@\n \n             return result\n \n+        # Make TensorFlow understand our decoration:\n+        tf.compat.v1.flags.tf_decorator.make_decorator(func, wrapped)\n+\n+        update_wrapper(wrapped, func)\n         set_check_shapes(wrapped, _check_shapes)\n-        wrapped.__doc__ = parse_and_rewrite_docstring(\n-            wrapped.__doc__, func_spec, bound_error_context\n-        )\n+        wrapped.__doc__ = parse_and_rewrite_docstring(func.__doc__, func_spec, bound_error_context)\n         return cast(C, wrapped)\n \n     return _check_shapes\ndiff --git a/gpflow/experimental/check_shapes/exceptions.py b/gpflow/experimental/check_shapes/exceptions.py\n--- a/gpflow/experimental/check_shapes/exceptions.py\n+++ b/gpflow/experimental/check_shapes/exceptions.py\n@@ -39,6 +39,9 @@\n \n         self.context = context\n \n+        # Prevent Keras from rewriting our exception:\n+        self._keras_call_info_injected = True\n+\n \n class VariableTypeError(CheckShapesError):\n     \"\"\"\n", "issue": "Bug check_shapes with keras layers\nHi @jesnie,\r\n\r\nI've been using the `check_shapes` decorator in eager the last couple of days and I really love it!\r\n\r\nI've encountered a small inconvenience when using it in combination with Keras layers. The following snippets show the weird behaviour:\r\n\r\n1. Custom keras layer without `check_shapes` decorator runs:\r\n```python\r\nfrom gpflow.experimental.check_shapes import check_shapes\r\nimport tensorflow as tf\r\n\r\nclass MyModel(tf.keras.Model):\r\n\r\n    def call(self, x, y):\r\n        return x + y\r\n\r\nlayer = MyModel()\r\nx = tf.random.normal((32, 3))\r\ny = tf.random.normal((32, 1))\r\nlayer(x, y)\r\n```\r\n\r\n2. Adding `check_shapes` decorator breaks the code:\r\n```python\r\nfrom gpflow.experimental.check_shapes import check_shapes\r\nimport tensorflow as tf\r\n\r\nclass MyModel(tf.keras.Model):\r\n\r\n    @check_shapes(\r\n        \"x: [batch, input_dim]\",\r\n        \"y: [batch, 1]\",\r\n        \"return: [batch, input_dim]\",\r\n    )\r\n    def call(self, x, y):\r\n        return x + y\r\n\r\nlayer = MyModel()\r\nx = tf.random.normal((32, 3))\r\ny = tf.random.normal((32, 1))\r\nlayer(x, y)\r\n```\r\nas now the layer expects a `training` kwarg to be added. The following error is thrown:\r\n```\r\n[TypeError: call() got an unexpected keyword argument 'training']()\r\n```\r\n\r\n3. After adding a `training` kwarg the code runs again. However, this isn't necessary when only using \"native\" keras.\r\n```python\r\nfrom gpflow.experimental.check_shapes import check_shapes\r\nimport tensorflow as tf\r\n\r\nclass MyModel(tf.keras.Model):\r\n\r\n    @check_shapes(\r\n        \"x: [batch, input_dim]\",\r\n        \"y: [batch, 1]\",\r\n        \"return: [batch, input_dim]\",\r\n    )\r\n    def call(self, x, y, training = None):\r\n        return x + y\r\n\r\nlayer = MyModel()\r\nx = tf.random.normal((32, 3))\r\ny = tf.random.normal((32, 1))\r\nlayer(x, y)\r\n```\r\n\r\nI'm using the following versions with python 3.8.13:\r\n```\r\ngpflow                  2.4.0\r\ntensorflow              2.5.3\r\n```\r\n\r\nMany thanks,\r\nVincent\n", "before_files": [{"content": "# Copyright 2022 The GPflow Contributors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nDecorator for checking the shapes of function using tf Tensors.\n\"\"\"\nimport inspect\nfrom functools import wraps\nfrom typing import Any, Callable, Sequence, cast\n\nfrom ..utils import experimental\nfrom .accessors import set_check_shapes\nfrom .argument_ref import RESULT_TOKEN\nfrom .base_types import C\nfrom .checker import ShapeChecker\nfrom .checker_context import set_shape_checker\nfrom .config import get_enable_check_shapes\nfrom .error_contexts import (\n    ConditionContext,\n    FunctionCallContext,\n    FunctionDefinitionContext,\n    NoteContext,\n    ParallelContext,\n    StackContext,\n)\nfrom .parser import parse_and_rewrite_docstring, parse_function_spec\nfrom .specs import ParsedArgumentSpec\n\n\ndef null_check_shapes(func: C) -> C:\n    \"\"\"\n    Annotates the given function so that it looks like it has shape checks, but without actually\n    checking anything.\n\n    This is necessary not to break `@inherit_check_shapes` when shape checking is disabled.\n    \"\"\"\n    set_check_shapes(func, null_check_shapes)\n    return func\n\n\n@experimental\ndef check_shapes(*specs: str) -> Callable[[C], C]:\n    \"\"\"\n    Decorator that checks the shapes of tensor arguments.\n\n    :param specs: Specification of arguments to check. See: `Check specification`_.\n    \"\"\"\n    if not get_enable_check_shapes():\n        return null_check_shapes\n\n    unbound_error_context = FunctionCallContext(check_shapes)\n\n    func_spec = parse_function_spec(specs, unbound_error_context)\n\n    pre_specs = [spec for spec in func_spec.arguments if not spec.argument_ref.is_result]\n    post_specs = [spec for spec in func_spec.arguments if spec.argument_ref.is_result]\n    note_specs = func_spec.notes\n\n    def _check_shapes(func: C) -> C:\n        bound_error_context = FunctionDefinitionContext(func)\n        signature = inspect.signature(func)\n\n        @wraps(func)\n        def wrapped(*args: Any, **kwargs: Any) -> Any:\n            if not get_enable_check_shapes():\n                return func(*args, **kwargs)\n\n            try:\n                bound_arguments = signature.bind(*args, **kwargs)\n            except TypeError as e:\n                # TypeError is raised if *args and **kwargs don't actually match the arguments of\n                # `func`. In that case we just call `func` normally, which will also result in an\n                # error, but an error with the error message the user is used to.\n                func(*args, **kwargs)\n                raise AssertionError(\n                    \"The above line should fail so this line should never be reached.\"\n                ) from e\n            bound_arguments.apply_defaults()\n            arg_map = bound_arguments.arguments\n\n            checker = ShapeChecker()\n            for note_spec in note_specs:\n                checker.add_context(StackContext(bound_error_context, NoteContext(note_spec)))\n\n            def _check_specs(specs: Sequence[ParsedArgumentSpec]) -> None:\n                processed_specs = []\n\n                for arg_spec in specs:\n                    for arg_value, relative_arg_context in arg_spec.argument_ref.get(\n                        arg_map, bound_error_context\n                    ):\n                        arg_context = StackContext(bound_error_context, relative_arg_context)\n\n                        if arg_spec.condition is not None:\n                            condition, condition_context = arg_spec.condition.get(\n                                arg_map,\n                                StackContext(arg_context, ConditionContext(arg_spec.condition)),\n                            )\n                            if not condition:\n                                continue\n                            arg_context = StackContext(\n                                bound_error_context,\n                                ParallelContext(\n                                    (\n                                        StackContext(\n                                            relative_arg_context,\n                                            StackContext(\n                                                ConditionContext(arg_spec.condition),\n                                                condition_context,\n                                            ),\n                                        ),\n                                    )\n                                ),\n                            )\n\n                        processed_specs.append((arg_value, arg_spec.tensor, arg_context))\n\n                checker.check_shapes(processed_specs)\n\n            _check_specs(pre_specs)\n\n            with set_shape_checker(checker):\n                result = func(*args, **kwargs)\n            arg_map[RESULT_TOKEN] = result\n\n            _check_specs(post_specs)\n\n            return result\n\n        set_check_shapes(wrapped, _check_shapes)\n        wrapped.__doc__ = parse_and_rewrite_docstring(\n            wrapped.__doc__, func_spec, bound_error_context\n        )\n        return cast(C, wrapped)\n\n    return _check_shapes\n", "path": "gpflow/experimental/check_shapes/decorator.py"}, {"content": "# Copyright 2022 The GPflow Contributors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# pylint: disable=broad-except\n\n\"\"\"\nExceptions generated by `check_shapes`.\n\nThese rely heavily on the infrastructure in `error_contexts.py`.\n\"\"\"\nfrom .error_contexts import ErrorContext, MessageBuilder\n\n\nclass CheckShapesError(Exception):\n    \"\"\"\n    Common super class for `check_shapes` errors.\n    \"\"\"\n\n    error_message: str\n\n    def __init__(self, context: ErrorContext) -> None:\n        builder = MessageBuilder()\n        builder.add_line(\"\")\n        builder.add_line(self.error_message)\n        with builder.indent() as b:\n            context.print(b)\n        super().__init__(builder.build())\n\n        self.context = context\n\n\nclass VariableTypeError(CheckShapesError):\n    \"\"\"\n    Error raised if a variable is used both as a rank-1 and a variable-rank variable.\n    \"\"\"\n\n    error_message = (\n        \"Cannot use the same variable to bind both a single dimension\"\n        \" and a variable number of dimensions.\"\n    )\n\n\nclass SpecificationParseError(CheckShapesError):\n    \"\"\"\n    Error raised if there was an error parsing the shape specification.\n    \"\"\"\n\n    error_message = \"Unable to parse shape specification.\"\n\n\nclass DocstringParseError(CheckShapesError):\n    \"\"\"\n    Error raised if there was an error parsing the shape specification.\n    \"\"\"\n\n    error_message = \"Unable to parse docstring.\"\n\n\nclass ArgumentReferenceError(CheckShapesError):\n    \"\"\"\n    Error raised if the argument to check the shape of could not be resolved.\n    \"\"\"\n\n    error_message = \"Unable to resolve argument / missing argument.\"\n\n\nclass ShapeMismatchError(CheckShapesError):\n    \"\"\"\n    Error raised if a function is called with tensors of the wrong shape.\n    \"\"\"\n\n    error_message = \"Tensor shape mismatch.\"\n\n\nclass NoShapeError(CheckShapesError):\n    \"\"\"\n    Error raised if we are trying to get the shape of an object that does not have a shape.\n    \"\"\"\n\n    error_message = \"Unable to determine shape of object.\"\n", "path": "gpflow/experimental/check_shapes/exceptions.py"}], "after_files": [{"content": "# Copyright 2022 The GPflow Contributors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nDecorator for checking the shapes of function using tf Tensors.\n\"\"\"\nimport inspect\nfrom functools import update_wrapper\nfrom typing import Any, Callable, Sequence, cast\n\nimport tensorflow as tf\n\nfrom ..utils import experimental\nfrom .accessors import set_check_shapes\nfrom .argument_ref import RESULT_TOKEN\nfrom .base_types import C\nfrom .checker import ShapeChecker\nfrom .checker_context import set_shape_checker\nfrom .config import get_enable_check_shapes\nfrom .error_contexts import (\n    ConditionContext,\n    FunctionCallContext,\n    FunctionDefinitionContext,\n    NoteContext,\n    ParallelContext,\n    StackContext,\n)\nfrom .parser import parse_and_rewrite_docstring, parse_function_spec\nfrom .specs import ParsedArgumentSpec\n\n\ndef null_check_shapes(func: C) -> C:\n    \"\"\"\n    Annotates the given function so that it looks like it has shape checks, but without actually\n    checking anything.\n\n    This is necessary not to break `@inherit_check_shapes` when shape checking is disabled.\n    \"\"\"\n    set_check_shapes(func, null_check_shapes)\n    return func\n\n\n@experimental\ndef check_shapes(*specs: str) -> Callable[[C], C]:\n    \"\"\"\n    Decorator that checks the shapes of tensor arguments.\n\n    :param specs: Specification of arguments to check. See: `Check specification`_.\n    \"\"\"\n    if not get_enable_check_shapes():\n        return null_check_shapes\n\n    unbound_error_context = FunctionCallContext(check_shapes)\n\n    func_spec = parse_function_spec(specs, unbound_error_context)\n\n    pre_specs = [spec for spec in func_spec.arguments if not spec.argument_ref.is_result]\n    post_specs = [spec for spec in func_spec.arguments if spec.argument_ref.is_result]\n    note_specs = func_spec.notes\n\n    def _check_shapes(func: C) -> C:\n        bound_error_context = FunctionDefinitionContext(func)\n        signature = inspect.signature(func)\n\n        def wrapped(*args: Any, **kwargs: Any) -> Any:\n            if not get_enable_check_shapes():\n                return func(*args, **kwargs)\n\n            try:\n                bound_arguments = signature.bind(*args, **kwargs)\n            except TypeError as e:\n                # TypeError is raised if *args and **kwargs don't actually match the arguments of\n                # `func`. In that case we just call `func` normally, which will also result in an\n                # error, but an error with the error message the user is used to.\n                func(*args, **kwargs)\n                raise AssertionError(\n                    \"The above line should fail so this line should never be reached.\"\n                ) from e\n            bound_arguments.apply_defaults()\n            arg_map = bound_arguments.arguments\n\n            checker = ShapeChecker()\n            for note_spec in note_specs:\n                checker.add_context(StackContext(bound_error_context, NoteContext(note_spec)))\n\n            def _check_specs(specs: Sequence[ParsedArgumentSpec]) -> None:\n                processed_specs = []\n\n                for arg_spec in specs:\n                    for arg_value, relative_arg_context in arg_spec.argument_ref.get(\n                        arg_map, bound_error_context\n                    ):\n                        arg_context = StackContext(bound_error_context, relative_arg_context)\n\n                        if arg_spec.condition is not None:\n                            condition, condition_context = arg_spec.condition.get(\n                                arg_map,\n                                StackContext(arg_context, ConditionContext(arg_spec.condition)),\n                            )\n                            if not condition:\n                                continue\n                            arg_context = StackContext(\n                                bound_error_context,\n                                ParallelContext(\n                                    (\n                                        StackContext(\n                                            relative_arg_context,\n                                            StackContext(\n                                                ConditionContext(arg_spec.condition),\n                                                condition_context,\n                                            ),\n                                        ),\n                                    )\n                                ),\n                            )\n\n                        processed_specs.append((arg_value, arg_spec.tensor, arg_context))\n\n                checker.check_shapes(processed_specs)\n\n            _check_specs(pre_specs)\n\n            with set_shape_checker(checker):\n                result = func(*args, **kwargs)\n            arg_map[RESULT_TOKEN] = result\n\n            _check_specs(post_specs)\n\n            return result\n\n        # Make TensorFlow understand our decoration:\n        tf.compat.v1.flags.tf_decorator.make_decorator(func, wrapped)\n\n        update_wrapper(wrapped, func)\n        set_check_shapes(wrapped, _check_shapes)\n        wrapped.__doc__ = parse_and_rewrite_docstring(func.__doc__, func_spec, bound_error_context)\n        return cast(C, wrapped)\n\n    return _check_shapes\n", "path": "gpflow/experimental/check_shapes/decorator.py"}, {"content": "# Copyright 2022 The GPflow Contributors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# pylint: disable=broad-except\n\n\"\"\"\nExceptions generated by `check_shapes`.\n\nThese rely heavily on the infrastructure in `error_contexts.py`.\n\"\"\"\nfrom .error_contexts import ErrorContext, MessageBuilder\n\n\nclass CheckShapesError(Exception):\n    \"\"\"\n    Common super class for `check_shapes` errors.\n    \"\"\"\n\n    error_message: str\n\n    def __init__(self, context: ErrorContext) -> None:\n        builder = MessageBuilder()\n        builder.add_line(\"\")\n        builder.add_line(self.error_message)\n        with builder.indent() as b:\n            context.print(b)\n        super().__init__(builder.build())\n\n        self.context = context\n\n        # Prevent Keras from rewriting our exception:\n        self._keras_call_info_injected = True\n\n\nclass VariableTypeError(CheckShapesError):\n    \"\"\"\n    Error raised if a variable is used both as a rank-1 and a variable-rank variable.\n    \"\"\"\n\n    error_message = (\n        \"Cannot use the same variable to bind both a single dimension\"\n        \" and a variable number of dimensions.\"\n    )\n\n\nclass SpecificationParseError(CheckShapesError):\n    \"\"\"\n    Error raised if there was an error parsing the shape specification.\n    \"\"\"\n\n    error_message = \"Unable to parse shape specification.\"\n\n\nclass DocstringParseError(CheckShapesError):\n    \"\"\"\n    Error raised if there was an error parsing the shape specification.\n    \"\"\"\n\n    error_message = \"Unable to parse docstring.\"\n\n\nclass ArgumentReferenceError(CheckShapesError):\n    \"\"\"\n    Error raised if the argument to check the shape of could not be resolved.\n    \"\"\"\n\n    error_message = \"Unable to resolve argument / missing argument.\"\n\n\nclass ShapeMismatchError(CheckShapesError):\n    \"\"\"\n    Error raised if a function is called with tensors of the wrong shape.\n    \"\"\"\n\n    error_message = \"Tensor shape mismatch.\"\n\n\nclass NoShapeError(CheckShapesError):\n    \"\"\"\n    Error raised if we are trying to get the shape of an object that does not have a shape.\n    \"\"\"\n\n    error_message = \"Unable to determine shape of object.\"\n", "path": "gpflow/experimental/check_shapes/exceptions.py"}]}