# 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/gaussian_encoder_model.py
# class GaussianEncoderModel(object):
#   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 sample_from_latent_distribution(self, z_mean, z_logvar):
# def export_as_tf_hub(gaussian_encoder_model,
#                      observation_shape,
#                      checkpoint_path,
#                      export_path,
#                      drop_collections=None):
#   def module_fn(is_training):


import math
import tensorflow.compat.v1 as tf
import gin.tf
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 gaussian_encoder_model
from six.moves import range
from six.moves import zip
from tensorflow.contrib import tpu as contrib_tpu
    """TPUEstimator compatible model function."""
    del labels
    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    data_shape = features.get_shape().as_list()[1:]
    z_mean, z_logvar = self.gaussian_encoder(features, is_training=is_training)
    z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar)
    z_shuffle = 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)
    reconstruction_loss = tf.reduce_mean(per_sample_loss)
    kl_loss = compute_gaussian_kl(z_mean, z_logvar)
    standard_vae_loss = tf.add(reconstruction_loss, kl_loss, name="VAE_loss")
    # tc = E[log(p_real)-log(p_fake)] = E[logit_real - logit_fake]
    tc_loss_per_sample = logits_z[:, 0] - logits_z[:, 1]
    tc_loss = tf.reduce_mean(tc_loss_per_sample, axis=0)
    regularizer = kl_loss + self.gamma * tc_loss
    factor_vae_loss = tf.add(
        standard_vae_loss, self.gamma * tc_loss, name="factor_VAE_loss")
    discr_loss = tf.add(
        0.5 * tf.reduce_mean(tf.log(probs_z[:, 0])),
        0.5 * tf.reduce_mean(tf.log(probs_z_shuffle[:, 1])),
        name="discriminator_loss")
    if mode == tf.estimator.ModeKeys.TRAIN:
      optimizer_vae = optimizers.make_vae_optimizer()

# 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/gaussian_encoder_model.py
# class GaussianEncoderModel(object):
#   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 sample_from_latent_distribution(self, z_mean, z_logvar):
# def export_as_tf_hub(gaussian_encoder_model,
#                      observation_shape,
#                      checkpoint_path,
#                      export_path,
#                      drop_collections=None):
#   def module_fn(is_training):


# Path: disentanglement_lib/methods/unsupervised/vae.py
import math
import tensorflow.compat.v1 as tf
import gin.tf
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 gaussian_encoder_model
from six.moves import range
from six.moves import zip
from tensorflow.contrib import tpu as contrib_tpu

    """TPUEstimator compatible model function."""
    del labels
    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    data_shape = features.get_shape().as_list()[1:]
    z_mean, z_logvar = self.gaussian_encoder(features, is_training=is_training)
    z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar)
    z_shuffle = 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)
    reconstruction_loss = tf.reduce_mean(per_sample_loss)
    kl_loss = compute_gaussian_kl(z_mean, z_logvar)
    standard_vae_loss = tf.add(reconstruction_loss, kl_loss, name="VAE_loss")
    # tc = E[log(p_real)-log(p_fake)] = E[logit_real - logit_fake]
    tc_loss_per_sample = logits_z[:, 0] - logits_z[:, 1]
    tc_loss = tf.reduce_mean(tc_loss_per_sample, axis=0)
    regularizer = kl_loss + self.gamma * tc_loss
    factor_vae_loss = tf.add(
        standard_vae_loss, self.gamma * tc_loss, name="factor_VAE_loss")
    discr_loss = tf.add(
        0.5 * tf.reduce_mean(tf.log(probs_z[:, 0])),
        0.5 * tf.reduce_mean(tf.log(probs_z_shuffle[:, 1])),
        name="discriminator_loss")
    if mode == tf.estimator.ModeKeys.TRAIN:
      optimizer_vae = optimizers.make_vae_optimizer()

      optimizer_discriminator = optimizers.make_discriminator_optimizer()