# Path: disentanglement_lib/data/ground_truth/named_data.py
# def get_named_ground_truth_data(name):
# 
# Path: disentanglement_lib/evaluation/udr/metrics/udr.py
# def relative_strength_disentanglement(corr_matrix):
# def spearman_correlation_conv(vec1, vec2):
# def lasso_correlation_matrix(vec1, vec2, random_state=None):
# def _generate_representation_batch(ground_truth_data, representation_functions,
#                                    batch_size, random_state):
# def _generate_representation_dataset(ground_truth_data,
#                                      representation_functions, batch_size,
#                                      num_data_points, random_state):
# def compute_udr_sklearn(ground_truth_data,
#                         representation_functions,
#                         random_state,
#                         batch_size,
#                         num_data_points,
#                         correlation_matrix="lasso",
#                         filter_low_kl=True,
#                         include_raw_correlations=True,
#                         kl_filter_threshold=0.01):
# 
# Path: disentanglement_lib/utils/results.py
# def update_result_directory(result_directory,
#                             step_name,
#                             results_dict,
#                             old_result_directory=None):
# def _copy_recursively(path_to_old_dir, path_to_new_dir):
# def copydir(path_to_old_dir, path_to_new_dir):
# def save_gin(config_path):
#   def default(self, obj):
# def save_dict(config_path, dict_with_info):
# def gin_dict(config_path=None):
# def namespaced_dict(base_dict=None, **named_dicts):
# def aggregate_json_results(base_path):
# class Encoder(json.JSONEncoder):


import contextlib
import os
import time
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
import gin.tf
from absl import flags
from disentanglement_lib.data.ground_truth import named_data
from disentanglement_lib.evaluation.udr.metrics import udr  # pylint: disable=unused-import
from disentanglement_lib.utils import results
    representation_functions = []
    eval_functions = [
        stack.enter_context(
            hub.eval_function_for_module(os.path.join(model_dir, "tfhub")))
        for model_dir in model_dirs
    ]
    for f in eval_functions:

      def _representation_function(x, f=f):

        def compute_gaussian_kl(z_mean, z_logvar):
          return np.mean(
              0.5 * (np.square(z_mean) + np.exp(z_logvar) - z_logvar - 1),
              axis=0)

        encoding = f(dict(images=x), signature="gaussian_encoder", as_dict=True)

        return np.array(encoding["mean"]), compute_gaussian_kl(
            np.array(encoding["mean"]), np.array(encoding["logvar"]))

      representation_functions.append(_representation_function)

    results_dict = evaluation_fn(
        dataset,
        representation_functions,
        random_state=np.random.RandomState(random_seed))

  original_results_dir = os.path.join(model_dirs[0], "results")
  results_dir = os.path.join(output_dir, "results")
  results_dict["elapsed_time"] = time.time() - experiment_timer

# Path: disentanglement_lib/data/ground_truth/named_data.py
# def get_named_ground_truth_data(name):
# 
# Path: disentanglement_lib/evaluation/udr/metrics/udr.py
# def relative_strength_disentanglement(corr_matrix):
# def spearman_correlation_conv(vec1, vec2):
# def lasso_correlation_matrix(vec1, vec2, random_state=None):
# def _generate_representation_batch(ground_truth_data, representation_functions,
#                                    batch_size, random_state):
# def _generate_representation_dataset(ground_truth_data,
#                                      representation_functions, batch_size,
#                                      num_data_points, random_state):
# def compute_udr_sklearn(ground_truth_data,
#                         representation_functions,
#                         random_state,
#                         batch_size,
#                         num_data_points,
#                         correlation_matrix="lasso",
#                         filter_low_kl=True,
#                         include_raw_correlations=True,
#                         kl_filter_threshold=0.01):
# 
# Path: disentanglement_lib/utils/results.py
# def update_result_directory(result_directory,
#                             step_name,
#                             results_dict,
#                             old_result_directory=None):
# def _copy_recursively(path_to_old_dir, path_to_new_dir):
# def copydir(path_to_old_dir, path_to_new_dir):
# def save_gin(config_path):
#   def default(self, obj):
# def save_dict(config_path, dict_with_info):
# def gin_dict(config_path=None):
# def namespaced_dict(base_dict=None, **named_dicts):
# def aggregate_json_results(base_path):
# class Encoder(json.JSONEncoder):


# Path: disentanglement_lib/evaluation/udr/evaluate.py
import contextlib
import os
import time
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
import gin.tf
from absl import flags
from disentanglement_lib.data.ground_truth import named_data
from disentanglement_lib.evaluation.udr.metrics import udr  # pylint: disable=unused-import
from disentanglement_lib.utils import results

    representation_functions = []
    eval_functions = [
        stack.enter_context(
            hub.eval_function_for_module(os.path.join(model_dir, "tfhub")))
        for model_dir in model_dirs
    ]
    for f in eval_functions:

      def _representation_function(x, f=f):

        def compute_gaussian_kl(z_mean, z_logvar):
          return np.mean(
              0.5 * (np.square(z_mean) + np.exp(z_logvar) - z_logvar - 1),
              axis=0)

        encoding = f(dict(images=x), signature="gaussian_encoder", as_dict=True)

        return np.array(encoding["mean"]), compute_gaussian_kl(
            np.array(encoding["mean"]), np.array(encoding["logvar"]))

      representation_functions.append(_representation_function)

    results_dict = evaluation_fn(
        dataset,
        representation_functions,
        random_state=np.random.RandomState(random_seed))

  original_results_dir = os.path.join(model_dirs[0], "results")
  results_dir = os.path.join(output_dir, "results")
  results_dict["elapsed_time"] = time.time() - experiment_timer

  results.update_result_directory(results_dir, "evaluation", results_dict,