# Path: disentanglement_lib/methods/shared/architectures.py
# def make_gaussian_encoder(input_tensor,
#                           is_training=True,
#                           num_latent=gin.REQUIRED,
#                           encoder_fn=gin.REQUIRED):
# def make_decoder(latent_tensor,
#                  output_shape,
#                  is_training=True,
#                  decoder_fn=gin.REQUIRED):
# def make_discriminator(input_tensor,
#                        is_training=False,
#                        discriminator_fn=gin.REQUIRED):
# def fc_encoder(input_tensor, num_latent, is_training=True):
# def conv_encoder(input_tensor, num_latent, is_training=True):
# def fc_decoder(latent_tensor, output_shape, is_training=True):
# def deconv_decoder(latent_tensor, output_shape, is_training=True):
# def fc_discriminator(input_tensor, is_training=True):
# def test_encoder(input_tensor, num_latent, is_training):
# def test_decoder(latent_tensor, output_shape, is_training=False):
# 
# 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 architectures  # pylint: disable=unused-import
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 numpy as np
import tensorflow.compat.v1 as tf
import gin.tf
    (https://arxiv.org/pdf/1802.05983).

    Args:
      factor_sizes: Size of each factor of variation.
      gamma: Hyperparameter for the unsupervised regularizer.
      gamma_sup: Hyperparameter for the supervised regularizer.
    """
    self.gamma = gamma
    self.gamma_sup = gamma_sup
    super(S2FactorVAE, self).__init__(factor_sizes)

  def model_fn(self, features, labels, mode, params):
    """TPUEstimator compatible model function."""
    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    labelled_features = labels[0]
    labels = tf.to_float(labels[1])
    data_shape = features.get_shape().as_list()[1:]
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      z_mean, z_logvar = self.gaussian_encoder(
          features, is_training=is_training)
      z_mean_labelled, _ = self.gaussian_encoder(
          labelled_features, is_training=is_training)
    z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar)
    z_shuffle = vae.shuffle_codes(z_sampled)
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      logits_z, probs_z = architectures.make_discriminator(
          z_sampled, is_training=is_training)
      _, probs_z_shuffle = architectures.make_discriminator(
          z_shuffle, is_training=is_training)
    reconstructions = self.decode(z_sampled, data_shape, is_training)

# Path: disentanglement_lib/methods/shared/architectures.py
# def make_gaussian_encoder(input_tensor,
#                           is_training=True,
#                           num_latent=gin.REQUIRED,
#                           encoder_fn=gin.REQUIRED):
# def make_decoder(latent_tensor,
#                  output_shape,
#                  is_training=True,
#                  decoder_fn=gin.REQUIRED):
# def make_discriminator(input_tensor,
#                        is_training=False,
#                        discriminator_fn=gin.REQUIRED):
# def fc_encoder(input_tensor, num_latent, is_training=True):
# def conv_encoder(input_tensor, num_latent, is_training=True):
# def fc_decoder(latent_tensor, output_shape, is_training=True):
# def deconv_decoder(latent_tensor, output_shape, is_training=True):
# def fc_discriminator(input_tensor, is_training=True):
# def test_encoder(input_tensor, num_latent, is_training):
# def test_decoder(latent_tensor, output_shape, is_training=False):
# 
# 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/semi_supervised/semi_supervised_vae.py
from disentanglement_lib.methods.shared import architectures  # pylint: disable=unused-import
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 numpy as np
import tensorflow.compat.v1 as tf
import gin.tf

    (https://arxiv.org/pdf/1802.05983).

    Args:
      factor_sizes: Size of each factor of variation.
      gamma: Hyperparameter for the unsupervised regularizer.
      gamma_sup: Hyperparameter for the supervised regularizer.
    """
    self.gamma = gamma
    self.gamma_sup = gamma_sup
    super(S2FactorVAE, self).__init__(factor_sizes)

  def model_fn(self, features, labels, mode, params):
    """TPUEstimator compatible model function."""
    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    labelled_features = labels[0]
    labels = tf.to_float(labels[1])
    data_shape = features.get_shape().as_list()[1:]
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      z_mean, z_logvar = self.gaussian_encoder(
          features, is_training=is_training)
      z_mean_labelled, _ = self.gaussian_encoder(
          labelled_features, is_training=is_training)
    z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar)
    z_shuffle = vae.shuffle_codes(z_sampled)
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      logits_z, probs_z = architectures.make_discriminator(
          z_sampled, is_training=is_training)
      _, probs_z_shuffle = architectures.make_discriminator(
          z_shuffle, is_training=is_training)
    reconstructions = self.decode(z_sampled, data_shape, is_training)

    per_sample_loss = losses.make_reconstruction_loss(features, reconstructions)