# Path: fedjax/core/dataclasses.py
# def dataclass(clz: type):
#   def replace(self, **updates):
#   def iterate_clz(x):
#   def clz_from_iterable(meta, data):
# 
# Path: fedjax/core/util.py
# def safe_div(a, b):
#   def __getattr__(self, _):
# def import_tf():
# class FakeTensorFlowModule:
# 
# Path: fedjax/core/typing.py

# 
# Path: fedjax/core/typing.py

# 
# Path: fedjax/core/typing.py

# 
# Path: fedjax/core/typing.py



import abc
import functools
import jax
import jax.numpy as jnp
from typing import Optional, Tuple
from fedjax.core import dataclasses
from fedjax.core import util
from fedjax.core.typing import BatchExample
from fedjax.core.typing import BatchPrediction
from fedjax.core.typing import SingleExample
from fedjax.core.typing import SinglePrediction
      This is typically used to ignore padding values in computation.
  """
  target_key: str = 'y'
  pred_key: Optional[str] = None
  masked_target_values: Tuple[int, ...] = (0,)

  def zero(self) -> MeanStat:
    return MeanStat.new(0., 0.)

  def evaluate_example(self, example: SingleExample,
                       prediction: SinglePrediction) -> MeanStat:
    """Computes total cross entropy loss for a single sequence example.

    Args:
      example: One example with target in range [0, num_classes) of shape [1].
      prediction: Unnormalized prediction for ``example`` of shape [num_classes]

    Returns:
      MeanStat for total loss for a single sequence example.
    """
    target = example[self.target_key]
    pred = prediction if self.pred_key is None else prediction[self.pred_key]
    target_weight = get_target_weight(target, self.masked_target_values)
    token_loss = unreduced_cross_entropy_loss(target, pred)
    # Change weight from number of non masked target tokens to 1 if the sequence
    # contains any non masked tokens or 0 if the entire sequence is masked.
    return MeanStat.new(
        jnp.sum(token_loss * target_weight), jnp.any(target_weight))



# Path: fedjax/core/dataclasses.py
# def dataclass(clz: type):
#   def replace(self, **updates):
#   def iterate_clz(x):
#   def clz_from_iterable(meta, data):
# 
# Path: fedjax/core/util.py
# def safe_div(a, b):
#   def __getattr__(self, _):
# def import_tf():
# class FakeTensorFlowModule:
# 
# Path: fedjax/core/typing.py

# 
# Path: fedjax/core/typing.py

# 
# Path: fedjax/core/typing.py

# 
# Path: fedjax/core/typing.py



# Path: fedjax/core/metrics.py
import abc
import functools
import jax
import jax.numpy as jnp
from typing import Optional, Tuple
from fedjax.core import dataclasses
from fedjax.core import util
from fedjax.core.typing import BatchExample
from fedjax.core.typing import BatchPrediction
from fedjax.core.typing import SingleExample
from fedjax.core.typing import SinglePrediction

      This is typically used to ignore padding values in computation.
  """
  target_key: str = 'y'
  pred_key: Optional[str] = None
  masked_target_values: Tuple[int, ...] = (0,)

  def zero(self) -> MeanStat:
    return MeanStat.new(0., 0.)

  def evaluate_example(self, example: SingleExample,
                       prediction: SinglePrediction) -> MeanStat:
    """Computes total cross entropy loss for a single sequence example.

    Args:
      example: One example with target in range [0, num_classes) of shape [1].
      prediction: Unnormalized prediction for ``example`` of shape [num_classes]

    Returns:
      MeanStat for total loss for a single sequence example.
    """
    target = example[self.target_key]
    pred = prediction if self.pred_key is None else prediction[self.pred_key]
    target_weight = get_target_weight(target, self.masked_target_values)
    token_loss = unreduced_cross_entropy_loss(target, pred)
    # Change weight from number of non masked target tokens to 1 if the sequence
    # contains any non masked tokens or 0 if the entire sequence is masked.
    return MeanStat.new(
        jnp.sum(token_loss * target_weight), jnp.any(target_weight))



@dataclasses.dataclass