# Path: disentanglement_lib/data/ground_truth/named_data.py
# def get_named_ground_truth_data(name):
# 
# 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/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/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/visualize/visualize_util.py
# def save_image(image, image_path):
# def grid_save_images(images, image_path):
# def padded_grid(images, num_rows=None, padding_px=10, value=None):
# def padded_stack(images, padding_px=10, axis=0, value=None):
# def padding_array(image, padding_px, axis, value=None):
# def best_num_rows(num_elements, max_ratio=4):
# def pad_around(image, padding_px=10, axis=None, value=None):
# def add_below(image, padding_px=10, value=None):
# def save_animation(list_of_animated_images, image_path, fps):
# def cycle_factor(starting_index, num_indices, num_frames):
# def cycle_gaussian(starting_value, num_frames, loc=0., scale=1.):
# def cycle_interval(starting_value, num_frames, min_val, max_val):


import os
import time
import numpy as np
import tensorflow as tf
import gin.tf.external_configurables  # pylint: disable=unused-import
import gin.tf
from disentanglement_lib.data.ground_truth import named_data
from disentanglement_lib.methods.unsupervised import gaussian_encoder_model
from disentanglement_lib.methods.weak import weak_vae  # pylint: disable=unused-import
from disentanglement_lib.utils import results
from disentanglement_lib.visualize import visualize_util
from tensorflow_estimator.python.estimator.tpu import tpu_config
from tensorflow_estimator.python.estimator.tpu.tpu_estimator import TPUEstimator
      raise ValueError("Directory already exists and overwrite is False.")

  # Create a numpy random state. We will sample the random seeds for training
  # and evaluation from this.
  random_state = np.random.RandomState(random_seed)

  # Obtain the dataset.
  dataset = named_data.get_named_ground_truth_data()
  # We create a TPUEstimator based on the provided model. This is primarily so
  # that we could switch to TPU training in the future. For now, we train
  # locally on GPUs.
  run_config = tpu_config.RunConfig(
      tf_random_seed=random_seed,
      keep_checkpoint_max=1,
      tpu_config=tpu_config.TPUConfig(iterations_per_loop=500))
  tpu_estimator = TPUEstimator(
      use_tpu=False,
      model_fn=model.model_fn,
      model_dir=model_dir,
      train_batch_size=batch_size,
      eval_batch_size=batch_size,
      config=run_config)

  # Set up time to keep track of elapsed time in results.
  experiment_timer = time.time()
  # Do the actual training.
  tpu_estimator.train(
      input_fn=_make_input_fn(dataset, random_state.randint(2**32)),
      steps=training_steps)
  # Save model as a TFHub module.

# Path: disentanglement_lib/data/ground_truth/named_data.py
# def get_named_ground_truth_data(name):
# 
# 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/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/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/visualize/visualize_util.py
# def save_image(image, image_path):
# def grid_save_images(images, image_path):
# def padded_grid(images, num_rows=None, padding_px=10, value=None):
# def padded_stack(images, padding_px=10, axis=0, value=None):
# def padding_array(image, padding_px, axis, value=None):
# def best_num_rows(num_elements, max_ratio=4):
# def pad_around(image, padding_px=10, axis=None, value=None):
# def add_below(image, padding_px=10, value=None):
# def save_animation(list_of_animated_images, image_path, fps):
# def cycle_factor(starting_index, num_indices, num_frames):
# def cycle_gaussian(starting_value, num_frames, loc=0., scale=1.):
# def cycle_interval(starting_value, num_frames, min_val, max_val):


# Path: disentanglement_lib/methods/weak/train_weak_lib.py
import os
import time
import numpy as np
import tensorflow as tf
import gin.tf.external_configurables  # pylint: disable=unused-import
import gin.tf
from disentanglement_lib.data.ground_truth import named_data
from disentanglement_lib.methods.unsupervised import gaussian_encoder_model
from disentanglement_lib.methods.weak import weak_vae  # pylint: disable=unused-import
from disentanglement_lib.utils import results
from disentanglement_lib.visualize import visualize_util
from tensorflow_estimator.python.estimator.tpu import tpu_config
from tensorflow_estimator.python.estimator.tpu.tpu_estimator import TPUEstimator

      raise ValueError("Directory already exists and overwrite is False.")

  # Create a numpy random state. We will sample the random seeds for training
  # and evaluation from this.
  random_state = np.random.RandomState(random_seed)

  # Obtain the dataset.
  dataset = named_data.get_named_ground_truth_data()
  # We create a TPUEstimator based on the provided model. This is primarily so
  # that we could switch to TPU training in the future. For now, we train
  # locally on GPUs.
  run_config = tpu_config.RunConfig(
      tf_random_seed=random_seed,
      keep_checkpoint_max=1,
      tpu_config=tpu_config.TPUConfig(iterations_per_loop=500))
  tpu_estimator = TPUEstimator(
      use_tpu=False,
      model_fn=model.model_fn,
      model_dir=model_dir,
      train_batch_size=batch_size,
      eval_batch_size=batch_size,
      config=run_config)

  # Set up time to keep track of elapsed time in results.
  experiment_timer = time.time()
  # Do the actual training.
  tpu_estimator.train(
      input_fn=_make_input_fn(dataset, random_state.randint(2**32)),
      steps=training_steps)
  # Save model as a TFHub module.

  output_shape = named_data.get_named_ground_truth_data().observation_shape