# 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
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Library of losses for weakly-supervised disentanglement learning.

Implementation of weakly-supervised VAE based models from the paper
"Weakly-Supervised Disentanglement Without Compromises"
https://arxiv.org/pdf/2002.02886.pdf.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




@gin.configurable("weak_loss", blacklist=["z1", "z2", "labels"])
def make_weak_loss(z1, z2, labels, loss_fn=gin.REQUIRED):
  """Wrapper that creates weakly-supervised losses."""

  return loss_fn(z1, z2, labels)


@gin.configurable("group_vae")

# 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

#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Library of losses for weakly-supervised disentanglement learning.

Implementation of weakly-supervised VAE based models from the paper
"Weakly-Supervised Disentanglement Without Compromises"
https://arxiv.org/pdf/2002.02886.pdf.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




@gin.configurable("weak_loss", blacklist=["z1", "z2", "labels"])
def make_weak_loss(z1, z2, labels, loss_fn=gin.REQUIRED):
  """Wrapper that creates weakly-supervised losses."""

  return loss_fn(z1, z2, labels)


@gin.configurable("group_vae")

class GroupVAEBase(vae.BaseVAE):