# 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/losses/single_index.py
# class SingleIndexLossFn(typing_extensions.Protocol):
# class L2Loss(SingleIndexLossFn):
# class XentLoss(SingleIndexLossFn):
# class AccuracyErrorLoss(SingleIndexLossFn):
# class ElboLoss(SingleIndexLossFn):
# class VaeLoss(SingleIndexLossFn):
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
# def average_single_index_loss(single_loss: SingleIndexLossFn,
#                               num_index_samples: int = 1) -> base.LossFn:
#   def loss_fn(enn: base.EpistemicNetwork,
#               params: hk.Params,
#               batch: base.Batch,
#               key: base.RngKey) -> Tuple[base.Array, base.LossMetrics]:
# def add_data_noise(single_loss: SingleIndexLossFn,
#                    noise_fn: data_noise.DataNoise) -> SingleIndexLossFn:
#   def noisy_loss(apply: base.ApplyFn,
#                  params: hk.Params,
#                  batch: base.Batch,
#                  index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __post_init__(self):
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:


import dataclasses
import haiku as hk
import jax
import jax.numpy as jnp
from typing import Any, Dict, Callable, List, Optional, Sequence, Tuple, Union
from enn import base
from enn.losses import single_index
# 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.
# ============================================================================

"""Helpful functions relating to losses."""



# Maps Haiku params (module_name, name, value) -> include or not
PredicateFn = Callable[[str, str, Any], bool]


def l2_weights_with_predicate(
    params: hk.Params,
    predicate: Optional[PredicateFn] = None) -> jnp.DeviceArray:
  """Sum of squares of parameter weights that passes predicate_fn."""
  if predicate is not None:
    params = hk.data_structures.filter(predicate, params)
  return sum(jnp.sum(jnp.square(p)) for p in jax.tree_leaves(params))



# 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/losses/single_index.py
# class SingleIndexLossFn(typing_extensions.Protocol):
# class L2Loss(SingleIndexLossFn):
# class XentLoss(SingleIndexLossFn):
# class AccuracyErrorLoss(SingleIndexLossFn):
# class ElboLoss(SingleIndexLossFn):
# class VaeLoss(SingleIndexLossFn):
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
# def average_single_index_loss(single_loss: SingleIndexLossFn,
#                               num_index_samples: int = 1) -> base.LossFn:
#   def loss_fn(enn: base.EpistemicNetwork,
#               params: hk.Params,
#               batch: base.Batch,
#               key: base.RngKey) -> Tuple[base.Array, base.LossMetrics]:
# def add_data_noise(single_loss: SingleIndexLossFn,
#                    noise_fn: data_noise.DataNoise) -> SingleIndexLossFn:
#   def noisy_loss(apply: base.ApplyFn,
#                  params: hk.Params,
#                  batch: base.Batch,
#                  index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __post_init__(self):
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:


# Path: enn/losses/utils.py
import dataclasses
import haiku as hk
import jax
import jax.numpy as jnp
from typing import Any, Dict, Callable, List, Optional, Sequence, Tuple, Union
from enn import base
from enn.losses import single_index

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

"""Helpful functions relating to losses."""



# Maps Haiku params (module_name, name, value) -> include or not
PredicateFn = Callable[[str, str, Any], bool]


def l2_weights_with_predicate(
    params: hk.Params,
    predicate: Optional[PredicateFn] = None) -> jnp.DeviceArray:
  """Sum of squares of parameter weights that passes predicate_fn."""
  if predicate is not None:
    params = hk.data_structures.filter(predicate, params)
  return sum(jnp.sum(jnp.square(p)) for p in jax.tree_leaves(params))



def add_l2_weight_decay(loss_fn: base.LossFn,