TensorFlow Skipping loop optimization warning

[Thomas-Christie]

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, the track_state`` argument of the :meth:~trieste.bayesian_optimizer.BayesianOptimizer.optimize method should be set to `False`. 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

[uri-granta]

Thanks for reporting this and sorry that it somehow fell between the gaps and was ignored for so long!

I've not managed to reproduce the warning with your sample script, which ran to completion (I tensorflow/core/kernels/logging_ops.cc:171] Optimization completed without errors) without that warning, suggesting it's either been fixed or is perhaps connected to running on macOS. There is definitely nothing that looks wrong with the script, and TF is often trigger happy with its warnings, so if the results are as expected I would recommend ignoring it.