# Path: disentanglement_lib/methods/weak/weak_vae.py
# def make_weak_loss(z1, z2, labels, loss_fn=gin.REQUIRED):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
#   def model_fn(self, features, labels, mode, params):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
#   def model_fn(self, features, labels, mode, params):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
# def aggregate_labels(z_mean, z_logvar, new_mean, new_log_var, labels,
#                      kl_per_point):
# def aggregate_argmax(z_mean, z_logvar, new_mean, new_log_var, labels,
#                      kl_per_point):
# def discretize_in_bins(x):
# def compute_kl(z_1, z_2, logvar_1, logvar_2):
# def make_metric_fn(*names):
#   def metric_fn(*args):
# class GroupVAEBase(vae.BaseVAE):
# class GroupVAELabels(GroupVAEBase):
# class GroupVAEArgmax(GroupVAEBase):
# class MLVae(vae.BaseVAE):
# class MLVaeLabels(MLVae):
# class MLVaeArgmax(MLVae):


from absl.testing import parameterized
from disentanglement_lib.methods.weak import weak_vae  # pylint: disable=unused-import
import numpy as np
import tensorflow as tf
from __future__ import division
from __future__ import print_function



class WeakVaeTest(parameterized.TestCase, tf.test.TestCase):

  @parameterized.parameters(
      (np.zeros([64, 10]),
       np.zeros([64, 10]),
       np.ones([64, 10]),
       np.ones([64, 10]),
       np.concatenate((np.zeros([64, 5]), np.ones([64, 5])), axis=1),
       np.concatenate((np.ones([64, 5]), np.zeros([64, 5])), axis=1)),
      (np.array([[1, 1]]),
       np.array([[1, 1]]),
       np.array([[0, 0]]),
       np.array([[0, 0]]),
       np.array([[0, 0.1]]),
       np.array([[0, 1]]))
      )
  def test_aggregate_argmax(self, z_mean, z_logvar, new_mean, new_log_var,
                            kl_per_point, target):

    mean_tf = tf.convert_to_tensor(z_mean, dtype=np.float32)
    logvar_tf = tf.convert_to_tensor(z_logvar, dtype=np.float32)
    new_mean_tf = tf.convert_to_tensor(new_mean, dtype=np.float32)
    new_log_var_tf = tf.convert_to_tensor(new_log_var, dtype=np.float32)
    kl_per_point_tf = tf.convert_to_tensor(kl_per_point, dtype=np.float32)
    with self.session() as sess:

# Path: disentanglement_lib/methods/weak/weak_vae.py
# def make_weak_loss(z1, z2, labels, loss_fn=gin.REQUIRED):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
#   def model_fn(self, features, labels, mode, params):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
#   def model_fn(self, features, labels, mode, params):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
# def aggregate_labels(z_mean, z_logvar, new_mean, new_log_var, labels,
#                      kl_per_point):
# def aggregate_argmax(z_mean, z_logvar, new_mean, new_log_var, labels,
#                      kl_per_point):
# def discretize_in_bins(x):
# def compute_kl(z_1, z_2, logvar_1, logvar_2):
# def make_metric_fn(*names):
#   def metric_fn(*args):
# class GroupVAEBase(vae.BaseVAE):
# class GroupVAELabels(GroupVAEBase):
# class GroupVAEArgmax(GroupVAEBase):
# class MLVae(vae.BaseVAE):
# class MLVaeLabels(MLVae):
# class MLVaeArgmax(MLVae):


# Path: disentanglement_lib/methods/weak/weak_vae_test.py
from absl.testing import parameterized
from disentanglement_lib.methods.weak import weak_vae  # pylint: disable=unused-import
import numpy as np
import tensorflow as tf

from __future__ import division
from __future__ import print_function



class WeakVaeTest(parameterized.TestCase, tf.test.TestCase):

  @parameterized.parameters(
      (np.zeros([64, 10]),
       np.zeros([64, 10]),
       np.ones([64, 10]),
       np.ones([64, 10]),
       np.concatenate((np.zeros([64, 5]), np.ones([64, 5])), axis=1),
       np.concatenate((np.ones([64, 5]), np.zeros([64, 5])), axis=1)),
      (np.array([[1, 1]]),
       np.array([[1, 1]]),
       np.array([[0, 0]]),
       np.array([[0, 0]]),
       np.array([[0, 0.1]]),
       np.array([[0, 1]]))
      )
  def test_aggregate_argmax(self, z_mean, z_logvar, new_mean, new_log_var,
                            kl_per_point, target):

    mean_tf = tf.convert_to_tensor(z_mean, dtype=np.float32)
    logvar_tf = tf.convert_to_tensor(z_logvar, dtype=np.float32)
    new_mean_tf = tf.convert_to_tensor(new_mean, dtype=np.float32)
    new_log_var_tf = tf.convert_to_tensor(new_log_var, dtype=np.float32)
    kl_per_point_tf = tf.convert_to_tensor(kl_per_point, dtype=np.float32)
    with self.session() as sess:

      test_value = sess.run(weak_vae.aggregate_argmax(