We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
print_summary does not display information on parameter priors
### `print_summary` does not display information on parameter priors

GPflow 2 (`develop-2.0` branch):
`print_summary` on a model with priors does not have a column for the `prior` attribute.

MWE:
```python
import gpflow
import tensorflow_probability as tfp
from gpflow.utilities import to_default_float
from gpflow.utilities import print_summary

k = gpflow.kernels.Matern32()
k.variance.prior = tfp.distributions.Gamma(
    to_default_float(2), to_default_float(3))

print_summary(k)
```
print_summary does not display information on parameter priors
### `print_summary` does not display information on parameter priors

GPflow 2 (`develop-2.0` branch):
`print_summary` on a model with priors does not have a column for the `prior` attribute.

MWE:
```python
import gpflow
import tensorflow_probability as tfp
from gpflow.utilities import to_default_float
from gpflow.utilities import print_summary

k = gpflow.kernels.Matern32()
k.variance.prior = tfp.distributions.Gamma(
    to_default_float(2), to_default_float(3))

print_summary(k)
```

--- 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/utilities/utilities.py`
Content:
```
1 import re
2 from copy import deepcopy
3 from functools import lru_cache
4 from typing import Callable, Dict, List, Optional, Union, TypeVar, Any, Tuple
5 
6 import numpy as np
7 import tensorflow as tf
8 import tensorflow_probability as tfp
9 from tabulate import tabulate
10 
11 from ..base import Parameter
12 from ..config import default_summary_fmt, default_float, default_int
13 
14 __all__ = [
15     "set_trainable",
16     "multiple_assign",
17     "training_loop",
18     "print_summary",
19     "deepcopy_components",
20     "leaf_components",
21     "parameter_dict",
22     "read_values",
23     "to_default_float",
24     "to_default_int",
25 ]
26 
27 TraverseInput = TypeVar("TraverseInput", tf.Variable, tf.Module, Parameter)
28 State = Any
29 Path = str
30 Accumulator = Tuple[Path, State]
31 TraverseUpdateCallable = Callable[[TraverseInput, Path, State], State]
32 
33 
34 def to_default_int(x):
35     return tf.cast(x, dtype=default_int())
36 
37 
38 def to_default_float(x):
39     return tf.cast(x, dtype=default_float())
40 
41 
42 def set_trainable(model: tf.Module, flag: bool):
43     """
44     Set trainable flag for all `tf.Variable`s and `gpflow.Parameter`s in a module.
45     """
46     for variable in model.variables:
47         variable._trainable = flag
48 
49 
50 def multiple_assign(module: tf.Module, parameters: Dict[str, tf.Tensor]):
51     """
52     Multiple assign takes a dictionary with new values. Dictionary keys are paths to the
53     `tf.Variable`s or `gpflow.Parameter` of the input module.
54 
55     :param module: `tf.Module`.
56     :param parameters: a dictionary with keys of the form ".module.path.to.variable" and new value tensors.
57     """
58     reference_var_dict = parameter_dict(module)
59     for path, value in parameters.items():
60         reference_var_dict[path].assign(value)
61 
62 
63 def read_values(module: tf.Module) -> Dict[str, np.ndarray]:
64     """Returns a dictionary of numpy values of the module parameters (variables)."""
65     return {k: v.numpy() for k, v in parameter_dict(module).items()}
66 
67 
68 def parameter_dict(module: tf.Module) -> Dict[str, Union[Parameter, tf.Variable]]:
69     """
70     Returns a dictionary of parameters (variables) for the `tf.Module` component.
71     Dictionary keys are relative paths to the attributes to which parameters (variables) assigned to.
72 
73         class SubModule(tf.Module):
74             def __init__(self):
75                 self.parameter = gpflow.Parameter(1.0)
76                 self.variable = tf.Variable(1.0)
77 
78         class Module(tf.Module):
79             def __init__(self):
80                 self.submodule = SubModule()
81 
82         m = Module()
83         params = parameter_dict(m)
84         # {
85         #   ".submodule.parameter": <parameter object>,
86         #   ".submodule.variable": <variable object>
87         # }
88     """
89     param_dict = leaf_components(module)
90     return {f".{key.split('.', 1)[-1]}": value for key, value in param_dict.items()}
91 
92 
93 def training_loop(closure: Callable[..., tf.Tensor],
94                   optimizer: Optional[tf.optimizers.Optimizer] = None,
95                   var_list: List[tf.Variable] = None,
96                   maxiter=1e3,
97                   jit=False):
98     """
99     Simple generic training loop. At each iteration uses a GradientTape to compute
100     the gradients of a loss function with respect to a set of variables.
101 
102     :param closure: Callable that constructs a loss function based on data and model being trained
103     :param optimizer: tf.optimizers or tf.keras.optimizers that updates variables by applying the
104         corresponding loss gradients. Adam is a default optimizer with default settings.
105     :param var_list: List of model variables to be learnt during training
106     :param maxiter: Maximum number of
107     :return:
108     """
109 
110     optimizer = tf.optimizers.Adam() if optimizer is None else optimizer
111 
112     def optimization_step():
113         with tf.GradientTape() as tape:
114             tape.watch(var_list)
115             loss = closure()
116         grads = tape.gradient(loss, var_list)
117         optimizer.apply_gradients(zip(grads, var_list))
118 
119     if jit:
120         optimization_step = tf.function(optimization_step)
121 
122     for _ in range(int(maxiter)):
123         optimization_step()
124 
125 
126 def print_summary(module: tf.Module, fmt: str = None):
127     """
128     Prints a summary of the parameters and variables contained in a tf.Module.
129     """
130     fmt = fmt if fmt is not None else default_summary_fmt()
131     if fmt == "notebook":
132         from IPython.core.display import display, HTML
133         tab = tabulate_module_summary(module, "html")
134         display(HTML(tab))
135     else:
136         print(tabulate_module_summary(module, fmt))
137 
138 
139 def tabulate_module_summary(module: tf.Module, tablefmt: Optional[str] = None) -> str:
140     column_names = ['name', 'class', 'transform', 'trainable', 'shape', 'dtype', 'value']
141 
142     def get_name(v):
143         return v.__class__.__name__
144 
145     def get_transform(v):
146         if hasattr(v, "transform") and v.transform is not None:
147             if isinstance(v.transform, tfp.bijectors.Chain):
148                 return " + ".join(b.__class__.__name__ for b in v.transform.bijectors[::-1])
149             return v.transform.__class__.__name__
150         return None
151 
152     merged_leaf_components = _merge_leaf_components(leaf_components(module))
153 
154     column_values = [[
155         path,
156         get_name(variable),
157         get_transform(variable), variable.trainable, variable.shape, variable.dtype.name,
158         _str_tensor_value(variable.numpy())
159     ] for path, variable in merged_leaf_components.items()]
160     return tabulate(column_values, headers=column_names, tablefmt=tablefmt)
161 
162 
163 def leaf_components(input: tf.Module):
164     return _get_leaf_components(input)
165 
166 
167 def _merge_leaf_components(
168         input: Dict[str, Union[tf.Variable, tf.Tensor, Parameter]]
169 ) -> Dict[str, Union[tf.Variable, tf.Tensor, Parameter]]:
170 
171     input_values = set(
172         [value.experimental_ref() for value in input.values()]
173     )
174     if len(input_values) == len(input):
175         return input
176     tmp_dict = dict()  # Type: Dict[ref, str]
177     for key, variable in input.items():
178         ref = variable.experimental_ref()
179         if ref in tmp_dict:
180             tmp_dict[ref] = f"{tmp_dict[ref]}\n{key}"
181         else:
182             tmp_dict[ref] = key
183     return {key: ref.deref() for ref, key in tmp_dict.items()}
184 
185 
186 def _get_leaf_components(input_module: tf.Module):
187     """
188     Returns a list of tuples each corresponding to a gpflow.Parameter or tf.Variable in the each
189     submodules of a given tf.Module. Each tuple consists of an specific Parameter (or Variable) and
190     its relative path inside the module, which is constructed recursively by adding a prefix with
191     the path to the current module. Designed to be used as a helper for the method 'print_summary'.
192 
193     :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.
194     :return:
195     """
196     target_types = (Parameter, tf.Variable)
197     input_name, state = input_module.__class__.__name__, dict()
198     accumulator = (input_name, state)
199 
200     def update_state(parameter_or_variable, path, state):
201         state[path] = parameter_or_variable
202         return state
203 
204     state = traverse_module(input_module, accumulator, update_state, target_types)
205     return state
206 
207 
208 def reset_cache_bijectors(input_module: tf.Module) -> tf.Module:
209     """
210     Recursively finds tfp.bijectors.Bijector-s inside the components of the tf.Module using `traverse_component`.
211     Resets the caches stored inside each tfp.bijectors.Bijector.
212 
213     :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.
214     :return:
215     """
216     target_types = (tfp.bijectors.Bijector, )
217     accumulator = ('', None)
218 
219     def clear_bijector(bijector, _, state):
220         bijector._from_x.clear()
221         bijector._from_y.clear()
222         return state
223 
224     _ = traverse_module(input_module, accumulator, clear_bijector, target_types)
225     return input_module
226 
227 
228 def deepcopy_components(input_module: tf.Module) -> tf.Module:
229     """
230     Returns a deepcopy of the input tf.Module. To do that first resets the caches stored inside each
231     tfp.bijectors.Bijector to allow the deepcopy of the tf.Module.
232 
233     :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.
234     :return:
235     """
236     return deepcopy(reset_cache_bijectors(input_module))
237 
238 
239 def traverse_module(m: TraverseInput, acc: Accumulator, update_cb: TraverseUpdateCallable,
240                     target_types: tuple) -> Accumulator:
241     """
242     Recursively traverses `m`, accumulating in `acc` a path and a state until it finds an object of type
243     in `target_types` to apply `update_cb` to update the accumulator `acc` and/or the object.
244 
245     :param m: tf.Module, tf.Variable or gpflow.Parameter
246     :param acc: Tuple of path and state
247     :param update_cb: Callable
248     :param target_types: target class types
249     :return:
250     """
251     path, state = acc
252 
253     new_state = state
254 
255     if isinstance(m, target_types):
256         return update_cb(m, path, state)
257 
258     if isinstance(m, (list, tuple)):
259         for term_idx, subterm in enumerate(m):
260             new_acc = (f"{path}[{term_idx}]", new_state)
261             new_state = traverse_module(subterm, new_acc, update_cb, target_types)
262     elif isinstance(m, dict):
263         for term_idx, subterm in m.items():
264             new_acc = (f"{path}['{term_idx}']", new_state)
265             new_state = traverse_module(subterm, new_acc, update_cb, target_types)
266     elif isinstance(m, tf.Module):
267         for name, submodule in vars(m).items():
268             if name in tf.Module._TF_MODULE_IGNORED_PROPERTIES:
269                 continue
270             new_acc = (f"{path}.{name}", new_state)
271             new_state = traverse_module(submodule, new_acc, update_cb, target_types)
272     return new_state
273 
274 
275 @lru_cache()
276 def _first_three_elements_regexp():
277     num_re = r"[+\-]?(?:0|[1-9]\d*)(?:\.\d*)?(?:[eE][+\-]?\d+)?"
278     pat_re = rf"^(?:(\[+)\s*)?({num_re})(?:\s+({num_re})(?:\s+({num_re}))?)?.*?"
279     return re.compile(pat_re)
280 
281 
282 def _str_tensor_value(value: np.ndarray):
283     value_str = str(value)
284     if value.size <= 3:
285         return value_str
286 
287     max_chars = 500
288     value_str = value_str[:max_chars]
289     regexp = _first_three_elements_regexp()
290     match = regexp.match(value_str)
291     assert match is not None
292     brackets, elem1, elem2, elem3 = match.groups()
293 
294     out = f"{elem1}"
295     if elem2 is not None:
296         out = f"{out}{f', {elem2}'}"
297         if elem3 is not None:
298             out = f"{out}{f', {elem3}'}"
299     if brackets is not None:
300         out = f"{brackets}{out}..."
301 
302     return out
303
```
--- 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/utilities/utilities.py b/gpflow/utilities/utilities.py
--- a/gpflow/utilities/utilities.py
+++ b/gpflow/utilities/utilities.py
@@ -137,26 +137,37 @@
 
 
 def tabulate_module_summary(module: tf.Module, tablefmt: Optional[str] = None) -> str:
-    column_names = ['name', 'class', 'transform', 'trainable', 'shape', 'dtype', 'value']
-
-    def get_name(v):
-        return v.__class__.__name__
+    def get_transform(path, var):
+        if hasattr(var, 'transform') and var.transform is not None:
+            if isinstance(var.transform, tfp.bijectors.Chain):
+                return " + ".join(b.__class__.__name__ for b in var.transform.bijectors[::-1])
+            return var.transform.__class__.__name__
+        return None
 
-    def get_transform(v):
-        if hasattr(v, "transform") and v.transform is not None:
-            if isinstance(v.transform, tfp.bijectors.Chain):
-                return " + ".join(b.__class__.__name__ for b in v.transform.bijectors[::-1])
-            return v.transform.__class__.__name__
+    def get_prior(path, var):
+        if hasattr(var, 'prior') and var.prior is not None:
+            return var.prior.name
         return None
 
+    # list of (column_name: str, column_getter: Callable[[tf.Variable], str]) tuples:
+    column_definition = [
+        ("name", lambda path, var: path),
+        ("class", lambda path, var: var.__class__.__name__),
+        ("transform", get_transform),
+        ("prior", get_prior),
+        ("trainable", lambda path, var: var.trainable),
+        ("shape", lambda path, var: var.shape),
+        ("dtype", lambda path, var: var.dtype.name),
+        ("value", lambda path, var: _str_tensor_value(var.numpy())),
+        ]
+    column_names, column_getters = zip(*column_definition)
+
     merged_leaf_components = _merge_leaf_components(leaf_components(module))
 
-    column_values = [[
-        path,
-        get_name(variable),
-        get_transform(variable), variable.trainable, variable.shape, variable.dtype.name,
-        _str_tensor_value(variable.numpy())
-    ] for path, variable in merged_leaf_components.items()]
+    column_values = [
+        [getter(path, variable) for getter in column_getters]
+        for path, variable in merged_leaf_components.items()
+    ]
     return tabulate(column_values, headers=column_names, tablefmt=tablefmt)

{"golden_diff": "diff --git a/gpflow/utilities/utilities.py b/gpflow/utilities/utilities.py\n--- a/gpflow/utilities/utilities.py\n+++ b/gpflow/utilities/utilities.py\n@@ -137,26 +137,37 @@\n \n \n def tabulate_module_summary(module: tf.Module, tablefmt: Optional[str] = None) -> str:\n-    column_names = ['name', 'class', 'transform', 'trainable', 'shape', 'dtype', 'value']\n-\n-    def get_name(v):\n-        return v.__class__.__name__\n+    def get_transform(path, var):\n+        if hasattr(var, 'transform') and var.transform is not None:\n+            if isinstance(var.transform, tfp.bijectors.Chain):\n+                return \" + \".join(b.__class__.__name__ for b in var.transform.bijectors[::-1])\n+            return var.transform.__class__.__name__\n+        return None\n \n-    def get_transform(v):\n-        if hasattr(v, \"transform\") and v.transform is not None:\n-            if isinstance(v.transform, tfp.bijectors.Chain):\n-                return \" + \".join(b.__class__.__name__ for b in v.transform.bijectors[::-1])\n-            return v.transform.__class__.__name__\n+    def get_prior(path, var):\n+        if hasattr(var, 'prior') and var.prior is not None:\n+            return var.prior.name\n         return None\n \n+    # list of (column_name: str, column_getter: Callable[[tf.Variable], str]) tuples:\n+    column_definition = [\n+        (\"name\", lambda path, var: path),\n+        (\"class\", lambda path, var: var.__class__.__name__),\n+        (\"transform\", get_transform),\n+        (\"prior\", get_prior),\n+        (\"trainable\", lambda path, var: var.trainable),\n+        (\"shape\", lambda path, var: var.shape),\n+        (\"dtype\", lambda path, var: var.dtype.name),\n+        (\"value\", lambda path, var: _str_tensor_value(var.numpy())),\n+        ]\n+    column_names, column_getters = zip(*column_definition)\n+\n     merged_leaf_components = _merge_leaf_components(leaf_components(module))\n \n-    column_values = [[\n-        path,\n-        get_name(variable),\n-        get_transform(variable), variable.trainable, variable.shape, variable.dtype.name,\n-        _str_tensor_value(variable.numpy())\n-    ] for path, variable in merged_leaf_components.items()]\n+    column_values = [\n+        [getter(path, variable) for getter in column_getters]\n+        for path, variable in merged_leaf_components.items()\n+    ]\n     return tabulate(column_values, headers=column_names, tablefmt=tablefmt)\n", "issue": "print_summary does not display information on parameter priors\n### `print_summary` does not display information on parameter priors\r\n\r\nGPflow 2 (`develop-2.0` branch):\r\n`print_summary` on a model with priors does not have a column for the `prior` attribute.\r\n\r\nMWE:\r\n```python\r\nimport gpflow\r\nimport tensorflow_probability as tfp\r\nfrom gpflow.utilities import to_default_float\r\nfrom gpflow.utilities import print_summary\r\n\r\nk = gpflow.kernels.Matern32()\r\nk.variance.prior = tfp.distributions.Gamma(\r\n    to_default_float(2), to_default_float(3))\r\n\r\nprint_summary(k)\r\n```\nprint_summary does not display information on parameter priors\n### `print_summary` does not display information on parameter priors\r\n\r\nGPflow 2 (`develop-2.0` branch):\r\n`print_summary` on a model with priors does not have a column for the `prior` attribute.\r\n\r\nMWE:\r\n```python\r\nimport gpflow\r\nimport tensorflow_probability as tfp\r\nfrom gpflow.utilities import to_default_float\r\nfrom gpflow.utilities import print_summary\r\n\r\nk = gpflow.kernels.Matern32()\r\nk.variance.prior = tfp.distributions.Gamma(\r\n    to_default_float(2), to_default_float(3))\r\n\r\nprint_summary(k)\r\n```\n", "before_files": [{"content": "import re\nfrom copy import deepcopy\nfrom functools import lru_cache\nfrom typing import Callable, Dict, List, Optional, Union, TypeVar, Any, Tuple\n\nimport numpy as np\nimport tensorflow as tf\nimport tensorflow_probability as tfp\nfrom tabulate import tabulate\n\nfrom ..base import Parameter\nfrom ..config import default_summary_fmt, default_float, default_int\n\n__all__ = [\n    \"set_trainable\",\n    \"multiple_assign\",\n    \"training_loop\",\n    \"print_summary\",\n    \"deepcopy_components\",\n    \"leaf_components\",\n    \"parameter_dict\",\n    \"read_values\",\n    \"to_default_float\",\n    \"to_default_int\",\n]\n\nTraverseInput = TypeVar(\"TraverseInput\", tf.Variable, tf.Module, Parameter)\nState = Any\nPath = str\nAccumulator = Tuple[Path, State]\nTraverseUpdateCallable = Callable[[TraverseInput, Path, State], State]\n\n\ndef to_default_int(x):\n    return tf.cast(x, dtype=default_int())\n\n\ndef to_default_float(x):\n    return tf.cast(x, dtype=default_float())\n\n\ndef set_trainable(model: tf.Module, flag: bool):\n    \"\"\"\n    Set trainable flag for all `tf.Variable`s and `gpflow.Parameter`s in a module.\n    \"\"\"\n    for variable in model.variables:\n        variable._trainable = flag\n\n\ndef multiple_assign(module: tf.Module, parameters: Dict[str, tf.Tensor]):\n    \"\"\"\n    Multiple assign takes a dictionary with new values. Dictionary keys are paths to the\n    `tf.Variable`s or `gpflow.Parameter` of the input module.\n\n    :param module: `tf.Module`.\n    :param parameters: a dictionary with keys of the form \".module.path.to.variable\" and new value tensors.\n    \"\"\"\n    reference_var_dict = parameter_dict(module)\n    for path, value in parameters.items():\n        reference_var_dict[path].assign(value)\n\n\ndef read_values(module: tf.Module) -> Dict[str, np.ndarray]:\n    \"\"\"Returns a dictionary of numpy values of the module parameters (variables).\"\"\"\n    return {k: v.numpy() for k, v in parameter_dict(module).items()}\n\n\ndef parameter_dict(module: tf.Module) -> Dict[str, Union[Parameter, tf.Variable]]:\n    \"\"\"\n    Returns a dictionary of parameters (variables) for the `tf.Module` component.\n    Dictionary keys are relative paths to the attributes to which parameters (variables) assigned to.\n\n        class SubModule(tf.Module):\n            def __init__(self):\n                self.parameter = gpflow.Parameter(1.0)\n                self.variable = tf.Variable(1.0)\n\n        class Module(tf.Module):\n            def __init__(self):\n                self.submodule = SubModule()\n\n        m = Module()\n        params = parameter_dict(m)\n        # {\n        #   \".submodule.parameter\": <parameter object>,\n        #   \".submodule.variable\": <variable object>\n        # }\n    \"\"\"\n    param_dict = leaf_components(module)\n    return {f\".{key.split('.', 1)[-1]}\": value for key, value in param_dict.items()}\n\n\ndef training_loop(closure: Callable[..., tf.Tensor],\n                  optimizer: Optional[tf.optimizers.Optimizer] = None,\n                  var_list: List[tf.Variable] = None,\n                  maxiter=1e3,\n                  jit=False):\n    \"\"\"\n    Simple generic training loop. At each iteration uses a GradientTape to compute\n    the gradients of a loss function with respect to a set of variables.\n\n    :param closure: Callable that constructs a loss function based on data and model being trained\n    :param optimizer: tf.optimizers or tf.keras.optimizers that updates variables by applying the\n        corresponding loss gradients. Adam is a default optimizer with default settings.\n    :param var_list: List of model variables to be learnt during training\n    :param maxiter: Maximum number of\n    :return:\n    \"\"\"\n\n    optimizer = tf.optimizers.Adam() if optimizer is None else optimizer\n\n    def optimization_step():\n        with tf.GradientTape() as tape:\n            tape.watch(var_list)\n            loss = closure()\n        grads = tape.gradient(loss, var_list)\n        optimizer.apply_gradients(zip(grads, var_list))\n\n    if jit:\n        optimization_step = tf.function(optimization_step)\n\n    for _ in range(int(maxiter)):\n        optimization_step()\n\n\ndef print_summary(module: tf.Module, fmt: str = None):\n    \"\"\"\n    Prints a summary of the parameters and variables contained in a tf.Module.\n    \"\"\"\n    fmt = fmt if fmt is not None else default_summary_fmt()\n    if fmt == \"notebook\":\n        from IPython.core.display import display, HTML\n        tab = tabulate_module_summary(module, \"html\")\n        display(HTML(tab))\n    else:\n        print(tabulate_module_summary(module, fmt))\n\n\ndef tabulate_module_summary(module: tf.Module, tablefmt: Optional[str] = None) -> str:\n    column_names = ['name', 'class', 'transform', 'trainable', 'shape', 'dtype', 'value']\n\n    def get_name(v):\n        return v.__class__.__name__\n\n    def get_transform(v):\n        if hasattr(v, \"transform\") and v.transform is not None:\n            if isinstance(v.transform, tfp.bijectors.Chain):\n                return \" + \".join(b.__class__.__name__ for b in v.transform.bijectors[::-1])\n            return v.transform.__class__.__name__\n        return None\n\n    merged_leaf_components = _merge_leaf_components(leaf_components(module))\n\n    column_values = [[\n        path,\n        get_name(variable),\n        get_transform(variable), variable.trainable, variable.shape, variable.dtype.name,\n        _str_tensor_value(variable.numpy())\n    ] for path, variable in merged_leaf_components.items()]\n    return tabulate(column_values, headers=column_names, tablefmt=tablefmt)\n\n\ndef leaf_components(input: tf.Module):\n    return _get_leaf_components(input)\n\n\ndef _merge_leaf_components(\n        input: Dict[str, Union[tf.Variable, tf.Tensor, Parameter]]\n) -> Dict[str, Union[tf.Variable, tf.Tensor, Parameter]]:\n\n    input_values = set(\n        [value.experimental_ref() for value in input.values()]\n    )\n    if len(input_values) == len(input):\n        return input\n    tmp_dict = dict()  # Type: Dict[ref, str]\n    for key, variable in input.items():\n        ref = variable.experimental_ref()\n        if ref in tmp_dict:\n            tmp_dict[ref] = f\"{tmp_dict[ref]}\\n{key}\"\n        else:\n            tmp_dict[ref] = key\n    return {key: ref.deref() for ref, key in tmp_dict.items()}\n\n\ndef _get_leaf_components(input_module: tf.Module):\n    \"\"\"\n    Returns a list of tuples each corresponding to a gpflow.Parameter or tf.Variable in the each\n    submodules of a given tf.Module. Each tuple consists of an specific Parameter (or Variable) and\n    its relative path inside the module, which is constructed recursively by adding a prefix with\n    the path to the current module. Designed to be used as a helper for the method 'print_summary'.\n\n    :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.\n    :return:\n    \"\"\"\n    target_types = (Parameter, tf.Variable)\n    input_name, state = input_module.__class__.__name__, dict()\n    accumulator = (input_name, state)\n\n    def update_state(parameter_or_variable, path, state):\n        state[path] = parameter_or_variable\n        return state\n\n    state = traverse_module(input_module, accumulator, update_state, target_types)\n    return state\n\n\ndef reset_cache_bijectors(input_module: tf.Module) -> tf.Module:\n    \"\"\"\n    Recursively finds tfp.bijectors.Bijector-s inside the components of the tf.Module using `traverse_component`.\n    Resets the caches stored inside each tfp.bijectors.Bijector.\n\n    :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.\n    :return:\n    \"\"\"\n    target_types = (tfp.bijectors.Bijector, )\n    accumulator = ('', None)\n\n    def clear_bijector(bijector, _, state):\n        bijector._from_x.clear()\n        bijector._from_y.clear()\n        return state\n\n    _ = traverse_module(input_module, accumulator, clear_bijector, target_types)\n    return input_module\n\n\ndef deepcopy_components(input_module: tf.Module) -> tf.Module:\n    \"\"\"\n    Returns a deepcopy of the input tf.Module. To do that first resets the caches stored inside each\n    tfp.bijectors.Bijector to allow the deepcopy of the tf.Module.\n\n    :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.\n    :return:\n    \"\"\"\n    return deepcopy(reset_cache_bijectors(input_module))\n\n\ndef traverse_module(m: TraverseInput, acc: Accumulator, update_cb: TraverseUpdateCallable,\n                    target_types: tuple) -> Accumulator:\n    \"\"\"\n    Recursively traverses `m`, accumulating in `acc` a path and a state until it finds an object of type\n    in `target_types` to apply `update_cb` to update the accumulator `acc` and/or the object.\n\n    :param m: tf.Module, tf.Variable or gpflow.Parameter\n    :param acc: Tuple of path and state\n    :param update_cb: Callable\n    :param target_types: target class types\n    :return:\n    \"\"\"\n    path, state = acc\n\n    new_state = state\n\n    if isinstance(m, target_types):\n        return update_cb(m, path, state)\n\n    if isinstance(m, (list, tuple)):\n        for term_idx, subterm in enumerate(m):\n            new_acc = (f\"{path}[{term_idx}]\", new_state)\n            new_state = traverse_module(subterm, new_acc, update_cb, target_types)\n    elif isinstance(m, dict):\n        for term_idx, subterm in m.items():\n            new_acc = (f\"{path}['{term_idx}']\", new_state)\n            new_state = traverse_module(subterm, new_acc, update_cb, target_types)\n    elif isinstance(m, tf.Module):\n        for name, submodule in vars(m).items():\n            if name in tf.Module._TF_MODULE_IGNORED_PROPERTIES:\n                continue\n            new_acc = (f\"{path}.{name}\", new_state)\n            new_state = traverse_module(submodule, new_acc, update_cb, target_types)\n    return new_state\n\n\n@lru_cache()\ndef _first_three_elements_regexp():\n    num_re = r\"[+\\-]?(?:0|[1-9]\\d*)(?:\\.\\d*)?(?:[eE][+\\-]?\\d+)?\"\n    pat_re = rf\"^(?:(\\[+)\\s*)?({num_re})(?:\\s+({num_re})(?:\\s+({num_re}))?)?.*?\"\n    return re.compile(pat_re)\n\n\ndef _str_tensor_value(value: np.ndarray):\n    value_str = str(value)\n    if value.size <= 3:\n        return value_str\n\n    max_chars = 500\n    value_str = value_str[:max_chars]\n    regexp = _first_three_elements_regexp()\n    match = regexp.match(value_str)\n    assert match is not None\n    brackets, elem1, elem2, elem3 = match.groups()\n\n    out = f\"{elem1}\"\n    if elem2 is not None:\n        out = f\"{out}{f', {elem2}'}\"\n        if elem3 is not None:\n            out = f\"{out}{f', {elem3}'}\"\n    if brackets is not None:\n        out = f\"{brackets}{out}...\"\n\n    return out\n", "path": "gpflow/utilities/utilities.py"}], "after_files": [{"content": "import re\nfrom copy import deepcopy\nfrom functools import lru_cache\nfrom typing import Callable, Dict, List, Optional, Union, TypeVar, Any, Tuple\n\nimport numpy as np\nimport tensorflow as tf\nimport tensorflow_probability as tfp\nfrom tabulate import tabulate\n\nfrom ..base import Parameter\nfrom ..config import default_summary_fmt, default_float, default_int\n\n__all__ = [\n    \"set_trainable\",\n    \"multiple_assign\",\n    \"training_loop\",\n    \"print_summary\",\n    \"deepcopy_components\",\n    \"leaf_components\",\n    \"parameter_dict\",\n    \"read_values\",\n    \"to_default_float\",\n    \"to_default_int\",\n]\n\nTraverseInput = TypeVar(\"TraverseInput\", tf.Variable, tf.Module, Parameter)\nState = Any\nPath = str\nAccumulator = Tuple[Path, State]\nTraverseUpdateCallable = Callable[[TraverseInput, Path, State], State]\n\n\ndef to_default_int(x):\n    return tf.cast(x, dtype=default_int())\n\n\ndef to_default_float(x):\n    return tf.cast(x, dtype=default_float())\n\n\ndef set_trainable(model: tf.Module, flag: bool):\n    \"\"\"\n    Set trainable flag for all `tf.Variable`s and `gpflow.Parameter`s in a module.\n    \"\"\"\n    for variable in model.variables:\n        variable._trainable = flag\n\n\ndef multiple_assign(module: tf.Module, parameters: Dict[str, tf.Tensor]):\n    \"\"\"\n    Multiple assign takes a dictionary with new values. Dictionary keys are paths to the\n    `tf.Variable`s or `gpflow.Parameter` of the input module.\n\n    :param module: `tf.Module`.\n    :param parameters: a dictionary with keys of the form \".module.path.to.variable\" and new value tensors.\n    \"\"\"\n    reference_var_dict = parameter_dict(module)\n    for path, value in parameters.items():\n        reference_var_dict[path].assign(value)\n\n\ndef read_values(module: tf.Module) -> Dict[str, np.ndarray]:\n    \"\"\"Returns a dictionary of numpy values of the module parameters (variables).\"\"\"\n    return {k: v.numpy() for k, v in parameter_dict(module).items()}\n\n\ndef parameter_dict(module: tf.Module) -> Dict[str, Union[Parameter, tf.Variable]]:\n    \"\"\"\n    Returns a dictionary of parameters (variables) for the `tf.Module` component.\n    Dictionary keys are relative paths to the attributes to which parameters (variables) assigned to.\n\n        class SubModule(tf.Module):\n            def __init__(self):\n                self.parameter = gpflow.Parameter(1.0)\n                self.variable = tf.Variable(1.0)\n\n        class Module(tf.Module):\n            def __init__(self):\n                self.submodule = SubModule()\n\n        m = Module()\n        params = parameter_dict(m)\n        # {\n        #   \".submodule.parameter\": <parameter object>,\n        #   \".submodule.variable\": <variable object>\n        # }\n    \"\"\"\n    param_dict = leaf_components(module)\n    return {f\".{key.split('.', 1)[-1]}\": value for key, value in param_dict.items()}\n\n\ndef training_loop(closure: Callable[..., tf.Tensor],\n                  optimizer: Optional[tf.optimizers.Optimizer] = None,\n                  var_list: List[tf.Variable] = None,\n                  maxiter=1e3,\n                  jit=False):\n    \"\"\"\n    Simple generic training loop. At each iteration uses a GradientTape to compute\n    the gradients of a loss function with respect to a set of variables.\n\n    :param closure: Callable that constructs a loss function based on data and model being trained\n    :param optimizer: tf.optimizers or tf.keras.optimizers that updates variables by applying the\n        corresponding loss gradients. Adam is a default optimizer with default settings.\n    :param var_list: List of model variables to be learnt during training\n    :param maxiter: Maximum number of\n    :return:\n    \"\"\"\n\n    optimizer = tf.optimizers.Adam() if optimizer is None else optimizer\n\n    def optimization_step():\n        with tf.GradientTape() as tape:\n            tape.watch(var_list)\n            loss = closure()\n        grads = tape.gradient(loss, var_list)\n        optimizer.apply_gradients(zip(grads, var_list))\n\n    if jit:\n        optimization_step = tf.function(optimization_step)\n\n    for _ in range(int(maxiter)):\n        optimization_step()\n\n\ndef print_summary(module: tf.Module, fmt: str = None):\n    \"\"\"\n    Prints a summary of the parameters and variables contained in a tf.Module.\n    \"\"\"\n    fmt = fmt if fmt is not None else default_summary_fmt()\n    if fmt == \"notebook\":\n        from IPython.core.display import display, HTML\n        tab = tabulate_module_summary(module, \"html\")\n        display(HTML(tab))\n    else:\n        print(tabulate_module_summary(module, fmt))\n\n\ndef tabulate_module_summary(module: tf.Module, tablefmt: Optional[str] = None) -> str:\n    def get_transform(path, var):\n        if hasattr(var, 'transform') and var.transform is not None:\n            if isinstance(var.transform, tfp.bijectors.Chain):\n                return \" + \".join(b.__class__.__name__ for b in var.transform.bijectors[::-1])\n            return var.transform.__class__.__name__\n        return None\n\n    def get_prior(path, var):\n        if hasattr(var, 'prior') and var.prior is not None:\n            return var.prior.name\n        return None\n\n    # list of (column_name: str, column_getter: Callable[[tf.Variable], str]) tuples:\n    column_definition = [\n        (\"name\", lambda path, var: path),\n        (\"class\", lambda path, var: var.__class__.__name__),\n        (\"transform\", get_transform),\n        (\"prior\", get_prior),\n        (\"trainable\", lambda path, var: var.trainable),\n        (\"shape\", lambda path, var: var.shape),\n        (\"dtype\", lambda path, var: var.dtype.name),\n        (\"value\", lambda path, var: _str_tensor_value(var.numpy())),\n        ]\n    column_names, column_getters = zip(*column_definition)\n\n    merged_leaf_components = _merge_leaf_components(leaf_components(module))\n\n    column_values = [\n        [getter(path, variable) for getter in column_getters]\n        for path, variable in merged_leaf_components.items()\n    ]\n    return tabulate(column_values, headers=column_names, tablefmt=tablefmt)\n\n\ndef leaf_components(input: tf.Module):\n    return _get_leaf_components(input)\n\n\ndef _merge_leaf_components(\n        input: Dict[str, Union[tf.Variable, tf.Tensor, Parameter]]\n) -> Dict[str, Union[tf.Variable, tf.Tensor, Parameter]]:\n\n    input_values = set(\n        [value.experimental_ref() for value in input.values()]\n    )\n    if len(input_values) == len(input):\n        return input\n    tmp_dict = dict()  # Type: Dict[ref, str]\n    for key, variable in input.items():\n        ref = variable.experimental_ref()\n        if ref in tmp_dict:\n            tmp_dict[ref] = f\"{tmp_dict[ref]}\\n{key}\"\n        else:\n            tmp_dict[ref] = key\n    return {key: ref.deref() for ref, key in tmp_dict.items()}\n\n\ndef _get_leaf_components(input_module: tf.Module):\n    \"\"\"\n    Returns a list of tuples each corresponding to a gpflow.Parameter or tf.Variable in the each\n    submodules of a given tf.Module. Each tuple consists of an specific Parameter (or Variable) and\n    its relative path inside the module, which is constructed recursively by adding a prefix with\n    the path to the current module. Designed to be used as a helper for the method 'print_summary'.\n\n    :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.\n    :return:\n    \"\"\"\n    target_types = (Parameter, tf.Variable)\n    input_name, state = input_module.__class__.__name__, dict()\n    accumulator = (input_name, state)\n\n    def update_state(parameter_or_variable, path, state):\n        state[path] = parameter_or_variable\n        return state\n\n    state = traverse_module(input_module, accumulator, update_state, target_types)\n    return state\n\n\ndef reset_cache_bijectors(input_module: tf.Module) -> tf.Module:\n    \"\"\"\n    Recursively finds tfp.bijectors.Bijector-s inside the components of the tf.Module using `traverse_component`.\n    Resets the caches stored inside each tfp.bijectors.Bijector.\n\n    :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.\n    :return:\n    \"\"\"\n    target_types = (tfp.bijectors.Bijector, )\n    accumulator = ('', None)\n\n    def clear_bijector(bijector, _, state):\n        bijector._from_x.clear()\n        bijector._from_y.clear()\n        return state\n\n    _ = traverse_module(input_module, accumulator, clear_bijector, target_types)\n    return input_module\n\n\ndef deepcopy_components(input_module: tf.Module) -> tf.Module:\n    \"\"\"\n    Returns a deepcopy of the input tf.Module. To do that first resets the caches stored inside each\n    tfp.bijectors.Bijector to allow the deepcopy of the tf.Module.\n\n    :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.\n    :return:\n    \"\"\"\n    return deepcopy(reset_cache_bijectors(input_module))\n\n\ndef traverse_module(m: TraverseInput, acc: Accumulator, update_cb: TraverseUpdateCallable,\n                    target_types: tuple) -> Accumulator:\n    \"\"\"\n    Recursively traverses `m`, accumulating in `acc` a path and a state until it finds an object of type\n    in `target_types` to apply `update_cb` to update the accumulator `acc` and/or the object.\n\n    :param m: tf.Module, tf.Variable or gpflow.Parameter\n    :param acc: Tuple of path and state\n    :param update_cb: Callable\n    :param target_types: target class types\n    :return:\n    \"\"\"\n    path, state = acc\n\n    new_state = state\n\n    if isinstance(m, target_types):\n        return update_cb(m, path, state)\n\n    if isinstance(m, (list, tuple)):\n        for term_idx, subterm in enumerate(m):\n            new_acc = (f\"{path}[{term_idx}]\", new_state)\n            new_state = traverse_module(subterm, new_acc, update_cb, target_types)\n    elif isinstance(m, dict):\n        for term_idx, subterm in m.items():\n            new_acc = (f\"{path}['{term_idx}']\", new_state)\n            new_state = traverse_module(subterm, new_acc, update_cb, target_types)\n    elif isinstance(m, tf.Module):\n        for name, submodule in vars(m).items():\n            if name in tf.Module._TF_MODULE_IGNORED_PROPERTIES:\n                continue\n            new_acc = (f\"{path}.{name}\", new_state)\n            new_state = traverse_module(submodule, new_acc, update_cb, target_types)\n    return new_state\n\n\n@lru_cache()\ndef _first_three_elements_regexp():\n    num_re = r\"[+\\-]?(?:0|[1-9]\\d*)(?:\\.\\d*)?(?:[eE][+\\-]?\\d+)?\"\n    pat_re = rf\"^(?:(\\[+)\\s*)?({num_re})(?:\\s+({num_re})(?:\\s+({num_re}))?)?.*?\"\n    return re.compile(pat_re)\n\n\ndef _str_tensor_value(value: np.ndarray):\n    value_str = str(value)\n    if value.size <= 3:\n        return value_str\n\n    max_chars = 500\n    value_str = value_str[:max_chars]\n    regexp = _first_three_elements_regexp()\n    match = regexp.match(value_str)\n    assert match is not None\n    brackets, elem1, elem2, elem3 = match.groups()\n\n    out = f\"{elem1}\"\n    if elem2 is not None:\n        out = f\"{out}{f', {elem2}'}\"\n        if elem3 is not None:\n            out = f\"{out}{f', {elem3}'}\"\n    if brackets is not None:\n        out = f\"{brackets}{out}...\"\n\n    return out\n", "path": "gpflow/utilities/utilities.py"}]}