from keras.src import backend

from keras.src import initializers

from keras.src import losses

from keras.src import ops

from keras.src.api_export import keras_export

from keras.src.metrics.metric import Metric

from keras.src.saving import serialization_lib

def reduce_to_samplewise_values(values, sample_weight, reduce_fn, dtype):

dtype = dtype or backend.floatx()

mask = getattr(values, "_keras_mask", None)

values = ops.cast(values, dtype=dtype)

if sample_weight is not None:

sample_weight = ops.convert_to_tensor(sample_weight, dtype=dtype)

if mask is not None:

sample_weight = losses.loss.apply_mask(

sample_weight, mask, dtype=dtype, reduction="sum"

)

# Update dimensions of weights to match with values if possible.

values, sample_weight = losses.loss.squeeze_or_expand_to_same_rank(

values, sample_weight

)

# Reduce values to same ndim as weight array

weight_ndim = len(sample_weight.shape)

values_ndim = len(values.shape)

if values_ndim > weight_ndim:

values = reduce_fn(

values, axis=list(range(weight_ndim, values_ndim))

)

values = values * sample_weight

if values_ndim > 1:

sample_weight = reduce_fn(

sample_weight, axis=list(range(1, weight_ndim))

)

values_ndim = len(values.shape)

if values_ndim > 1:

values = reduce_fn(values, axis=list(range(1, values_ndim)))

return values, sample_weight

return values, sample_weight

@keras_export("keras.metrics.Sum")

class Sum(Metric):

"""Compute the (weighted) sum of the given values.

For example, if `values` is `[1, 3, 5, 7]` then their sum is 16.

If `sample_weight` was specified as `[1, 1, 0, 0]` then the sum would be 4.

This metric creates one variable, `total`.

This is ultimately returned as the sum value.

Args:

name: (Optional) string name of the metric instance.

dtype: (Optional) data type of the metric result.

Example:

>>> m = metrics.Sum()

>>> m.update_state([1, 3, 5, 7])

>>> m.result()

16.0

>>> m = metrics.Sum()

>>> m.update_state([1, 3, 5, 7], sample_weight=[1, 1, 0, 0])

>>> m.result()

4.0

"""

def __init__(self, name="sum", dtype=None):

super().__init__(name=name, dtype=dtype)

self.total = self.add_variable(

shape=(),

initializer=initializers.Zeros(),

dtype=self.dtype,

name="total",

)

def update_state(self, values, sample_weight=None):

values, _ = reduce_to_samplewise_values(

values, sample_weight, reduce_fn=ops.sum, dtype=self.dtype

)

self.total.assign(self.total + ops.sum(values))

def reset_state(self):

self.total.assign(0.0)

def result(self):

return ops.cast(self.total, self.dtype)

@keras_export("keras.metrics.Mean")

class Mean(Metric):

"""Compute the (weighted) mean of the given values.

For example, if values is `[1, 3, 5, 7]` then the mean is 4.

If `sample_weight` was specified as `[1, 1, 0, 0]` then the mean would be 2.

This metric creates two variables, `total` and `count`.

The mean value returned is simply `total` divided by `count`.

Args:

name: (Optional) string name of the metric instance.

dtype: (Optional) data type of the metric result.

Example:

>>> m = Mean()

>>> m.update_state([1, 3, 5, 7])

>>> m.result()

4.0

>>> m.reset_state()

>>> m.update_state([1, 3, 5, 7], sample_weight=[1, 1, 0, 0])

>>> m.result()

2.0

```

"""

def __init__(self, name="mean", dtype=None):

super().__init__(name=name, dtype=dtype)

self.total = self.add_variable(

shape=(),

initializer=initializers.Zeros(),

dtype=self.dtype,

name="total",

)

self.count = self.add_variable(

shape=(),

initializer=initializers.Zeros(),

dtype=self.dtype,

name="count",

)

def update_state(self, values, sample_weight=None):

values, sample_weight = reduce_to_samplewise_values(

values, sample_weight, reduce_fn=ops.mean, dtype=self.dtype

)

self.total.assign(self.total + ops.sum(values))

if len(values.shape) >= 1:

num_samples = ops.shape(values)[0]

else:

num_samples = 1

if sample_weight is not None:

num_samples = ops.sum(sample_weight)

self.count.assign(self.count + ops.cast(num_samples, dtype=self.dtype))

def reset_state(self):

self.total.assign(0.0)

self.count.assign(0)

def result(self):

return ops.divide_no_nan(

self.total, ops.cast(self.count, dtype=self.dtype)

)

@keras_export("keras.metrics.MeanMetricWrapper")

class MeanMetricWrapper(Mean):

"""Wrap a stateless metric function with the `Mean` metric.

You could use this class to quickly build a mean metric from a function. The

function needs to have the signature `fn(y_true, y_pred)` and return a

per-sample loss array. `MeanMetricWrapper.result()` will return

the average metric value across all samples seen so far.

For example:

```python

def mse(y_true, y_pred):

return (y_true - y_pred) ** 2

mse_metric = MeanMetricWrapper(fn=mse)

```

Args:

fn: The metric function to wrap, with signature

`fn(y_true, y_pred, **kwargs)`.

name: (Optional) string name of the metric instance.

dtype: (Optional) data type of the metric result.

**kwargs: Keyword arguments to pass on to `fn`.

"""

def __init__(self, fn, name=None, dtype=None, **kwargs):

super().__init__(name=name, dtype=dtype)

self._fn = fn

self._fn_kwargs = kwargs

# If we are wrapping a Keras loss, register the metric's

# direction as "down" (needs to be minimized during training).

if (

self._fn in losses.ALL_OBJECTS

or hasattr(self._fn, "__class__")

and self._fn.__class__ in losses.ALL_OBJECTS

):

self._direction = "down"

def update_state(self, y_true, y_pred, sample_weight=None):

mask = getattr(y_pred, "_keras_mask", None)

values = self._fn(y_true, y_pred, **self._fn_kwargs)

if sample_weight is not None and mask is not None:

sample_weight = losses.loss.apply_mask(

sample_weight, mask, dtype=self.dtype, reduction="sum"

)

return super().update_state(values, sample_weight=sample_weight)

def get_config(self):

base_config = super().get_config()

config = {"fn": serialization_lib.serialize_keras_object(self._fn)}

config.update(serialization_lib.serialize_keras_object(self._fn_kwargs))

return {**base_config, **config}

@classmethod

def from_config(cls, config):

if "fn" in config:

config = serialization_lib.deserialize_keras_object(config)

return cls(**config)