class_docstr: Implements a natural gradient descent optimizer for variational models
  that are based on a distribution q(u) = N(q_mu, q_sqrt q_sqrtᵀ) that is
  parameterized by mean q_mu and lower-triangular Cholesky factor q_sqrt
  of the covariance.
  
  Note that this optimizer does not implement the standard API of
  tf.optimizers.Optimizer. Its only public method is minimize(), which has
  a custom signature (var_list needs to be a list of (q_mu, q_sqrt) tuples,
  where q_mu and q_sqrt are gpflow.Parameter instances, not tf.Variable).
  
  Note furthermore that the natural gradients are implemented only for the
  full covariance case (i.e., q_diag=True is NOT supported).
  
  When using in your work, please cite :cite:t:`salimbeni18`.
class_docstr_tok: ["Implements","a","natural","gradient","descent","optimizer","for","variational","models","that","are","based","on","a","distribution","q","(","u",")","=","N","(","q_mu",",","q_sqrt","q_sqrt\u1d40",")","that","is","parameterized","by","mean","q_mu","and","lower-triangular","Cholesky","factor","q_sqrt","of","the","covariance",".","Note","that","this","optimizer","does","not","implement","the","standard","API","of","tf.optimizers.Optimizer",".","Its","only","public","method","is","minimize","(",")",",","which","has","a","custom","signature","(","var_list","needs","to","be","a","list","of","(","q_mu",",","q_sqrt",")","tuples",",","where","q_mu","and","q_sqrt","are","gpflow.Parameter","instances",",","not","tf.Variable",")",".","Note","furthermore","that","the","natural","gradients","are","implemented","only","for","the","full","covariance","case","(","i.e.",",","q_diag=True","is","NOT","supported",")",".","When","using","in","your","work",",","please","cite",":","cite",":","t",":","`","salimbeni18","`","."]
docstr: This function does the backward step on the q_mu and q_sqrt parameters,
  given the gradients of the loss function with respect to their unconstrained
  variables. I.e., it expects the arguments to come from
  
      with tf.GradientTape() as tape:
          loss = loss_function()
      q_mu_grad, q_mu_sqrt = tape.gradient(loss, [q_mu, q_sqrt])
  
  (Note that tape.gradient() returns the gradients in *unconstrained* space!)
  
  Implements equation [10] from :cite:t:`salimbeni18`.
  
  In addition, for convenience with the rest of GPflow, this code computes ∂L/∂η using
  the chain rule (the following assumes a numerator layout where the gradient is a row
  vector; note that TensorFlow actually returns a column vector), where L is the loss:
  
  ∂L/∂η = (∂L / ∂[q_mu, q_sqrt])(∂[q_mu, q_sqrt] / ∂η)
  
  In total there are three derivative calculations:
  natgrad of L w.r.t ξ  = (∂ξ / ∂θ) [(∂L / ∂[q_mu, q_sqrt]) (∂[q_mu, q_sqrt] / ∂η)]ᵀ
  
  Note that if ξ = θ (i.e. [q_mu, q_sqrt]) some of these calculations are the identity.
  In the code η = eta, ξ = xi, θ = nat.
  
  :param q_mu_grad: gradient of loss w.r.t. q_mu (in unconstrained space)
  :param q_sqrt_grad: gradient of loss w.r.t. q_sqrt (in unconstrained space)
  :param q_mu: parameter for the mean of q(u) with shape [M, L]
  :param q_sqrt: parameter for the square root of the covariance of q(u)
      with shape [L, M, M] (the diagonal parametrization, q_diag=True, is NOT supported)
  :param xi_transform: the ξ transform to use (self.xi_transform if not specified)
docstr_tok: ["This","function","does","the","backward","step","on","the","q_mu","and","q_sqrt","parameters",",","given","the","gradients","of","the","loss","function","with","respect","to","their","unconstrained","variables",".","I.e.",",","it","expects","the","arguments","to","come","from","with","tf.GradientTape","(",")","as","tape",":","loss","=","loss_function","(",")","q_mu_grad",",","q_mu_sqrt","=","tape.gradient","(","loss",",","[","q_mu",",","q_sqrt","]",")","(","Note","that","tape.gradient","(",")","returns","the","gradients","in","*","unconstrained","*","space","!",")","Implements","equation","[","10","]","from",":","cite",":","t",":","`","salimbeni18","`",".","In","addition",",","for","convenience","with","the","rest","of","GPflow",",","this","code","computes","\u2202L\/\u2202\u03b7","using","the","chain","rule","(","the","following","assumes","a","numerator","layout","where","the","gradient","is","a","row","vector",";","note","that","TensorFlow","actually","returns","a","column","vector",")",",","where","L","is","the","loss",":","\u2202L\/\u2202\u03b7","=","(","\u2202L","\/","\u2202","[","q_mu",",","q_sqrt","]",")","(","\u2202","[","q_mu",",","q_sqrt","]","\/","\u2202\u03b7",")","In","total","there","are","three","derivative","calculations",":","natgrad","of","L","w.r.t","\u03be","=","(","\u2202\u03be","\/","\u2202\u03b8",")","[","(","\u2202L","\/","\u2202","[","q_mu",",","q_sqrt","]",")","(","\u2202","[","q_mu",",","q_sqrt","]","\/","\u2202\u03b7",")","]","\u1d40","Note","that","if","\u03be","=","\u03b8","(","i.e",".","[","q_mu",",","q_sqrt","]",")","some","of","these","calculations","are","the","identity",".","In","the","code","\u03b7","=","eta",",","\u03be","=","xi",",","\u03b8","=","nat",".",":","param","q_mu_grad",":","gradient","of","loss","w.r.t",".","q_mu","(","in","unconstrained","space",")",":","param","q_sqrt_grad",":","gradient","of","loss","w.r.t",".","q_sqrt","(","in","unconstrained","space",")",":","param","q_mu",":","parameter","for","the","mean","of","q","(","u",")","with","shape","[","M",",","L","]",":","param","q_sqrt",":","parameter","for","the","square","root","of","the","covariance","of","q","(","u",")","with","shape","[","L",",","M",",","M","]","(","the","diagonal","parametrization",",","q_diag=True",",","is","NOT","supported",")",":","param","xi_transform",":","the","\u03be","transform","to","use","(","self.xi_transform","if","not","specified",")"]
code: def _natgrad_apply_gradients(
      self,
      q_mu_grad: tf.Tensor,
      q_sqrt_grad: tf.Tensor,
      q_mu: Parameter,
      q_sqrt: Parameter,
      xi_transform: Optional[XiTransform] = None,
  ) -> None:
      """
      This function does the backward step on the q_mu and q_sqrt parameters,
      given the gradients of the loss function with respect to their unconstrained
      variables. I.e., it expects the arguments to come from
  
          with tf.GradientTape() as tape:
              loss = loss_function()
          q_mu_grad, q_mu_sqrt = tape.gradient(loss, [q_mu, q_sqrt])
  
      (Note that tape.gradient() returns the gradients in *unconstrained* space!)
  
      Implements equation [10] from :cite:t:`salimbeni18`.
  
      In addition, for convenience with the rest of GPflow, this code computes ∂L/∂η using
      the chain rule (the following assumes a numerator layout where the gradient is a row
      vector; note that TensorFlow actually returns a column vector), where L is the loss:
  
      ∂L/∂η = (∂L / ∂[q_mu, q_sqrt])(∂[q_mu, q_sqrt] / ∂η)
  
      In total there are three derivative calculations:
      natgrad of L w.r.t ξ  = (∂ξ / ∂θ) [(∂L / ∂[q_mu, q_sqrt]) (∂[q_mu, q_sqrt] / ∂η)]ᵀ
  
      Note that if ξ = θ (i.e. [q_mu, q_sqrt]) some of these calculations are the identity.
      In the code η = eta, ξ = xi, θ = nat.
  
      :param q_mu_grad: gradient of loss w.r.t. q_mu (in unconstrained space)
      :param q_sqrt_grad: gradient of loss w.r.t. q_sqrt (in unconstrained space)
      :param q_mu: parameter for the mean of q(u) with shape [M, L]
      :param q_sqrt: parameter for the square root of the covariance of q(u)
          with shape [L, M, M] (the diagonal parametrization, q_diag=True, is NOT supported)
      :param xi_transform: the ξ transform to use (self.xi_transform if not specified)
      """
      if xi_transform is None:
          xi_transform = self.xi_transform
  
      # 1) the ordinary gpflow gradient
      dL_dmean = _to_constrained(q_mu_grad, q_mu.transform)
      dL_dvarsqrt = _to_constrained(q_sqrt_grad, q_sqrt.transform)
  
      with tf.GradientTape(
          persistent=True, watch_accessed_variables=False
      ) as tape:
          tape.watch(
              [
                  q_mu.unconstrained_variable,
                  q_sqrt.unconstrained_variable,
              ]
          )
  
          # the three parameterizations as functions of [q_mu, q_sqrt]
          eta1, eta2 = meanvarsqrt_to_expectation(q_mu, q_sqrt)
          # we need these to calculate the relevant gradients
          meanvarsqrt = expectation_to_meanvarsqrt(eta1, eta2)
  
          if not isinstance(xi_transform, XiNat):
              nat1, nat2 = meanvarsqrt_to_natural(q_mu, q_sqrt)
              xi1_nat, xi2_nat = xi_transform.naturals_to_xi(nat1, nat2)
              dummy_tensors = tf.ones_like(xi1_nat), tf.ones_like(
                  xi2_nat
              )
              with tf.GradientTape(
                  watch_accessed_variables=False
              ) as forward_tape:
                  forward_tape.watch(dummy_tensors)
                  dummy_gradients = tape.gradient(
                      [xi1_nat, xi2_nat],
                      [nat1, nat2],
                      output_gradients=dummy_tensors,
                  )
  
      # 2) the chain rule to get ∂L/∂η, where η (eta) are the expectation parameters
      dL_deta1, dL_deta2 = tape.gradient(
          meanvarsqrt,
          [eta1, eta2],
          output_gradients=[dL_dmean, dL_dvarsqrt],
      )
  
      if not isinstance(xi_transform, XiNat):
          nat_dL_xi1, nat_dL_xi2 = forward_tape.gradient(
              dummy_gradients,
              dummy_tensors,
              output_gradients=[dL_deta1, dL_deta2],
          )
      else:
          nat_dL_xi1, nat_dL_xi2 = dL_deta1, dL_deta2
  
      del tape  # Remove "persistent" tape
  
      xi1, xi2 = xi_transform.meanvarsqrt_to_xi(q_mu, q_sqrt)
      xi1_new = xi1 - self.gamma * nat_dL_xi1
      xi2_new = xi2 - self.gamma * nat_dL_xi2
  
      # Transform back to the model parameters [q_mu, q_sqrt]
      mean_new, varsqrt_new = xi_transform.xi_to_meanvarsqrt(
          xi1_new, xi2_new
      )
  
      q_mu.assign(mean_new)
      q_sqrt.assign(varsqrt_new)
  
imports: import abc
  import functools
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
  import tensorflow
  from check_shapes import check_shapes
  from ..base import AnyNDArray, Parameter, _to_constrained
label_desc: Use image node_id 4 for calling the NaturalGradient obj's underlying member method code with example usage: obj._natgrad_apply_gradients(q_mu_grad, q_sqrt_grad, q_mu, q_sqrt, xi_transform) without return types