Loss Base

base

CLASS	DESCRIPTION
LossType	Categories for loss functions.
Loss	Base class for loss functions with common parameter validation.

Classes

LossType

Bases: str, Enum

Categories for loss functions.

Used to validate loss compatibility with specific models and provide semantic grouping of loss functions.

ATTRIBUTE	DESCRIPTION
SPECTRAL	Eigenvalue/eigenvector-based losses (e.g., eigenvalue gap, alignment). TYPE: str
RECONSTRUCTION	Reconstruction error losses (e.g., MSE, MAE, Huber). TYPE: str
PROBABILISTIC	Probability distribution losses (e.g., KL divergence, ELBO, entropy). TYPE: str
CONTRASTIVE	Contrastive and metric learning losses (e.g., triplet, contrastive). TYPE: str
REGULARIZATION	Regularization penalties (e.g., orthogonality, Laplacian smoothness). TYPE: str
METRIC	Evaluation metrics adapted as losses (e.g., R2, cosine similarity). TYPE: str
DISTILLATION	Distillation losses (e.g., knowledge distillation, attention distillation). TYPE: str
CLUSTERING	Clustering losses (e.g., min cut). TYPE: str

Loss(reduce: Literal['sum', 'mean'] = 'mean', sign: float = 1.0, context_prefix: str | None = None, *args, **kwargs)

Bases: Module

Base class for loss functions with common parameter validation.

All loss functions inherit from this class and must implement forward() method that accepts a context dictionary containing tensors required for loss computation.

Child classes should override allowed_reduce and loss_types at class level to define custom reduction operations and semantic types.

PARAMETER	DESCRIPTION
reduce	Reduction method. TYPE: str, optional, by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. Use -1.0 for maximization, 1.0 for minimization. TYPE: float, optional, by default 1.0 DEFAULT: 1.0

ATTRIBUTE	DESCRIPTION
allowed_reduce	Allowed reduction operations. Child classes can override this class attribute to support different reduction methods. TYPE: list[str], by default ["sum", "mean"]
loss_types	Loss function categories. Child classes should override this class attribute to define their semantic type(s). TYPE: list[LossType], by default []

Examples:

Define a custom loss with specific reduction options

>>> class MyLoss(Loss):
...     allowed_reduce = ["sum", "mean"]
...     loss_types = [LossType.RECONSTRUCTION]
...
...     def __init__(self, reduce="mean", sign=1.0):
...         super().__init__(reduce=reduce, sign=sign)
...
...     def forward(self, context, **kwargs):
...         X = context.get("X")
...         loss = (X**2).sum()
...         if self.reduce == "mean":
...             loss = loss.mean()
...         return self.sign * loss

METHOD	DESCRIPTION
exists	Check if loss matches the given pattern.
forward	Compute loss from context dictionary.

Source code in spectre/loss/base.py

def __init__(
    self,
    reduce: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
    context_prefix: str | None = None,
    *args,
    **kwargs,
):
    super().__init__(*args, **kwargs)

    if not isinstance(reduce, str):
        raise TypeError(f"Expected reduce to be str, got {type(reduce).__name__}.")
    if reduce not in self.allowed_reduce:
        raise ValueError(
            f"Unknown reduce operation '{reduce}'. Options: {self.allowed_reduce}."
        )
    self.reduce = reduce

    if not isinstance(sign, (int, float)):
        raise TypeError(
            f"Expected sign to be int or float, got {type(sign).__name__}."
        )
    self.sign = sign

    if context_prefix is not None and not isinstance(context_prefix, str):
        raise TypeError(
            f"Expected context_prefix to be str or None, got {type(context_prefix).__name__}."
        )
    self.context_prefix = context_prefix

Functions

exists(pattern: str) -> bool

Check if loss matches the given pattern.

For base Loss, checks if the pattern is contained in the loss name. CompositeLoss overrides this to check all constituent losses.

PARAMETER	DESCRIPTION
pattern	Pattern to match against loss name (e.g., "mincut", "eigenvalue"). TYPE: str

RETURNS	DESCRIPTION
bool	True if pattern is contained in the loss name, False otherwise.

Examples:

>>> loss = MinCutLoss()
>>> loss.exists("mincut")
True
>>> loss.exists("loss_mincut")
True
>>> loss.exists("mse")
False

Source code in spectre/loss/base.py

def exists(self, pattern: str) -> bool:
    """
    Check if loss matches the given pattern.

    For base Loss, checks if the pattern is contained in the loss name.
    CompositeLoss overrides this to check all constituent losses.

    Parameters
    ----------
    pattern : str
        Pattern to match against loss name (e.g., "mincut", "eigenvalue").

    Returns
    -------
    bool
        True if pattern is contained in the loss name, False otherwise.

    Examples
    --------
    >>> loss = MinCutLoss()
    >>> loss.exists("mincut")
    True
    >>> loss.exists("loss_mincut")
    True
    >>> loss.exists("mse")
    False
    """
    regex = r"[A-Z]+(?=[A-Z][a-z]|\b)|[A-Z][a-z]+"
    name_parts = re.findall(regex, self.__class__.__name__)[:-1]
    loss_name = f"loss_{'_'.join([n.lower() for n in name_parts])}"

    return pattern in loss_name

forward(context: dict[str, torch.Tensor], **kwargs) -> torch.Tensor abstractmethod

Compute loss from context dictionary.

PARAMETER	DESCRIPTION
context	Dictionary containing all tensors required as arguments for the inherited loss classes. TYPE: dict[str, Tensor]
**kwargs	Additional keyword arguments for specific loss implementations. TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Computed loss value.

Source code in spectre/loss/base.py

@abstractmethod
def forward(self, context: dict[str, torch.Tensor], **kwargs) -> torch.Tensor:
    """
    Compute loss from context dictionary.

    Parameters
    ----------
    context : dict[str, torch.Tensor]
        Dictionary containing all tensors required as arguments for the inherited
        loss classes.

    **kwargs : dict, optional
        Additional keyword arguments for specific loss implementations.

    Returns
    -------
    torch.Tensor
        Computed loss value.
    """
    raise NotImplementedError()

LossRegistry

Bases: Registry

Registry for loss classes with type metadata.

METHOD	DESCRIPTION
register	Register a loss class with the given name and types.
get_by_type	Return names of losses matching the given type.
get_types	Return types for a registered loss or loss instance.

Functions

register(name: str, types: list[LossType] | None = None) classmethod

Register a loss class with the given name and types.

Arguments

name : str Name to register the loss under. types : list[LossType] | None Loss type categories for this loss.

RETURNS	DESCRIPTION
Callable	Decorator that registers the loss class.

Source code in spectre/loss/base.py

@classmethod
def register(cls, name: str, types: list[LossType] | None = None):
    """
    Register a loss class with the given name and types.

    Arguments
    ---------
    name : str
        Name to register the loss under.
    types : list[LossType] | None
        Loss type categories for this loss.

    Returns
    -------
    Callable
        Decorator that registers the loss class.
    """
    parent_decorator = super().register(name)

    def decorator(loss_cls: type) -> type:
        result = parent_decorator(loss_cls)
        cls._types[name] = types if types is not None else []
        return result

    return decorator

get_by_type(loss_type: LossType) -> list[str] classmethod

Return names of losses matching the given type.

Arguments

loss_type : LossType Type to filter by.

RETURNS	DESCRIPTION
list[str]	Names of registered losses with the given type.

Source code in spectre/loss/base.py

@classmethod
def get_by_type(cls, loss_type: LossType) -> list[str]:
    """
    Return names of losses matching the given type.

    Arguments
    ---------
    loss_type : LossType
        Type to filter by.

    Returns
    -------
    list[str]
        Names of registered losses with the given type.
    """
    return [name for name, types in cls._types.items() if loss_type in types]

get_types(loss: str | Loss) -> list[LossType] classmethod

Return types for a registered loss or loss instance.

Arguments

loss : str | Loss Name of the registered loss or loss instance.

RETURNS	DESCRIPTION
list[LossType]	Types associated with the loss.

RAISES	DESCRIPTION
ValueError	If no loss is registered with the given name.

Source code in spectre/loss/base.py

@classmethod
def get_types(cls, loss: str | Loss) -> list[LossType]:
    """
    Return types for a registered loss or loss instance.

    Arguments
    ---------
    loss : str | Loss
        Name of the registered loss or loss instance.

    Returns
    -------
    list[LossType]
        Types associated with the loss.

    Raises
    ------
    ValueError
        If no loss is registered with the given name.
    """
    if isinstance(loss, str):
        name = loss
    elif isinstance(loss, Loss):
        # Reverse lookup: instance -> class -> name
        loss_class = type(loss)
        name = None
        for reg_name, registered_class in cls._registry.items():
            if registered_class == loss_class:
                name = reg_name
                break

        if name is None:
            raise ValueError(
                f"Loss class {loss_class.__name__} is not registered. "
                f"Available: {cls.available()}."
            )
    else:
        raise TypeError(f"Expected str or Loss, got {type(loss)}")

    if name not in cls._registry:
        raise ValueError(
            f"Unknown {cls._registry_name}: '{name}'. Available: {cls.available()}."
        )
    return cls._types.get(name, [])

Functions

initialize_loss_fn(loss_fn: Loss | str | None, loss_kwargs: dict[str, Any] | None = None) -> Loss

Initialize a loss from an instance or registry name.

Arguments

loss_fn : Loss | str | None Loss instance or name of registered loss. loss_kwargs : dict[str, Any] | None Keyword arguments passed to loss constructor (only for string names).

RETURNS	DESCRIPTION
Loss	Initialized loss instance.

RAISES	DESCRIPTION
ValueError	If loss_fn is None.
TypeError	If loss_fn is not a Loss instance or string.

Source code in spectre/loss/base.py

def initialize_loss_fn(
    loss_fn: Loss | str | None, loss_kwargs: dict[str, Any] | None = None
) -> Loss:
    """
    Initialize a loss from an instance or registry name.

    Arguments
    ---------
    loss_fn : Loss | str | None
        Loss instance or name of registered loss.
    loss_kwargs : dict[str, Any] | None
        Keyword arguments passed to loss constructor (only for string names).

    Returns
    -------
    Loss
        Initialized loss instance.

    Raises
    ------
    ValueError
        If `loss_fn` is None.
    TypeError
        If `loss_fn` is not a Loss instance or string.
    """
    return _initialize_from_registry(
        LossRegistry, loss_fn, loss_kwargs, param_name="loss_fn"
    )