# Buggy code
def elbo(self) -> tf.Tensor:
    """
    Construct a tensorflow function to compute the bound on the marginal
    likelihood.
    """
    Y_data = self.data

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

    num_inducing = len(self.inducing_variable)
    psi0 = tf.reduce_sum(expectation(pX, self.kernel))
    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,
    )
    cov_uu = covariances.Kuu(
        self.inducing_variable, self.kernel, jitter=default_jitter()
    )
    L = tf.linalg.cholesky(cov_uu)
    sigma2 = self.likelihood.variance
    sigma = tf.sqrt(sigma2)

    # Compute intermediate matrices
    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)
    log_det_B = 2.0 * tf.reduce_sum(tf.math.log(tf.linalg.diag_part(LB)))
    c = tf.linalg.triangular_solve(LB, tf.linalg.matmul(A, Y_data), lower=True) / sigma

    # KL[q(x) || p(x)]
    dX_data_var = (
        self.X_data_var
        if self.X_data_var.shape.ndims == 2
        else tf.linalg.diag_part(self.X_data_var)
    )
    NQ = to_default_float(tf.size(self.X_data_mean))
    D = to_default_float(tf.shape(Y_data)[1])
    KL = -0.5 * tf.reduce_sum(tf.math.log(dX_data_var))
    KL += 0.5 * tf.reduce_sum(tf.math.log(self.X_prior_var))
    KL -= 0.5 * NQ
    KL += 0.5 * tf.reduce_sum(
        (tf.square(self.X_data_mean - self.X_prior_mean) + dX_data_var)
        / self.X_prior_var
    )

    # compute log marginal bound
    ND = to_default_float(tf.size(Y_data))
    bound = -0.5 * ND * tf.math.log(2 * np.pi * sigma2)
    bound += -0.5 * D * log_det_B
    bound += -0.5 * tf.reduce_sum(tf.square(Y_data)) / sigma2
    bound += 0.5 * tf.reduce_sum(tf.square(c))
    bound += (
        -0.5
        * D
        * (tf.reduce_sum(psi0) / sigma2 - tf.reduce_sum(tf.linalg.diag_part(AAT)))
    )
    bound -= KL
    return bound


# 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 elbo(self) -> tf.Tensor:
    """
    Construct a tensorflow function to compute the bound on the marginal
    likelihood.
    """
    Y_data = self.data

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

    num_inducing = self.inducing_variable.num_inducing
    psi0 = tf.reduce_sum(expectation(pX, self.kernel))
    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,
    )
    cov_uu = covariances.Kuu(
        self.inducing_variable, self.kernel, jitter=default_jitter()
    )
    L = tf.linalg.cholesky(cov_uu)
    sigma2 = self.likelihood.variance
    sigma = tf.sqrt(sigma2)

    # Compute intermediate matrices
    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)
    log_det_B = 2.0 * tf.reduce_sum(tf.math.log(tf.linalg.diag_part(LB)))
    c = tf.linalg.triangular_solve(LB, tf.linalg.matmul(A, Y_data), lower=True) / sigma

    # KL[q(x) || p(x)]
    dX_data_var = (
        self.X_data_var
        if self.X_data_var.shape.ndims == 2
        else tf.linalg.diag_part(self.X_data_var)
    )
    NQ = to_default_float(tf.size(self.X_data_mean))
    D = to_default_float(tf.shape(Y_data)[1])
    KL = -0.5 * tf.reduce_sum(tf.math.log(dX_data_var))
    KL += 0.5 * tf.reduce_sum(tf.math.log(self.X_prior_var))
    KL -= 0.5 * NQ
    KL += 0.5 * tf.reduce_sum(
        (tf.square(self.X_data_mean - self.X_prior_mean) + dX_data_var)
        / self.X_prior_var
    )

    # compute log marginal bound
    ND = to_default_float(tf.size(Y_data))
    bound = -0.5 * ND * tf.math.log(2 * np.pi * sigma2)
    bound += -0.5 * D * log_det_B
    bound += -0.5 * tf.reduce_sum(tf.square(Y_data)) / sigma2
    bound += 0.5 * tf.reduce_sum(tf.square(c))
    bound += (
        -0.5
        * D
        * (tf.reduce_sum(psi0) / sigma2 - tf.reduce_sum(tf.linalg.diag_part(AAT)))
    )
    bound -= KL
    return bound