# 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/networks/indexers.py
# class PrngIndexer(base.EpistemicIndexer):
# class EnsembleIndexer(base.EpistemicIndexer):
# class GaussianIndexer(base.EpistemicIndexer):
# class ScaledGaussianIndexer(base.EpistemicIndexer):
# class GaussianWithUnitIndexer(base.EpistemicIndexer):
# class DirichletIndexer(base.EpistemicIndexer):
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def mean_index(self) -> base.Array:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:


from typing import Optional, Sequence, List
from enn import base
from enn.networks import indexers
import haiku as hk
import jax
import jax.numpy as jnp
      # Note that we consider a dropout layer after the input to be
      # consistent with the paper "Dropout as a Bayesian Approximation:
      # Representing Model Uncertainty in Deep Learning" (2015),
      # https://github.com/yaringal/DropoutUncertaintyExps/blob/master/net/net.py
      if dropout_input:
        mask = masks[0]
        x *= mask

      for layer_index, output_size in enumerate(output_sizes):
        if layer_index == 0:
          if nonzero_bias:
            b_init_0 = hk.initializers.TruncatedNormal(
                stddev=(1. / jnp.sqrt(input_size)))
          else:
            b_init_0 = b_init
          x = hk.Linear(output_size, w_init=w_init, b_init=b_init_0)(x)
        else:
          x = hk.Linear(output_size, w_init=w_init, b_init=b_init)(x)
        if layer_index < len(output_sizes) - 1:
          mask = masks[layer_index + 1]
          x *= mask
          x = jax.nn.relu(x)
      return x

    # Note that our enn_fn is stochastic because we generate masks in it. But,
    # since we pass a base.RngKey directly to it, we can still wrap transformed
    # function with hk.without_apply_rng.
    transformed = hk.without_apply_rng(hk.transform(enn_fn))

    # We use a simple indexer which is basically an identity map.

# 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/networks/indexers.py
# class PrngIndexer(base.EpistemicIndexer):
# class EnsembleIndexer(base.EpistemicIndexer):
# class GaussianIndexer(base.EpistemicIndexer):
# class ScaledGaussianIndexer(base.EpistemicIndexer):
# class GaussianWithUnitIndexer(base.EpistemicIndexer):
# class DirichletIndexer(base.EpistemicIndexer):
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def mean_index(self) -> base.Array:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:


# Path: enn/networks/dropout.py
from typing import Optional, Sequence, List
from enn import base
from enn.networks import indexers
import haiku as hk
import jax
import jax.numpy as jnp

      # Note that we consider a dropout layer after the input to be
      # consistent with the paper "Dropout as a Bayesian Approximation:
      # Representing Model Uncertainty in Deep Learning" (2015),
      # https://github.com/yaringal/DropoutUncertaintyExps/blob/master/net/net.py
      if dropout_input:
        mask = masks[0]
        x *= mask

      for layer_index, output_size in enumerate(output_sizes):
        if layer_index == 0:
          if nonzero_bias:
            b_init_0 = hk.initializers.TruncatedNormal(
                stddev=(1. / jnp.sqrt(input_size)))
          else:
            b_init_0 = b_init
          x = hk.Linear(output_size, w_init=w_init, b_init=b_init_0)(x)
        else:
          x = hk.Linear(output_size, w_init=w_init, b_init=b_init)(x)
        if layer_index < len(output_sizes) - 1:
          mask = masks[layer_index + 1]
          x *= mask
          x = jax.nn.relu(x)
      return x

    # Note that our enn_fn is stochastic because we generate masks in it. But,
    # since we pass a base.RngKey directly to it, we can still wrap transformed
    # function with hk.without_apply_rng.
    transformed = hk.without_apply_rng(hk.transform(enn_fn))

    # We use a simple indexer which is basically an identity map.

    indexer = indexers.PrngIndexer()