# Path: disentanglement_lib/methods/shared/losses.py
# def bernoulli_loss(true_images,
#                    reconstructed_images,
#                    activation,
#                    subtract_true_image_entropy=False):
# def l2_loss(true_images, reconstructed_images, activation):
# def make_reconstruction_loss(true_images,
#                              reconstructed_images,
#                              loss_fn=gin.REQUIRED,
#                              activation="logits"):
# 
# Path: disentanglement_lib/methods/shared/optimizers.py
# def make_optimizer(optimizer_fn, learning_rate):
# def make_vae_optimizer(optimizer_fn=gin.REQUIRED, learning_rate=None):
# def make_discriminator_optimizer(optimizer_fn=gin.REQUIRED, learning_rate=None):
# 
# Path: disentanglement_lib/methods/unsupervised/vae.py
# class BaseVAE(gaussian_encoder_model.GaussianEncoderModel):
# class BetaVAE(BaseVAE):
# class AnnealedVAE(BaseVAE):
# class FactorVAE(BaseVAE):
# class DIPVAE(BaseVAE):
# class BetaTCVAE(BaseVAE):
#   def model_fn(self, features, labels, mode, params):
#   def gaussian_encoder(self, input_tensor, is_training):
#   def decode(self, latent_tensor, observation_shape, is_training):
# def shuffle_codes(z):
# def compute_gaussian_kl(z_mean, z_logvar):
# def make_metric_fn(*names):
#   def metric_fn(*args):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
# def anneal(c_max, step, iteration_threshold):
#   def __init__(self,
#                gamma=gin.REQUIRED,
#                c_max=gin.REQUIRED,
#                iteration_threshold=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
#   def __init__(self, gamma=gin.REQUIRED):
#   def model_fn(self, features, labels, mode, params):
# def compute_covariance_z_mean(z_mean):
# def regularize_diag_off_diag_dip(covariance_matrix, lambda_od, lambda_d):
#   def __init__(self,
#                lambda_od=gin.REQUIRED,
#                lambda_d_factor=gin.REQUIRED,
#                dip_type="i"):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
# def gaussian_log_density(samples, mean, log_var):
# def total_correlation(z, z_mean, z_logvar):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):


from disentanglement_lib.methods.shared import losses  # pylint: disable=unused-import
from disentanglement_lib.methods.shared import optimizers  # pylint: disable=unused-import
from disentanglement_lib.methods.unsupervised import vae
from six.moves import zip
from tensorflow_estimator.python.estimator.tpu.tpu_estimator import TPUEstimatorSpec
import tensorflow.compat.v1 as tf
import gin.tf
          mode=mode,
          loss=loss,
          eval_metrics=(make_metric_fn("reconstruction_loss", "elbo",
                                       "regularizer", "kl_loss"),
                        [reconstruction_loss, -elbo, regularizer, kl_loss]))
    else:
      raise NotImplementedError("Eval mode not supported.")


@gin.configurable("group_vae_labels")
class GroupVAELabels(GroupVAEBase):
  """Class implementing the group-VAE with labels on which factor is shared."""

  def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
                kl_per_point):
    return aggregate_labels(z_mean, z_logvar, new_mean, new_log_var, labels,
                            kl_per_point)


@gin.configurable("group_vae_argmax")
class GroupVAEArgmax(GroupVAEBase):
  """Class implementing the group-VAE without any label."""

  def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
                kl_per_point):
    return aggregate_argmax(z_mean, z_logvar, new_mean, new_log_var, labels,
                            kl_per_point)


@gin.configurable("mlvae")

# Path: disentanglement_lib/methods/shared/losses.py
# def bernoulli_loss(true_images,
#                    reconstructed_images,
#                    activation,
#                    subtract_true_image_entropy=False):
# def l2_loss(true_images, reconstructed_images, activation):
# def make_reconstruction_loss(true_images,
#                              reconstructed_images,
#                              loss_fn=gin.REQUIRED,
#                              activation="logits"):
# 
# Path: disentanglement_lib/methods/shared/optimizers.py
# def make_optimizer(optimizer_fn, learning_rate):
# def make_vae_optimizer(optimizer_fn=gin.REQUIRED, learning_rate=None):
# def make_discriminator_optimizer(optimizer_fn=gin.REQUIRED, learning_rate=None):
# 
# Path: disentanglement_lib/methods/unsupervised/vae.py
# class BaseVAE(gaussian_encoder_model.GaussianEncoderModel):
# class BetaVAE(BaseVAE):
# class AnnealedVAE(BaseVAE):
# class FactorVAE(BaseVAE):
# class DIPVAE(BaseVAE):
# class BetaTCVAE(BaseVAE):
#   def model_fn(self, features, labels, mode, params):
#   def gaussian_encoder(self, input_tensor, is_training):
#   def decode(self, latent_tensor, observation_shape, is_training):
# def shuffle_codes(z):
# def compute_gaussian_kl(z_mean, z_logvar):
# def make_metric_fn(*names):
#   def metric_fn(*args):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
# def anneal(c_max, step, iteration_threshold):
#   def __init__(self,
#                gamma=gin.REQUIRED,
#                c_max=gin.REQUIRED,
#                iteration_threshold=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
#   def __init__(self, gamma=gin.REQUIRED):
#   def model_fn(self, features, labels, mode, params):
# def compute_covariance_z_mean(z_mean):
# def regularize_diag_off_diag_dip(covariance_matrix, lambda_od, lambda_d):
#   def __init__(self,
#                lambda_od=gin.REQUIRED,
#                lambda_d_factor=gin.REQUIRED,
#                dip_type="i"):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
# def gaussian_log_density(samples, mean, log_var):
# def total_correlation(z, z_mean, z_logvar):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):


# Path: disentanglement_lib/methods/weak/weak_vae.py
from disentanglement_lib.methods.shared import losses  # pylint: disable=unused-import
from disentanglement_lib.methods.shared import optimizers  # pylint: disable=unused-import
from disentanglement_lib.methods.unsupervised import vae
from six.moves import zip
from tensorflow_estimator.python.estimator.tpu.tpu_estimator import TPUEstimatorSpec
import tensorflow.compat.v1 as tf
import gin.tf

          mode=mode,
          loss=loss,
          eval_metrics=(make_metric_fn("reconstruction_loss", "elbo",
                                       "regularizer", "kl_loss"),
                        [reconstruction_loss, -elbo, regularizer, kl_loss]))
    else:
      raise NotImplementedError("Eval mode not supported.")


@gin.configurable("group_vae_labels")
class GroupVAELabels(GroupVAEBase):
  """Class implementing the group-VAE with labels on which factor is shared."""

  def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
                kl_per_point):
    return aggregate_labels(z_mean, z_logvar, new_mean, new_log_var, labels,
                            kl_per_point)


@gin.configurable("group_vae_argmax")
class GroupVAEArgmax(GroupVAEBase):
  """Class implementing the group-VAE without any label."""

  def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
                kl_per_point):
    return aggregate_argmax(z_mean, z_logvar, new_mean, new_log_var, labels,
                            kl_per_point)


@gin.configurable("mlvae")

class MLVae(vae.BaseVAE):