class_docstr: Riemannian Hamiltonian VAE model.
  
  
  Args:
      model_config (RHVAEConfig): A model configuration setting the main parameters of the model.
  
      encoder (BaseEncoder): An instance of BaseEncoder (inheriting from `torch.nn.Module` which
          plays the role of encoder. This argument allows you to use your own neural networks
          architectures if desired. If None is provided, a simple Multi Layer Preception
          (https://en.wikipedia.org/wiki/Multilayer_perceptron) is used. Default: None.
  
      decoder (BaseDecoder): An instance of BaseDecoder (inheriting from `torch.nn.Module` which
          plays the role of decoder. This argument allows you to use your own neural networks
          architectures if desired. If None is provided, a simple Multi Layer Preception
          (https://en.wikipedia.org/wiki/Multilayer_perceptron) is used. Default: None.
  
  
  .. note::
      For high dimensional data we advice you to provide you own network architectures. With the
      provided MLP you may end up with a ``MemoryError``.
class_docstr_tok: ["Riemannian","Hamiltonian","VAE","model",".","Args",":","model_config","(","RHVAEConfig",")",":","A","model","configuration","setting","the","main","parameters","of","the","model",".","encoder","(","BaseEncoder",")",":","An","instance","of","BaseEncoder","(","inheriting","from","`","torch.nn.Module","`","which","plays","the","role","of","encoder",".","This","argument","allows","you","to","use","your","own","neural","networks","architectures","if","desired",".","If","None","is","provided",",","a","simple","Multi","Layer","Preception","(","https",":","\/\/en.wikipedia.org\/wiki\/Multilayer_perceptron",")","is","used",".","Default",":","None",".","decoder","(","BaseDecoder",")",":","An","instance","of","BaseDecoder","(","inheriting","from","`","torch.nn.Module","`","which","plays","the","role","of","decoder",".","This","argument","allows","you","to","use","your","own","neural","networks","architectures","if","desired",".","If","None","is","provided",",","a","simple","Multi","Layer","Preception","(","https",":","\/\/en.wikipedia.org\/wiki\/Multilayer_perceptron",")","is","used",".","Default",":","None",".","..","note",":",":","For","high","dimensional","data","we","advice","you","to","provide","you","own","network","architectures",".","With","the","provided","MLP","you","may","end","up","with","a","``","MemoryError","``","."]
docstr: The input data is first encoded. The reparametrization is used to produce a sample
  :math:`z_0` from the approximate posterior :math:`q_{\phi}(z|x)`. Then Riemannian
  Hamiltonian equations are solved using the generalized leapfrog integrator. In the meantime,
  the input data :math:`x` is fed to the metric network outputing the matrices
  :math:`L_{\psi}`. The metric is computed and used with the integrator.
  
  Args:
      inputs (BaseDataset): The training data with labels
  
  Returns:
      ModelOutput: An instance of ModelOutput containing all the relevant parameters
docstr_tok: ["The","input","data","is","first","encoded",".","The","reparametrization","is","used","to","produce","a","sample",":","math",":","`","z_0","`","from","the","approximate","posterior",":","math",":","`","q_","{","\\phi","}","(","z|x",")","`",".","Then","Riemannian","Hamiltonian","equations","are","solved","using","the","generalized","leapfrog","integrator",".","In","the","meantime",",","the","input","data",":","math",":","`","x","`","is","fed","to","the","metric","network","outputing","the","matrices",":","math",":","`","L_","{","\\psi","}","`",".","The","metric","is","computed","and","used","with","the","integrator",".","Args",":","inputs","(","BaseDataset",")",":","The","training","data","with","labels","Returns",":","ModelOutput",":","An","instance","of","ModelOutput","containing","all","the","relevant","parameters"]
code: def forward(self, inputs: BaseDataset, **kwargs) -> ModelOutput:
      r"""
      The input data is first encoded. The reparametrization is used to produce a sample
      :math:`z_0` from the approximate posterior :math:`q_{\phi}(z|x)`. Then Riemannian
      Hamiltonian equations are solved using the generalized leapfrog integrator. In the meantime,
      the input data :math:`x` is fed to the metric network outputing the matrices
      :math:`L_{\psi}`. The metric is computed and used with the integrator.
  
      Args:
          inputs (BaseDataset): The training data with labels
  
      Returns:
          ModelOutput: An instance of ModelOutput containing all the relevant parameters
      """
  
      x = inputs["data"]
  
      encoder_output = self.encoder(x)
      mu, log_var = (
          encoder_output.embedding,
          encoder_output.log_covariance,
      )
  
      std = torch.exp(0.5 * log_var)
      z0, eps0 = self._sample_gauss(mu, std)
  
      z = z0
  
      if self.training:
          # update the metric using batch data points
          L = self.metric(x)["L"]
  
          M = L @ torch.transpose(L, 1, 2)
  
          # store LL^T and mu(x_i) to update final metric
          self.M.append(M.detach().clone())
          self.centroids.append(mu.detach().clone())
  
          G_inv = (
              M.unsqueeze(0)
              * torch.exp(
                  -torch.norm(mu.unsqueeze(0) - z.unsqueeze(1), dim=-1)
                  ** 2
                  / (self.temperature**2)
              )
              .unsqueeze(-1)
              .unsqueeze(-1)
          ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(
              x.device
          )
  
      else:
          G = self.G(z)
          G_inv = self.G_inv(z)
          L = torch.linalg.cholesky(G)
  
      G_log_det = -torch.logdet(G_inv)
  
      gamma = torch.randn_like(z0, device=x.device)
      rho = gamma / self.beta_zero_sqrt
      beta_sqrt_old = self.beta_zero_sqrt
  
      # sample \rho from N(0, G)
      rho = (L @ rho.unsqueeze(-1)).squeeze(-1)
  
      recon_x = self.decoder(z)["reconstruction"]
  
      for k in range(self.n_lf):
          # perform leapfrog steps
  
          # step 1
          rho_ = self._leap_step_1(recon_x, x, z, rho, G_inv, G_log_det)
  
          # step 2
          z = self._leap_step_2(recon_x, x, z, rho_, G_inv, G_log_det)
  
          recon_x = self.decoder(z)["reconstruction"]
  
          if self.training:
              G_inv = (
                  M.unsqueeze(0)
                  * torch.exp(
                      -torch.norm(
                          mu.unsqueeze(0) - z.unsqueeze(1), dim=-1
                      )
                      ** 2
                      / (self.temperature**2)
                  )
                  .unsqueeze(-1)
                  .unsqueeze(-1)
              ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(
                  x.device
              )
  
          else:
              # compute metric value on new z using final metric
              G = self.G(z)
              G_inv = self.G_inv(z)
  
          G_log_det = -torch.logdet(G_inv)
  
          # step 3
          rho__ = self._leap_step_3(
              recon_x, x, z, rho_, G_inv, G_log_det
          )
  
          # tempering
          beta_sqrt = self._tempering(k + 1, self.n_lf)
          rho = (beta_sqrt_old / beta_sqrt) * rho__
          beta_sqrt_old = beta_sqrt
  
      loss = self.loss_function(
          recon_x,
          x,
          z0,
          z,
          rho,
          eps0,
          gamma,
          mu,
          log_var,
          G_inv,
          G_log_det,
      )
  
      output = ModelOutput(
          loss=loss,
          recon_x=recon_x,
          z=z,
          z0=z0,
          rho=rho,
          eps0=eps0,
          gamma=gamma,
          mu=mu,
          log_var=log_var,
          G_inv=G_inv,
          G_log_det=G_log_det,
      )
  
      return output
  
imports: import inspect
  import logging
  import os
  import warnings
  from collections import deque
  from copy import deepcopy
  from typing import Optional
  import cloudpickle
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.autograd import grad
  from ...customexception import BadInheritanceError
  from ...data.datasets import BaseDataset
  from ..base.base_utils import CPU_Unpickler, ModelOutput, hf_hub_is_available
  from ..nn import BaseDecoder, BaseEncoder, BaseMetric
  from ..nn.default_architectures import Metric_MLP
  from ..vae import VAE
  from .rhvae_config import RHVAEConfig
  from .rhvae_utils import create_inverse_metric, create_metric
label_desc: Use image node_id 4 for calling the RHVAE obj's underlying member method code with example usage: obj.forward(inputs) and returns: output