# 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/priors.py
# class EnnWithAdditivePrior(base.EpistemicNetwork):
# class NetworkWithAdditivePrior(hk.Module):
#   def __init__(self,
#                enn: base.EpistemicNetwork,
#                prior_fn: PriorFn,
#                prior_scale: float = 1.):
#     def apply_fn(params: hk.Params,
#                  inputs: base.Array,
#                  index: base.Index) -> base.OutputWithPrior:
# def convert_enn_to_prior_fn(enn: base.EpistemicNetwork,
#                             dummy_input: base.Array,
#                             key: base.RngKey) -> PriorFn:
#   def prior_fn(x: base.Array, z: base.Index) -> base.Output:
# def make_null_prior(output_dim: int) -> Callable[[base.Array], base.Array]:
#   def null_prior(inputs: base.Array) -> base.Array:
# def _parse_gamma(gamma: GpGamma,
#                  num_feat: int,
#                  key: base.RngKey) -> Union[float, base.Array]:
# def make_random_feat_gp(
#     input_dim: int,
#     output_dim: int,
#     num_feat: int,
#     key: base.RngKey,
#     gamma: GpGamma = 1.,
#     scale: float = 1.,
# ) -> Callable[[base.Array], base.Array]:
#   def gp_instance(inputs: base.Array) -> base.Array:
# def get_random_mlp_with_index(
#     x_sample: base.Array,
#     z_sample: base.Array,
#     rng: chex.PRNGKey,
#     prior_output_sizes: Optional[Iterable[int]] = None,
#     prior_weight_std: float = 3,
#     prior_bias_std: float = 0.1
# ) -> PriorFn:
#   def net_fn(x: base.Array, z: base.Array):
#   def __call__(self, *args, **kwargs) -> base.Array:
# WARN = ('WARNING: prior parameters will be included as hk.Params for module.'
#         'If possible, you should use EnnWithAdditivePrior instead.')


from typing import Callable, Optional, Sequence, Tuple
from enn import base
from enn.networks import indexers
from enn.networks import priors
import chex
import haiku as hk
import jax
import jax.numpy as jnp
    b_init: Optional[hk.initializers.Initializer] = None
) -> Sequence[Callable[[base.Array], base.Array]]:
  """Factory method for creating ensemble of prior functions."""
  rng = hk.PRNGSequence(seed)
  def net_fn(x):
    layers = [
        hk.Flatten(),
        hk.nets.MLP(output_sizes, w_init=w_init, b_init=b_init)
    ]
    return hk.Sequential(layers)(x)

  transformed = hk.without_apply_rng(hk.transform(net_fn))

  prior_fns = []
  for _ in range(num_ensemble):
    params = transformed.init(next(rng), dummy_input)
    prior_fns.append(lambda x, params=params: transformed.apply(params, x))
  return prior_fns


def wrap_sequence_as_prior(
    prior_fns: Sequence[Callable[[base.Array], base.Array]],
) -> Callable[[base.Array, base.Index], base.Array]:
  return lambda x, z: jax.lax.switch(z, prior_fns, x)


def make_random_gp_ensemble_prior_fns(
    input_dim: int,
    output_dim: int,
    num_feat: int,

# 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/priors.py
# class EnnWithAdditivePrior(base.EpistemicNetwork):
# class NetworkWithAdditivePrior(hk.Module):
#   def __init__(self,
#                enn: base.EpistemicNetwork,
#                prior_fn: PriorFn,
#                prior_scale: float = 1.):
#     def apply_fn(params: hk.Params,
#                  inputs: base.Array,
#                  index: base.Index) -> base.OutputWithPrior:
# def convert_enn_to_prior_fn(enn: base.EpistemicNetwork,
#                             dummy_input: base.Array,
#                             key: base.RngKey) -> PriorFn:
#   def prior_fn(x: base.Array, z: base.Index) -> base.Output:
# def make_null_prior(output_dim: int) -> Callable[[base.Array], base.Array]:
#   def null_prior(inputs: base.Array) -> base.Array:
# def _parse_gamma(gamma: GpGamma,
#                  num_feat: int,
#                  key: base.RngKey) -> Union[float, base.Array]:
# def make_random_feat_gp(
#     input_dim: int,
#     output_dim: int,
#     num_feat: int,
#     key: base.RngKey,
#     gamma: GpGamma = 1.,
#     scale: float = 1.,
# ) -> Callable[[base.Array], base.Array]:
#   def gp_instance(inputs: base.Array) -> base.Array:
# def get_random_mlp_with_index(
#     x_sample: base.Array,
#     z_sample: base.Array,
#     rng: chex.PRNGKey,
#     prior_output_sizes: Optional[Iterable[int]] = None,
#     prior_weight_std: float = 3,
#     prior_bias_std: float = 0.1
# ) -> PriorFn:
#   def net_fn(x: base.Array, z: base.Array):
#   def __call__(self, *args, **kwargs) -> base.Array:
# WARN = ('WARNING: prior parameters will be included as hk.Params for module.'
#         'If possible, you should use EnnWithAdditivePrior instead.')


# Path: enn/networks/ensembles.py
from typing import Callable, Optional, Sequence, Tuple
from enn import base
from enn.networks import indexers
from enn.networks import priors
import chex
import haiku as hk
import jax
import jax.numpy as jnp

    b_init: Optional[hk.initializers.Initializer] = None
) -> Sequence[Callable[[base.Array], base.Array]]:
  """Factory method for creating ensemble of prior functions."""
  rng = hk.PRNGSequence(seed)
  def net_fn(x):
    layers = [
        hk.Flatten(),
        hk.nets.MLP(output_sizes, w_init=w_init, b_init=b_init)
    ]
    return hk.Sequential(layers)(x)

  transformed = hk.without_apply_rng(hk.transform(net_fn))

  prior_fns = []
  for _ in range(num_ensemble):
    params = transformed.init(next(rng), dummy_input)
    prior_fns.append(lambda x, params=params: transformed.apply(params, x))
  return prior_fns


def wrap_sequence_as_prior(
    prior_fns: Sequence[Callable[[base.Array], base.Array]],
) -> Callable[[base.Array, base.Index], base.Array]:
  return lambda x, z: jax.lax.switch(z, prior_fns, x)


def make_random_gp_ensemble_prior_fns(
    input_dim: int,
    output_dim: int,
    num_feat: int,

    gamma: priors.GpGamma,