keras/src/metrics/reduction_metrics.py/MeanMetricWrapper/update_state
class MeanMetricWrapper:    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)