nvalchemi.training.ComposedLossFunction

class nvalchemi.training.ComposedLossFunction(components, *, weights=None, normalize_weights=True, dtype_policy=None)

Weighted sum of BaseLossFunction components.

This class owns the per-component weighting — leaves are weightless. Weights may be plain floats or LossWeightSchedule instances; they are resolved to floats at call time. By default the resolved weights are normalized to sum to 1.0 so scheduling controls relative contributions while the learning rate controls the absolute loss magnitude. Opt out with normalize_weights=False.

Components live in an torch.nn.ModuleList for .modules() / .state_dict() / nested-__repr__ support. When a component is itself a ComposedLossFunction, its components and weights are flattened into the parent element-wise so (A + B) + C is equivalent to A + B + C.

Parameters:

• components (Sequence[BaseLossFunction | ComposedLossFunction]) – Loss terms to combine; must contain at least one element.

• weights (Sequence[LossWeightSchedule | float | None] | None) – Optional per-component weights. When provided, weights must have the same length as components at construction time (i.e. top-level components — child weights inside nested compositions are multiplied element-wise by the parent weight during flattening). A None entry is shorthand for 1.0, so weights=[None, 2.0, None] means “component 1 gets 2x, others default”. Passing weights=None defaults every component to 1.0.

• normalize_weights (bool) – When True (default), resolved weights are divided by their sum at each call so the effective weights sum to 1.0. A zero-sum raises ValueError. When False, raw weighted sums are returned.

• dtype_policy (DTypePolicy | None) – Optional composed-level dtype policy applied at call time for components whose own dtype_policy is still "strict". This avoids mutating reusable leaf instances while allowing one composed loss to opt into automatic dtype alignment.

components

torch.nn.ModuleList of the flattened leaf components.

normalize_weights

Whether effective weights are renormalized to sum to 1.0.

dtype_policy

Composed-level dtype alignment policy, or None when each leaf controls dtype handling independently.

current_weight(step=0, epoch=None)

Resolve each component’s weight to a float for (step, epoch).

When normalize_weights is True the returned list sums to 1.0; otherwise it is the raw resolved weights. With normalization enabled the raw sum must be a strictly positive float or ValueError is raised.

Parameters:

• step (int) – Current global training step.

• epoch (int | None) – Current training epoch, or None when unused.

Returns:

One effective weight per component, in order.

Return type:

list[float]

Raises:

ValueError – If normalization is enabled and the raw weights do not sum to a strictly positive, finite float.

property dtype_policy: Literal['strict', 'prediction_to_target', 'target_to_prediction'] | None

Composed-level dtype policy applied to strict leaves at call time.

extra_repr()

Expose component count and normalization alongside the default repr.

Return type:

str

forward(predictions, targets, *, step=0, epoch=None, **kwargs)

Return the weighted total loss and per-component diagnostics.

Each component is called with the routed pred / target tensors, then its raw loss is scaled by the effective weight for this step. The output’s per_component_unweighted contains each raw component loss before effective weighting; per_component_weight holds the scalar weights that were applied (after normalization, if enabled); per_component_raw_weight holds the pre-normalization resolved weights so schedule ramps remain observable on single-component normalized compositions; see BaseLossFunction.per_sample_loss for the per_component_sample contract.

Parameters:

• predictions (Mapping[str, Tensor])

• targets (Mapping[str, Tensor])

• step (int)

• epoch (int | None)

• kwargs (Any)

Return type:

ComposedLossOutput

requires_eval_grad()

Whether evaluating this loss needs autograd enabled.

Inspects each leaf component’s requires_eval_grad flag. A component reporting True (e.g. a force/stress loss that differentiates the energy) forces gradient-enabled evaluation; components reporting False do not. A component reporting None is undeclared and cannot be inferred automatically.

Returns:

True when at least one component requires gradients, False when every component explicitly declares it does not.

Return type:

bool

Raises:

ValueError – When one or more components report requires_eval_grad=None and none require gradients, so the requirement is ambiguous.

weight_factors(step=0, epoch=None)

Return a flat {component_name: effective_weight} dict.

Duplicate class names get numeric suffixes (_0, _1, …) applied to all colliding entries, not only the duplicates.

Parameters:

• step (int)

• epoch (int | None)

Return type:

dict[str, float]