sage.factory.loss_adapters

Loss adapters — a uniform protocol so SageVanillaTraining can combine an arbitrary main loss with auxiliary losses under a GradientNormBalancer.

Different losses have different call signatures (BCEWithPEsigmaLoss takes a (ranking, point_estimates) tuple and sliced targets; ConsistencyNLLLoss takes per-detector tensors + a mask read from the batch context). An adapter wraps a loss so the trainer can call every loss the same way — adapter(out, targets, ctx) -> stacked components — and declares which component is the primary (reference / BCE) term and which are the auxiliary terms to balance.

Classes

LossAdapter	Base loss adapter.
MergedLossAdapter	Wraps BCEWithPEsigmaLoss (classification + merged heteroscedastic PE).
ConsistencyLossAdapter	Wraps ConsistencyNLLLoss (per-detector tc / mchirp NLL).

Module Contents

class LossAdapter

Base loss adapter.

Subclasses implement components() (the underlying loss call, returning a stacked component tensor whose [0] is the loss’s own total) and set:

primary_indexint or None

Index of the primary (reference) component, e.g. BCE. None if this adapter contributes only auxiliary terms.

aux_indicestuple[int]

Indices of the auxiliary components to balance against the primary.

primary_index = 0

aux_indices = ()

abstractmethod components(out, targets, ctx)

class MergedLossAdapter(loss)

Bases: LossAdapter

Wraps BCEWithPEsigmaLoss (classification + merged heteroscedastic PE).

Returns [total, bce, pe_reg, coupling]; the primary is BCE and the aux terms are pe_reg and coupling. targets[:, :mw] is the merged-width slice (PE params + class column), matching the loss’s existing contract.

primary_index = 1

aux_indices = (2, 3)

loss

mw

components(out, targets, ctx)

class ConsistencyLossAdapter(loss)

Bases: LossAdapter

Wraps ConsistencyNLLLoss (per-detector tc / mchirp NLL).

Returns [total, cons_tc, cons_mc] — both are auxiliary terms (no primary). Reads the per-detector supervision per_det_mask from ctx (populated by the masking callback) and slices the per-detector tc / mc targets from the appended target columns. tc is window-normalised (/ tc_scale) to match the model’s window-normalised mu_tc.

primary_index = None

aux_indices = (1, 2)

loss

D

tc_scale

components(out, targets, ctx)