Skip to content

Abstract Weights

popinn.AbstractWeights ¤

Abstract base class for per-term loss weighting schemes.

Subclasses store the weighting parameters in values and implement the combine method, which folds the per-term scalar residuals into the total weighted loss.

Fields:

  • values (dict): the weights, keyed by the set of corresponding ResidualTerm.name strings. Whether the weights are static constants or trainable leaves is decided by the concrete subclass.
Source code in src/popinn/loss.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
class AbstractWeights(eqx.Module):
    """
    Abstract base class for per-term loss weighting schemes.

    Subclasses store the weighting parameters in `values` and implement the
    `combine` method, which folds the per-term scalar residuals into the total weighted loss.

    **Fields**:

    - `values` (`dict`): the weights, keyed by the set of corresponding `ResidualTerm.name` strings. Whether the
        weights are static constants or trainable leaves is decided by the concrete subclass.
    """

    values: eqx.AbstractVar[dict]

    @abc.abstractmethod
    def combine(self, residuals: dict) -> Float[Array, ""]:
        """
        Combine per-term scalar residuals into the weighted total.

        Args:
            residuals (dict): A mapping from each `ResidualTerm.name` to that term's scalar
                value.

        Returns:
            (Float[Array, '']): The total weighted loss.
        """
        raise NotImplementedError
combine(residuals) ¤

Combine per-term scalar residuals into the weighted total.

Parameters:

Name Type Description Default
residuals dict

A mapping from each ResidualTerm.name to that term's scalar value.

 required

Returns:

Type Description
jaxtyping.Float[jax.Array, '']

The total weighted loss.
Source code in src/popinn/loss.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
@abc.abstractmethod
def combine(self, residuals: dict) -> Float[Array, ""]:
    """
    Combine per-term scalar residuals into the weighted total.

    Args:
        residuals (dict): A mapping from each `ResidualTerm.name` to that term's scalar
            value.

    Returns:
        (Float[Array, '']): The total weighted loss.
    """
    raise NotImplementedError