# Path: enn/utils.py
# FLAGS = flags.FLAGS
# def epistemic_network_from_module(
#     enn_ctor: Callable[[], base.EpistemicModule],
#     indexer: base.EpistemicIndexer,
# ) -> base.EpistemicNetwork:
#   def enn_fn(inputs: base.Array, index: base.Index) -> base.Output:
# def wrap_transformed_as_enn(
#     transformed: hk.Transformed) -> base.EpistemicNetwork:
# def wrap_enn_as_enn_with_state(
#     enn: base.EpistemicNetwork) -> base.EpistemicNetworkWithState:
#   def init(key, x, index):
#   def apply(params, unused_state, x, index):
# def make_centered_enn(enn: base.EpistemicNetwork,
#                       x_train: base.Array) -> base.EpistemicNetwork:
#   def centered_apply(params: hk.Params,
#                      x: base.Array,
#                      z: base.Index) -> base.Output:
# def parse_net_output(net_out: base.Output) -> base.Array:
# def make_batch_indexer(indexer: base.EpistemicIndexer,
#                        batch_size: int) -> base.EpistemicIndexer:
#   def batch_indexer(key: base.RngKey) -> base.Index:
# def _clean_batch_data(data: base.Batch) -> base.Batch:
# def make_batch_iterator(data: base.Batch,
#                         batch_size: Optional[int] = None,
#                         seed: int = 0) -> base.BatchIterator:
# def make_test_data(n_samples: int = 20) -> base.BatchIterator:
# 
# Path: enn/extra/vae.py
# class MeanLogVariance(NamedTuple):
# class TrainedVAE(NamedTuple):
# class MLPVAEConfig:
# def get_mlp_vae_encoder_decoder(
#     data_x: chex.Array,
#     config: MLPVAEConfig = MLPVAEConfig(),
# ) -> TrainedVAE:
#   def mlp_encoder(x: chex.Array) -> MeanLogVariance:
#   def mlp_decoder(x: chex.Array) -> MeanLogVariance:
# def make_vae_enn(encoder: PreTransformFn,
#                  decoder: PreTransformFn,
#                  latent_dim: int) -> base.EpistemicNetwork:
#   def net_fn(x: base.Array, z: base.Index) -> base.OutputWithPrior:
# def train_vae(encoder: PreTransformFn,
#               decoder: PreTransformFn,
#               latent_dim: int,
#               data_x: chex.Array,
#               log_likelihood_fn: losses.LogLikelihoodFn,
#               optimizer: optax.GradientTransformation,
#               num_batches: int = 10_000,
#               batch_size: int = 1_000) -> TrainedVAE:
#   def encoder_fn(x: chex.Array) -> MeanLogVariance:
#   def decoder_fn(x: chex.Array) -> MeanLogVariance:


from absl.testing import absltest
from absl.testing import parameterized
from enn import utils
from enn.extra import vae
import chex
# python3
# pylint: disable=g-bad-file-header
# Copyright 2021 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# 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."""



class VaeTest(parameterized.TestCase):

  @parameterized.product(bernoulli_decoder=[True, False], latent_dim=[1, 3])
  def test_vae_outputs(self, bernoulli_decoder: bool, latent_dim: int):
    """Train a VAE on a test dataset and test encoder decoder functions."""
    dataset = utils.make_test_data(10)
    data = next(dataset)
    num_train, input_dim = data.x.shape


# Path: enn/utils.py
# FLAGS = flags.FLAGS
# def epistemic_network_from_module(
#     enn_ctor: Callable[[], base.EpistemicModule],
#     indexer: base.EpistemicIndexer,
# ) -> base.EpistemicNetwork:
#   def enn_fn(inputs: base.Array, index: base.Index) -> base.Output:
# def wrap_transformed_as_enn(
#     transformed: hk.Transformed) -> base.EpistemicNetwork:
# def wrap_enn_as_enn_with_state(
#     enn: base.EpistemicNetwork) -> base.EpistemicNetworkWithState:
#   def init(key, x, index):
#   def apply(params, unused_state, x, index):
# def make_centered_enn(enn: base.EpistemicNetwork,
#                       x_train: base.Array) -> base.EpistemicNetwork:
#   def centered_apply(params: hk.Params,
#                      x: base.Array,
#                      z: base.Index) -> base.Output:
# def parse_net_output(net_out: base.Output) -> base.Array:
# def make_batch_indexer(indexer: base.EpistemicIndexer,
#                        batch_size: int) -> base.EpistemicIndexer:
#   def batch_indexer(key: base.RngKey) -> base.Index:
# def _clean_batch_data(data: base.Batch) -> base.Batch:
# def make_batch_iterator(data: base.Batch,
#                         batch_size: Optional[int] = None,
#                         seed: int = 0) -> base.BatchIterator:
# def make_test_data(n_samples: int = 20) -> base.BatchIterator:
# 
# Path: enn/extra/vae.py
# class MeanLogVariance(NamedTuple):
# class TrainedVAE(NamedTuple):
# class MLPVAEConfig:
# def get_mlp_vae_encoder_decoder(
#     data_x: chex.Array,
#     config: MLPVAEConfig = MLPVAEConfig(),
# ) -> TrainedVAE:
#   def mlp_encoder(x: chex.Array) -> MeanLogVariance:
#   def mlp_decoder(x: chex.Array) -> MeanLogVariance:
# def make_vae_enn(encoder: PreTransformFn,
#                  decoder: PreTransformFn,
#                  latent_dim: int) -> base.EpistemicNetwork:
#   def net_fn(x: base.Array, z: base.Index) -> base.OutputWithPrior:
# def train_vae(encoder: PreTransformFn,
#               decoder: PreTransformFn,
#               latent_dim: int,
#               data_x: chex.Array,
#               log_likelihood_fn: losses.LogLikelihoodFn,
#               optimizer: optax.GradientTransformation,
#               num_batches: int = 10_000,
#               batch_size: int = 1_000) -> TrainedVAE:
#   def encoder_fn(x: chex.Array) -> MeanLogVariance:
#   def decoder_fn(x: chex.Array) -> MeanLogVariance:


# Path: enn/extra/vae_test.py
from absl.testing import absltest
from absl.testing import parameterized
from enn import utils
from enn.extra import vae
import chex

# python3
# pylint: disable=g-bad-file-header
# Copyright 2021 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# 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."""



class VaeTest(parameterized.TestCase):

  @parameterized.product(bernoulli_decoder=[True, False], latent_dim=[1, 3])
  def test_vae_outputs(self, bernoulli_decoder: bool, latent_dim: int):
    """Train a VAE on a test dataset and test encoder decoder functions."""
    dataset = utils.make_test_data(10)
    data = next(dataset)
    num_train, input_dim = data.x.shape


    config = vae.MLPVAEConfig(hidden_sizes=[5, 2],