# Buggy code
def predict_f(
    self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
) -> MeanAndVariance:
    """
    Compute the mean and variance of the latent function at some new points.
    Note that this is very similar to the SGPR prediction, for which
    there are notes in the SGPR notebook.

    Note: This model does not allow full output covariances.

    :param Xnew: points at which to predict
    """
    if full_output_cov:
        raise NotImplementedError

    pX = DiagonalGaussian(self.X_data_mean, self.X_data_var)

    Y_data = self.data
    num_inducing = len(self.inducing_variable)
    psi1 = expectation(pX, (self.kernel, self.inducing_variable))
    psi2 = tf.reduce_sum(
        expectation(
            pX,
            (self.kernel, self.inducing_variable),
            (self.kernel, self.inducing_variable),
        ),
        axis=0,
    )
    jitter = default_jitter()
    Kus = covariances.Kuf(self.inducing_variable, self.kernel, Xnew)
    sigma2 = self.likelihood.variance
    sigma = tf.sqrt(sigma2)
    L = tf.linalg.cholesky(
        covariances.Kuu(self.inducing_variable, self.kernel, jitter=jitter)
    )

    A = tf.linalg.triangular_solve(L, tf.transpose(psi1), lower=True) / sigma
    tmp = tf.linalg.triangular_solve(L, psi2, lower=True)
    AAT = tf.linalg.triangular_solve(L, tf.transpose(tmp), lower=True) / sigma2
    B = AAT + tf.eye(num_inducing, dtype=default_float())
    LB = tf.linalg.cholesky(B)
    c = tf.linalg.triangular_solve(LB, tf.linalg.matmul(A, Y_data), lower=True) / sigma
    tmp1 = tf.linalg.triangular_solve(L, Kus, lower=True)
    tmp2 = tf.linalg.triangular_solve(LB, tmp1, lower=True)
    mean = tf.linalg.matmul(tmp2, c, transpose_a=True)
    if full_cov:
        var = (
            self.kernel(Xnew)
            + tf.linalg.matmul(tmp2, tmp2, transpose_a=True)
            - tf.linalg.matmul(tmp1, tmp1, transpose_a=True)
        )
        shape = tf.stack([1, 1, tf.shape(Y_data)[1]])
        var = tf.tile(tf.expand_dims(var, 2), shape)
    else:
        var = (
            self.kernel(Xnew, full_cov=False)
            + tf.reduce_sum(tf.square(tmp2), axis=0)
            - tf.reduce_sum(tf.square(tmp1), axis=0)
        )
        shape = tf.stack([1, tf.shape(Y_data)[1]])
        var = tf.tile(tf.expand_dims(var, 1), shape)
    return mean + self.mean_function(Xnew), var


# Traceback (TypeError)
TypeError                                 Traceback (most recent call last)
<ipython-input-24-9a082736eedc> in <module>
38
39
---> 40     optimization_step()
41
42

~/anaconda3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
778       else:
779         compiler = "nonXla"
--> 780         result = self._call(*args, **kwds)
781
782       new_tracing_count = self._get_tracing_count()

~/anaconda3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
821       # This is the first call of __call__, so we have to initialize.
822       initializers = []
--> 823       self._initialize(args, kwds, add_initializers_to=initializers)
824     finally:
825       # At this point we know that the initialization is complete (or less

~/anaconda3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
694     self._graph_deleter = FunctionDeleter(self._lifted_initializer_graph)
695     self._concrete_stateful_fn = (
--> 696         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
697             *args, **kwds))
698

~/anaconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
2853       args, kwargs = None, None
2854     with self._lock:
-> 2855       graph_function, _, _ = self._maybe_define_function(args, kwargs)
2856     return graph_function
2857

~/anaconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
3211
3212       self._function_cache.missed.add(call_context_key)
-> 3213       graph_function = self._create_graph_function(args, kwargs)
3214       self._function_cache.primary[cache_key] = graph_function
3215       return graph_function, args, kwargs

~/anaconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
3063     arg_names = base_arg_names + missing_arg_names
3064     graph_function = ConcreteFunction(
-> 3065         func_graph_module.func_graph_from_py_func(
3066             self._name,
3067             self._python_function,

~/anaconda3/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
984         _, original_func = tf_decorator.unwrap(python_func)
985
--> 986       func_outputs = python_func(*func_args, **func_kwargs)
987
988       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~/anaconda3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
598         # __wrapped__ allows AutoGraph to swap in a converted function. We give
599         # the function a weak reference to itself to avoid a reference cycle.
--> 600         return weak_wrapped_fn().__wrapped__(*args, **kwds)
601     weak_wrapped_fn = weakref.ref(wrapped_fn)
602

~/anaconda3/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
971           except Exception as e:  # pylint:disable=broad-except
972             if hasattr(e, "ag_error_metadata"):
--> 973               raise e.ag_error_metadata.to_exception(e)
974             else:
975               raise

TypeError: in user code:

<ipython-input-24-9a082736eedc>:32 optimization_step  *
optimizer.minimize(m.training_loss, m.trainable_variables)
/home/maltamirano/anaconda3/lib/python3.8/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:374 minimize  **
grads_and_vars = self._compute_gradients(
/home/maltamirano/anaconda3/lib/python3.8/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:429 _compute_gradients
loss_value = loss()
/home/maltamirano/anaconda3/lib/python3.8/site-packages/gpflow/models/training_mixins.py:64 training_loss
return self._training_loss()
/home/maltamirano/anaconda3/lib/python3.8/site-packages/gpflow/models/model.py:57 _training_loss
return -(self.maximum_log_likelihood_objective(*args, **kwargs) + self.log_prior_density())
/home/maltamirano/anaconda3/lib/python3.8/site-packages/gpflow/models/sgpr.py:154 maximum_log_likelihood_objective
return self.elbo()
/home/maltamirano/anaconda3/lib/python3.8/site-packages/gpflow/models/sgpr.py:164 elbo
num_inducing = len(self.inducing_variable)

TypeError: 'Tensor' object cannot be interpreted as an integer

# Fixed code
def predict_f(
    self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
) -> MeanAndVariance:
    """
    Compute the mean and variance of the latent function at some new points.
    Note that this is very similar to the SGPR prediction, for which
    there are notes in the SGPR notebook.

    Note: This model does not allow full output covariances.

    :param Xnew: points at which to predict
    """
    if full_output_cov:
        raise NotImplementedError

    pX = DiagonalGaussian(self.X_data_mean, self.X_data_var)

    Y_data = self.data
    num_inducing = self.inducing_variable.num_inducing
    psi1 = expectation(pX, (self.kernel, self.inducing_variable))
    psi2 = tf.reduce_sum(
        expectation(
            pX,
            (self.kernel, self.inducing_variable),
            (self.kernel, self.inducing_variable),
        ),
        axis=0,
    )
    jitter = default_jitter()
    Kus = covariances.Kuf(self.inducing_variable, self.kernel, Xnew)
    sigma2 = self.likelihood.variance
    sigma = tf.sqrt(sigma2)
    L = tf.linalg.cholesky(
        covariances.Kuu(self.inducing_variable, self.kernel, jitter=jitter)
    )

    A = tf.linalg.triangular_solve(L, tf.transpose(psi1), lower=True) / sigma
    tmp = tf.linalg.triangular_solve(L, psi2, lower=True)
    AAT = tf.linalg.triangular_solve(L, tf.transpose(tmp), lower=True) / sigma2
    B = AAT + tf.eye(num_inducing, dtype=default_float())
    LB = tf.linalg.cholesky(B)
    c = tf.linalg.triangular_solve(LB, tf.linalg.matmul(A, Y_data), lower=True) / sigma
    tmp1 = tf.linalg.triangular_solve(L, Kus, lower=True)
    tmp2 = tf.linalg.triangular_solve(LB, tmp1, lower=True)
    mean = tf.linalg.matmul(tmp2, c, transpose_a=True)
    if full_cov:
        var = (
            self.kernel(Xnew)
            + tf.linalg.matmul(tmp2, tmp2, transpose_a=True)
            - tf.linalg.matmul(tmp1, tmp1, transpose_a=True)
        )
        shape = tf.stack([1, 1, tf.shape(Y_data)[1]])
        var = tf.tile(tf.expand_dims(var, 2), shape)
    else:
        var = (
            self.kernel(Xnew, full_cov=False)
            + tf.reduce_sum(tf.square(tmp2), axis=0)
            - tf.reduce_sum(tf.square(tmp1), axis=0)
        )
        shape = tf.stack([1, tf.shape(Y_data)[1]])
        var = tf.tile(tf.expand_dims(var, 1), shape)
    return mean + self.mean_function(Xnew), var