5e599fb01df65d156a40f7a138ab6627a06a50db,gpflow/optimizers/natgrad.py,NaturalGradient,_natgrad_steps,#NaturalGradient#Any#Any#,113

Before Change

            self._natgrad_step(loss_fn, q_mu, q_sqrt, xi_transform)

        with tf.name_scope(f"{self._name}/natural_gradient_steps"):
            list(map(natural_gradient_step, *zip(*parameters)))

    def _natgrad_step(
        self, loss_fn: Callable, q_mu: Parameter, q_sqrt: Parameter, xi_transform: XiTransform
    ):

After Change

        :param parameters: List of tuples (q_mu, q_sqrt, xi_transform)
        
        q_mus, q_sqrts, xis = zip(*parameters)
        unconstrained_variables = [
            p.unconstrained_variable for params in (q_mus, q_sqrts) for p in params
        ]

        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(unconstrained_variables)
            loss = loss_fn()

Instances