TensorFlow `Skipping loop optimization` warning

Describe the bug Hi, thanks for the excellent library. I have been using the library for a while now and have found that whenever I'm dealing with fitting GaussianProcessRegression models TensorFlow gives me the following warning:

W tensorflow/core/grappler/optimizers/loop_optimizer.cc:907] Skipping loop optimization for Merge node with control input: cond/branch_executed/_9

This seems to occur when the models are getting fitted by the BayesianOptimizer (https://github.com/secondmind-labs/trieste/blob/develop/trieste/bayesian_optimizer.py#L751C5-L755).

I just wanted to double check that this is unavoidable, and not something to do with my setup?

As a side note - the expected improvement notebook also doesn't seem to work in its current state; running the BO loop gives the following error:

NotImplementedError: Failed to save the optimization state. Some models do not support deecopying or serialization and cannot be saved. (This is particularly common for deep neural network models, though some of the model wrappers accept a model closure as a workaround.) For these models, thetrack_state` argument of the :meth:~trieste.bayesian_optimizer.BayesianOptimizer.optimizemethod should be set toFalse. This means that only the final model will be available.

To reproduce Steps to reproduce the behaviour:

Behaviour can be reproduced by running the expected improvement notebook. Alternatively, you can run the following (which has been copied from the notebook):


import tensorflow as tf
from trieste.objectives import ScaledBranin
from trieste.experimental.plotting import plot_function_plotly
from trieste.space import Box
import trieste
import gpflow
import tensorflow_probability as tfp

from trieste.models.gpflow import GaussianProcessRegression

if __name__ == '__main__':
    np.random.seed(1793)
    tf.random.set_seed(1793)

    scaled_branin = ScaledBranin.objective
    search_space = ScaledBranin.search_space  # predefined search space
    search_space = Box([0, 0], [1, 1])  # define the search space directly

    observer = trieste.objectives.utils.mk_observer(scaled_branin)

    num_initial_points = 5
    initial_query_points = search_space.sample_sobol(num_initial_points)
    initial_data = observer(initial_query_points)

    def build_model(data):
        variance = tf.math.reduce_variance(data.observations)
        kernel = gpflow.kernels.Matern52(variance=variance, lengthscales=[0.2, 0.2])
        prior_scale = tf.cast(1.0, dtype=tf.float64)
        kernel.variance.prior = tfp.distributions.LogNormal(
            tf.cast(-2.0, dtype=tf.float64), prior_scale
        )
        kernel.lengthscales.prior = tfp.distributions.LogNormal(
            tf.math.log(kernel.lengthscales), prior_scale
        )
        gpr = gpflow.models.GPR(data.astuple(), kernel, noise_variance=1e-5)
        gpflow.set_trainable(gpr.likelihood, False)

        return GaussianProcessRegression(gpr, num_kernel_samples=100)

    model = build_model(initial_data)

    bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)

    num_steps = 15
    result = bo.optimize(num_steps, initial_data, model, track_state=False)
    dataset = result.try_get_final_dataset()```

**System information**
 - OS: macOS 13.4.1
 - Python version: 3.10.10
 - Trieste version: Latest commit (https://github.com/secondmind-labs/trieste/commit/56101c092f28ed87398c4cd63fdece2f16909451)
 - TensorFlow version: 2.12.1
 - GPflow version: 2.9.0

secondmind-labs / trieste

TensorFlow `Skipping loop optimization` warning #779