GPflow wrapper update may be incompatible with tf.function()

As mentioned on https://github.com/secondmind-labs/trieste/pull/54, there might be an issue with the current way of how Trieste updates data in GPflow models when you want to use tf.function to speed up the prediction step (currently, Trieste only uses tf.function for model optimization, and relies on eager mode for predictions).

VariationalGaussianProcess.update() sets the model’s data attribute to the new (X, Y) tuple. However, if model.predict_f() is wrapped in tf.function(), it will store the values of (X, Y) when first called in the cached graph. To enable data update with tf.function(), model.data must be set to a tuple of tf.Variables, which are then .assign()ed rather than overwritten in VariationalGaussianProcess.update().

[johnamcleod]

model.num_data in a variable causes the elbo method to fail

Unfortunately the tf.eye function only supports Tensors, Python Integral and np.integer types for the first input. If none of those work with tf.function you could consider creating a subclass of the VGP model in Trieste to turn the num_data attribute into a @property. Then it could be dynamically calculated from the size of the Xdata variable. If that works we could incorporate it into GPflow.

[uri-granta]

I’m currently looking at this (see #271) and am hitting an issue with updating the VGP wrapper. Putting model.data in two variables works fine, but putting model.num_data in a variable causes the elbo method to fail when it tries to execute tf.eye(self.num_data, dtype=default_float()), claiming the self.num_data isn’t an integer. Any idea if there’s a way round this? @johnamcleod