loss base class to be implemented by subclasses call contains the logic for loss calculation using ytrue ypred example subclass implementation python class meansquarederrorloss def callself ytrue ypred return ops meanops squareypred ytrue axis1 squeeze last dim if ranks differ from expected by exactly 1 x1rank lenx1 shape x2rank lenx2 shape if x1rank x2rank return x1 x2 if x1rank x2rank 1 if x1 shape1 1 x1 ops squeezex1 axis1 if x2rank x1rank 1 if x2 shape1 1 x2 ops squeezex2 axis1 return x1 x2 def reducevaluesvalues reductionsumoverbatchsize if reduction is none or reduction none or tuplevalues shape or tuplevalues shape 0 return values loss ops sumvalues if reduction sumoverbatchsize loss ops cast ops prodops converttotensorops shapevalues dtypeint32 loss dtype return loss def reduceweightedvalues values sampleweightnone masknone reductionsumoverbatchsize dtypenone reduction standardizereductionreduction values ops converttotensorvalues dtypedtype if sampleweight is not none sampleweight ops converttotensorsampleweight dtypedtype if mask is not none mask ops converttotensormask dtypedtype merge mask and sample weight into sample weight sampleweight applymask sampleweight mask dtypevalues dtype reductionreduction if sampleweight is not none sampleweight ops castsampleweight values dtype update dimensions of sampleweight to match losses values sampleweight squeezetosamerankvalues sampleweight values values sampleweight apply reduction function to the individual weighted losses loss reducevaluesvalues reduction return loss def applymasksampleweight mask dtype reduction valid entries have weight totalvalid while invalid ones have 0 when summed over batch they will be reduced to meanloss sampleweight total valid sumloss sampleweight total valid total sumloss sampleweight total total valid sumloss sampleweight valid loss base class to be implemented by subclasses call contains the logic for loss calculation using y_true y_pred example subclass implementation python class meansquarederror loss def call self y_true y_pred return ops mean ops square y_pred y_true axis 1 squeeze last dim if ranks differ from expected by exactly 1 merge mask and sample weight into sample weight update dimensions of sample_weight to match losses apply reduction function to the individual weighted losses applies any mask on predictions to sample weights valid entries have weight total valid while invalid ones have 0 when summed over batch they will be reduced to mean loss sample_weight total valid sum loss sample_weight total valid total sum loss sample_weight total total valid sum loss sample_weight valid may be 0
import tree

from keras import backend
from keras import ops
from keras.api_export import keras_export
from keras.utils.naming import auto_name


@keras_export(["keras.Loss", "keras.losses.Loss"])
class Loss:
    

    def __init__(self, name=None, reduction="sum_over_batch_size", dtype=None):
        self.name = name or auto_name(self.__class__.__name__)
        self.reduction = standardize_reduction(reduction)
        self.dtype = dtype or backend.floatx()

    def __call__(self, y_true, y_pred, sample_weight=None):
        in_mask = getattr(y_pred, "_keras_mask", None)

        with ops.name_scope(self.name):
            y_pred = tree.map_structure(
                lambda x: ops.convert_to_tensor(x, dtype=self.dtype), y_pred
            )
            y_true = tree.map_structure(
                lambda x: ops.convert_to_tensor(x, dtype=self.dtype), y_true
            )

            losses = self.call(y_true, y_pred)
            out_mask = getattr(losses, "_keras_mask", None)

            if in_mask is not None and out_mask is not None:
                mask = in_mask & out_mask
            elif in_mask is not None:
                mask = in_mask
            elif out_mask is not None:
                mask = out_mask
            else:
                mask = None

            return reduce_weighted_values(
                losses,
                sample_weight=sample_weight,
                mask=mask,
                reduction=self.reduction,
                dtype=self.dtype,
            )

    def call(self, y_true, y_pred):
        raise NotImplementedError

    def get_config(self):
        return {"name": self.name, "reduction": self.reduction}

    @classmethod
    def from_config(cls, config):
        return cls(**config)


def standardize_reduction(reduction):
    allowed = {"sum_over_batch_size", "sum", None, "none"}
    if reduction not in allowed:
        raise ValueError(
            "Invalid value for argument `reduction`. "
            f"Expected on of {allowed}. Received: "
            f"reduction={reduction}"
        )
    return reduction


def squeeze_to_same_rank(x1, x2):
    
    x1_rank = len(x1.shape)
    x2_rank = len(x2.shape)
    if x1_rank == x2_rank:
        return x1, x2
    if x1_rank == x2_rank + 1:
        if x1.shape[-1] == 1:
            x1 = ops.squeeze(x1, axis=-1)
    if x2_rank == x1_rank + 1:
        if x2.shape[-1] == 1:
            x2 = ops.squeeze(x2, axis=-1)
    return x1, x2


def reduce_values(values, reduction="sum_over_batch_size"):
    if (
        reduction is None
        or reduction == "none"
        or tuple(values.shape) == ()
        or tuple(values.shape) == (0,)
    ):
        return values
    loss = ops.sum(values)
    if reduction == "sum_over_batch_size":
        loss /= ops.cast(
            ops.prod(ops.convert_to_tensor(ops.shape(values), dtype="int32")),
            loss.dtype,
        )
    return loss


def reduce_weighted_values(
    values,
    sample_weight=None,
    mask=None,
    reduction="sum_over_batch_size",
    dtype=None,
):
    reduction = standardize_reduction(reduction)

    values = ops.convert_to_tensor(values, dtype=dtype)
    if sample_weight is not None:
        sample_weight = ops.convert_to_tensor(sample_weight, dtype=dtype)
    if mask is not None:
        mask = ops.convert_to_tensor(mask, dtype=dtype)

    
    sample_weight = apply_mask(
        sample_weight, mask, dtype=values.dtype, reduction=reduction
    )

    if sample_weight is not None:
        sample_weight = ops.cast(sample_weight, values.dtype)
        
        values, sample_weight = squeeze_to_same_rank(values, sample_weight)
        values = values * sample_weight

    
    loss = reduce_values(values, reduction)
    return loss


def apply_mask(sample_weight, mask, dtype, reduction):
    
    if mask is not None:
        mask = ops.cast(mask, dtype=dtype)
        if reduction == "sum_over_batch_size":
            
            
            
            
            
            
            
            total = ops.cast(
                ops.prod(ops.convert_to_tensor(ops.shape(mask), dtype="int32")),
                dtype,
            )
            valid = ops.sum(mask)  
            mask *= total / (valid + backend.epsilon())

        if sample_weight is not None:
            sample_weight = ops.cast(sample_weight, dtype=dtype)
            mask, sample_weight = squeeze_to_same_rank(mask, sample_weight)
            sample_weight *= mask
        else:
            sample_weight = mask
    return sample_weight