Why is GPflow's Scipy optimizer incompatible with decorating the optimization step with tf.function?

Question

I am supplying different minibatches to optimize a GPflow model (SVGP). If I decorate the optimization_step with tf.function I get the following error:

NotImplementedError: Cannot convert a symbolic Tensor (concat:0) to a numpy array. This error may indicate that you're trying to pass a Tensor to a NumPy call, which is not supported

In order for the optimizer to run I had to remove the tf.function decorator, losing the speed-up advantages. What do I need to change so that I can keep using the tf.function decorator?

The xAndY input shapes and types are all numpy arrays.

type(xAndY)
Out[71]: tuple
xAndY[0].shape
Out[72]: (245760, 2)
xAndY[1].shape
Out[73]: (245760, 1)
type(xAndY[0])
Out[74]: numpy.ndarray

def run_optimizer_on_minibatch_size(model, iterations, minibatch_size, xAndY):
    """
    Utility function running a Scipy optimizer
    :param model: GPflow model
    :param interations: number of iterations
    """
    N = xAndY[0].shape[0]
    tensor_data = tuple(map(tf.convert_to_tensor, xAndY))
    train_dataset = tf.data.Dataset.from_tensor_slices(tensor_data).repeat().shuffle(N)
    
    logf = []
    train_iter = iter(train_dataset.batch(minibatch_size))
    training_loss = model.training_loss_closure(train_iter, compile=True)
    optimizer = gpflow.optimizers.Scipy()

    @tf.function   # had to remove this decorator
    def optimization_step():
        optimizer.minimize(training_loss, model.trainable_variables)

    # step = 0
    for step in range(iterations):
        optimization_step()
        if step % 10 == 0:
            elbo = -training_loss().numpy()
            logf.append(elbo)
            print(elbo)
    return logf

from gpflow.ci_utils import ci_niter
maxiter = ci_niter(20000)
logf = run_optimizer_on_minibatch_size(m, maxiter, minibatch_size, (X,Y))

Answer 1

GPflow's gpflow.optimizers.Scipy() is a wrapper around Scipy's minimize(), and as it calls into non-TensorFlow operations, you cannot wrap it in tf.function. Moreover, the optimizers implemented in Scipy's minimize are second-order methods that assume that your gradients are not stochastic, and aren't compatible with minibatching.

If you want to do full-batch optimization with Scipy: The minimize() method of gpflow.optimizers.Scipy(), by default, does wrap the objective and gradient computation inside tf.function (see its compile argument with default True). It also does the full optimization, so you only have to call the minimize() method once (by default it runs until convergence or failure to continue optimization; you can supply a maximum number of iterations using the options=dict(maxiter=1000) argument).

If you want to use mini-batching: simply use one of the TensorFlow optimizers, such as tf.optimizers.Adam(), and then your code should run fine including the @tf.function decorator on your optimization_step() function (and in that case you do need to call it in a loop as in your example).