# 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
    mean_sample_2, log_var_sample_2 = self.aggregate(
        z_mean_2, z_logvar_2, new_mean, new_log_var, labels, kl_per_point)

    z_sampled_1 = self.sample_from_latent_distribution(
        mean_sample_1, log_var_sample_1)
    z_sampled_2 = self.sample_from_latent_distribution(
        mean_sample_2, log_var_sample_2)
    with tf.variable_scope(
        tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      reconstructions_1 = self.decode(z_sampled_1, data_shape, is_training)
      reconstructions_2 = self.decode(z_sampled_2, data_shape, is_training)
    per_sample_loss_1 = losses.make_reconstruction_loss(
        features_1, reconstructions_1)
    per_sample_loss_2 = losses.make_reconstruction_loss(
        features_2, reconstructions_2)
    reconstruction_loss_1 = tf.reduce_mean(per_sample_loss_1)
    reconstruction_loss_2 = tf.reduce_mean(per_sample_loss_2)
    reconstruction_loss = (0.5 * reconstruction_loss_1 +
                           0.5 * reconstruction_loss_2)
    kl_loss_1 = vae.compute_gaussian_kl(mean_sample_1, log_var_sample_1)
    kl_loss_2 = vae.compute_gaussian_kl(mean_sample_2, log_var_sample_2)
    kl_loss = 0.5 * kl_loss_1 + 0.5 * kl_loss_2
    regularizer = self.regularizer(
        kl_loss, None, None, None)

    loss = tf.add(reconstruction_loss,
                  regularizer,
                  name="loss")
    elbo = tf.add(reconstruction_loss, kl_loss, name="elbo")
    if mode == tf.estimator.ModeKeys.TRAIN:

# 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

    mean_sample_2, log_var_sample_2 = self.aggregate(
        z_mean_2, z_logvar_2, new_mean, new_log_var, labels, kl_per_point)

    z_sampled_1 = self.sample_from_latent_distribution(
        mean_sample_1, log_var_sample_1)
    z_sampled_2 = self.sample_from_latent_distribution(
        mean_sample_2, log_var_sample_2)
    with tf.variable_scope(
        tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      reconstructions_1 = self.decode(z_sampled_1, data_shape, is_training)
      reconstructions_2 = self.decode(z_sampled_2, data_shape, is_training)
    per_sample_loss_1 = losses.make_reconstruction_loss(
        features_1, reconstructions_1)
    per_sample_loss_2 = losses.make_reconstruction_loss(
        features_2, reconstructions_2)
    reconstruction_loss_1 = tf.reduce_mean(per_sample_loss_1)
    reconstruction_loss_2 = tf.reduce_mean(per_sample_loss_2)
    reconstruction_loss = (0.5 * reconstruction_loss_1 +
                           0.5 * reconstruction_loss_2)
    kl_loss_1 = vae.compute_gaussian_kl(mean_sample_1, log_var_sample_1)
    kl_loss_2 = vae.compute_gaussian_kl(mean_sample_2, log_var_sample_2)
    kl_loss = 0.5 * kl_loss_1 + 0.5 * kl_loss_2
    regularizer = self.regularizer(
        kl_loss, None, None, None)

    loss = tf.add(reconstruction_loss,
                  regularizer,
                  name="loss")
    elbo = tf.add(reconstruction_loss, kl_loss, name="elbo")
    if mode == tf.estimator.ModeKeys.TRAIN:

      optimizer = optimizers.make_vae_optimizer()