# 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 Callable, Sequence
from enn import base
from enn.networks import indexers
import haiku as hk
import haiku.experimental as hke
import jax
import jax.numpy as jnp
          and isinstance(context.module, hk.Linear))


def make_enn_creator(init_scale: float = 1.):
  """Make enn_creator initializing c unit to init_scale."""
  custom_init = lambda shape, dtype: init_scale * jnp.ones(shape, dtype)
  def enn_creator(next_creator, shape, dtype, init, context):
    """Create gaussian enn linear layer."""
    # TODO(author2): How to import hk._src.base types correctly?
    if _is_linear_bias(context):  # Append gaussian bias term
      standard_bias = next_creator(shape, dtype, init)
      gaussian_bias = next_creator(shape, dtype, custom_init)
      return standard_bias, gaussian_bias
    else:  # Return the usual creator
      return next_creator(shape, dtype, init)
  return enn_creator


def enn_getter(next_getter, value, context):
  """Get variables for gaussian enn linear layer."""
  # TODO(author2): How to import hk._src.base types correctly?
  if _is_linear_bias(context):
    standard_bias = next_getter(value[0])
    gaussian_bias = next_getter(value[1])
    noise = jax.random.normal(hk.next_rng_key(), standard_bias.shape)
    return standard_bias + gaussian_bias * noise
  else:
    return next_getter(value)



# 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/gaussian_enn.py
from typing import Callable, Sequence
from enn import base
from enn.networks import indexers
import haiku as hk
import haiku.experimental as hke
import jax
import jax.numpy as jnp

          and isinstance(context.module, hk.Linear))


def make_enn_creator(init_scale: float = 1.):
  """Make enn_creator initializing c unit to init_scale."""
  custom_init = lambda shape, dtype: init_scale * jnp.ones(shape, dtype)
  def enn_creator(next_creator, shape, dtype, init, context):
    """Create gaussian enn linear layer."""
    # TODO(author2): How to import hk._src.base types correctly?
    if _is_linear_bias(context):  # Append gaussian bias term
      standard_bias = next_creator(shape, dtype, init)
      gaussian_bias = next_creator(shape, dtype, custom_init)
      return standard_bias, gaussian_bias
    else:  # Return the usual creator
      return next_creator(shape, dtype, init)
  return enn_creator


def enn_getter(next_getter, value, context):
  """Get variables for gaussian enn linear layer."""
  # TODO(author2): How to import hk._src.base types correctly?
  if _is_linear_bias(context):
    standard_bias = next_getter(value[0])
    gaussian_bias = next_getter(value[1])
    noise = jax.random.normal(hk.next_rng_key(), standard_bias.shape)
    return standard_bias + gaussian_bias * noise
  else:
    return next_getter(value)



class GaussianNoiseEnn(base.EpistemicNetwork):