# Path: disentanglement_lib/data/ground_truth/named_data.py
# def get_named_ground_truth_data(name):
# 
# Path: disentanglement_lib/data/ground_truth/util.py
# def tf_data_set_from_ground_truth_data(ground_truth_data, random_seed):
#   def generator():
#   def __init__(self, factor_sizes, latent_factor_indices):
#   def num_latent_factors(self):
#   def sample_latent_factors(self, num, random_state):
#   def sample_all_factors(self, latent_factors, random_state):
#   def _sample_factor(self, i, num, random_state):
#   def __init__(self, factor_sizes, features):
#   def features_to_index(self, features):
#   def _features_to_state_space_index(self, features):
# class SplitDiscreteStateSpace(object):
# class StateSpaceAtomIndex(object):
# 
# 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/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/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 os
import time
import numpy as np
import tensorflow.compat.v1 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.data.ground_truth import util
from disentanglement_lib.methods.unsupervised import gaussian_encoder_model
from disentanglement_lib.methods.unsupervised import vae  # pylint: disable=unused-import
from disentanglement_lib.utils import results
from tensorflow.contrib import tpu as contrib_tpu
  # 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 = contrib_tpu.RunConfig(
      tf_random_seed=random_seed,
      keep_checkpoint_max=1,
      tpu_config=contrib_tpu.TPUConfig(iterations_per_loop=500))
  tpu_estimator = contrib_tpu.TPUEstimator(
      use_tpu=False,
      model_fn=model.model_fn,
      model_dir=os.path.join(model_dir, "tf_checkpoint"),
      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/data/ground_truth/util.py
# def tf_data_set_from_ground_truth_data(ground_truth_data, random_seed):
#   def generator():
#   def __init__(self, factor_sizes, latent_factor_indices):
#   def num_latent_factors(self):
#   def sample_latent_factors(self, num, random_state):
#   def sample_all_factors(self, latent_factors, random_state):
#   def _sample_factor(self, i, num, random_state):
#   def __init__(self, factor_sizes, features):
#   def features_to_index(self, features):
#   def _features_to_state_space_index(self, features):
# class SplitDiscreteStateSpace(object):
# class StateSpaceAtomIndex(object):
# 
# 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/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/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/methods/unsupervised/train.py
import os
import time
import numpy as np
import tensorflow.compat.v1 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.data.ground_truth import util
from disentanglement_lib.methods.unsupervised import gaussian_encoder_model
from disentanglement_lib.methods.unsupervised import vae  # pylint: disable=unused-import
from disentanglement_lib.utils import results
from tensorflow.contrib import tpu as contrib_tpu

  # 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 = contrib_tpu.RunConfig(
      tf_random_seed=random_seed,
      keep_checkpoint_max=1,
      tpu_config=contrib_tpu.TPUConfig(iterations_per_loop=500))
  tpu_estimator = contrib_tpu.TPUEstimator(
      use_tpu=False,
      model_fn=model.model_fn,
      model_dir=os.path.join(model_dir, "tf_checkpoint"),
      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