# 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 absl.testing import parameterized
from disentanglement_lib.methods.unsupervised import vae
import numpy as np
import tensorflow.compat.v1 as tf
import gin.tf.external_configurables  # pylint: disable=unused-import
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     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.

"""Tests for vae.py."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function


def _make_symmetric_psd(matrix):
  return 0.5 * (matrix + matrix.T) + np.diag(np.ones(10)) * 10.


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

  @parameterized.parameters((np.zeros([10, 10]), np.zeros([10, 10]), 0., 0.01),
                            (np.ones([10, 10]), np.zeros([10, 10]), 5., 5.01),
                            (np.ones([10, 10]), np.ones([10, 10]), 8.58, 8.6))
  def test_compute_gaussian_kl(self, mean, logvar, target_low, target_high):
    mean_tf = tf.convert_to_tensor(mean, dtype=np.float32)
    logvar_tf = tf.convert_to_tensor(logvar, dtype=np.float32)
    with self.test_session() as sess:

# 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/unsupervised/vae_test.py
from absl.testing import parameterized
from disentanglement_lib.methods.unsupervised import vae
import numpy as np
import tensorflow.compat.v1 as tf
import gin.tf.external_configurables  # pylint: disable=unused-import

# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     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.

"""Tests for vae.py."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function


def _make_symmetric_psd(matrix):
  return 0.5 * (matrix + matrix.T) + np.diag(np.ones(10)) * 10.


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

  @parameterized.parameters((np.zeros([10, 10]), np.zeros([10, 10]), 0., 0.01),
                            (np.ones([10, 10]), np.zeros([10, 10]), 5., 5.01),
                            (np.ones([10, 10]), np.ones([10, 10]), 8.58, 8.6))
  def test_compute_gaussian_kl(self, mean, logvar, target_low, target_high):
    mean_tf = tf.convert_to_tensor(mean, dtype=np.float32)
    logvar_tf = tf.convert_to_tensor(logvar, dtype=np.float32)
    with self.test_session() as sess:

      test_value = sess.run(vae.compute_gaussian_kl(mean_tf, logvar_tf))