# Path: enn/base.py
# class OutputWithPrior(NamedTuple):
# class EpistemicModule(abc.ABC, hk.Module):
# class ApplyFn(typing_extensions.Protocol):
# class InitFn(typing_extensions.Protocol):
# class EpistemicIndexer(typing_extensions.Protocol):
# class EpistemicNetwork:
# class Batch(NamedTuple):
# class LossFn(typing_extensions.Protocol):
# class EpistemicNetworkWithState:
# class LossFnWithState(typing_extensions.Protocol):
#   def preds(self) -> Array:
#   def __call__(self, inputs: Array, index: Index) -> Output:
#   def __call__(self, params: hk.Params, inputs: Array, index: Index) -> Output:
#   def __call__(self, rng_key: RngKey, inputs: Array, index: Index) -> hk.Params:
#   def __call__(self, key: RngKey) -> Index:
#   def __call__(self,
#                enn: EpistemicNetwork,
#                params: hk.Params,
#                batch: Batch,
#                key: RngKey) -> Tuple[Array, LossMetrics]:
#   def __call__(self,
#                enn: EpistemicNetworkWithState,
#                params: hk.Params,
#                state: hk.State,
#                batch: Batch,
#                key: RngKey) -> Tuple[Array, Tuple[hk.State, LossMetrics]]:
# 
# 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:


from absl import logging
from enn import base
from enn import utils
import dataclasses
import haiku as hk
import jax
import jax.numpy as jnp
import typing_extensions
# 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.
# ============================================================================

"""Prior losses are losses that regulate towards the prior.

These might take the form of weight regularization, or sampling "fake data".
These prior_losses are used in e.g. supervised/prior_experiment.py.
"""



class FakeInputGenerator(typing_extensions.Protocol):


# Path: enn/base.py
# class OutputWithPrior(NamedTuple):
# class EpistemicModule(abc.ABC, hk.Module):
# class ApplyFn(typing_extensions.Protocol):
# class InitFn(typing_extensions.Protocol):
# class EpistemicIndexer(typing_extensions.Protocol):
# class EpistemicNetwork:
# class Batch(NamedTuple):
# class LossFn(typing_extensions.Protocol):
# class EpistemicNetworkWithState:
# class LossFnWithState(typing_extensions.Protocol):
#   def preds(self) -> Array:
#   def __call__(self, inputs: Array, index: Index) -> Output:
#   def __call__(self, params: hk.Params, inputs: Array, index: Index) -> Output:
#   def __call__(self, rng_key: RngKey, inputs: Array, index: Index) -> hk.Params:
#   def __call__(self, key: RngKey) -> Index:
#   def __call__(self,
#                enn: EpistemicNetwork,
#                params: hk.Params,
#                batch: Batch,
#                key: RngKey) -> Tuple[Array, LossMetrics]:
#   def __call__(self,
#                enn: EpistemicNetworkWithState,
#                params: hk.Params,
#                state: hk.State,
#                batch: Batch,
#                key: RngKey) -> Tuple[Array, Tuple[hk.State, LossMetrics]]:
# 
# 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/losses/prior_losses.py
from absl import logging
from enn import base
from enn import utils
import dataclasses
import haiku as hk
import jax
import jax.numpy as jnp
import typing_extensions

# 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.
# ============================================================================

"""Prior losses are losses that regulate towards the prior.

These might take the form of weight regularization, or sampling "fake data".
These prior_losses are used in e.g. supervised/prior_experiment.py.
"""



class FakeInputGenerator(typing_extensions.Protocol):


  def __call__(self, batch: base.Batch, key: base.RngKey) -> base.Array: