# 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 typing import Callable, Tuple
from enn import base
from enn import utils
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import typing_extensions
      batch: base.Batch,
      index: base.Index,
  ) -> Tuple[base.Array, Tuple[hk.State, base.LossMetrics]]:
    """Computes a loss based on one batch of data and one index."""


def average_single_index_loss_with_state(
    single_loss: SingleIndexLossFnWithState,
    num_index_samples: int = 1,
) -> base.LossFnWithState:
  """Average a single index loss over multiple index samples.

  Note that the *network state* is also averaged over indices. This is not going
  to be equivalent to num_index_samples updates sequentially. We may want to
  think about alternative ways to do this, or set num_index_samples=1.

  Args:
    single_loss: loss function applied per epistemic index.
    num_index_samples: number of index samples to average.

  Returns:
    LossFnWithState that comprises the mean of both the loss and the metrics.
  """

  def loss_fn(
      enn: base.EpistemicNetworkWithState,
      params: hk.Params,
      state: hk.State,
      batch: base.Batch,
      key: base.RngKey) -> Tuple[base.Array, Tuple[hk.State, base.LossMetrics]]:

# 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/single_index_with_state.py
from typing import Callable, Tuple
from enn import base
from enn import utils
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import typing_extensions

      batch: base.Batch,
      index: base.Index,
  ) -> Tuple[base.Array, Tuple[hk.State, base.LossMetrics]]:
    """Computes a loss based on one batch of data and one index."""


def average_single_index_loss_with_state(
    single_loss: SingleIndexLossFnWithState,
    num_index_samples: int = 1,
) -> base.LossFnWithState:
  """Average a single index loss over multiple index samples.

  Note that the *network state* is also averaged over indices. This is not going
  to be equivalent to num_index_samples updates sequentially. We may want to
  think about alternative ways to do this, or set num_index_samples=1.

  Args:
    single_loss: loss function applied per epistemic index.
    num_index_samples: number of index samples to average.

  Returns:
    LossFnWithState that comprises the mean of both the loss and the metrics.
  """

  def loss_fn(
      enn: base.EpistemicNetworkWithState,
      params: hk.Params,
      state: hk.State,
      batch: base.Batch,
      key: base.RngKey) -> Tuple[base.Array, Tuple[hk.State, base.LossMetrics]]:

    batched_indexer = utils.make_batch_indexer(enn.indexer, num_index_samples)