# 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
    c = anneal(self.c_max, tf.train.get_global_step(), self.iteration_threshold)
    return self.gamma * tf.math.abs(kl_loss - c)


@gin.configurable("factor_vae")
class FactorVAE(BaseVAE):
  """FactorVAE model."""

  def __init__(self, gamma=gin.REQUIRED):
    """Creates a FactorVAE model.

    Implementing Eq. 2 of "Disentangling by Factorizing"
    (https://arxiv.org/pdf/1802.05983).

    Args:
      gamma: Hyperparameter for the regularizer.
    """
    self.gamma = gamma

  def model_fn(self, features, labels, mode, params):
    """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)

# 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

    c = anneal(self.c_max, tf.train.get_global_step(), self.iteration_threshold)
    return self.gamma * tf.math.abs(kl_loss - c)


@gin.configurable("factor_vae")
class FactorVAE(BaseVAE):
  """FactorVAE model."""

  def __init__(self, gamma=gin.REQUIRED):
    """Creates a FactorVAE model.

    Implementing Eq. 2 of "Disentangling by Factorizing"
    (https://arxiv.org/pdf/1802.05983).

    Args:
      gamma: Hyperparameter for the regularizer.
    """
    self.gamma = gamma

  def model_fn(self, features, labels, mode, params):
    """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(