Ctrl+K

nvalchemi.training.BaseLossFunction#

class nvalchemi.training.BaseLossFunction(*, dtype_policy='strict')[source]#

Abstract torch.nn.Module base for ALCHEMI loss functions.

BaseLossFunction implements a template-method forward() pipeline that orchestrates five overridable hooks:

  1. validate() — shape / dtype checks.

  2. normalize() — pre-process pred and target (e.g. per-atom energy division) and return a ReductionContext for downstream hooks.

  3. mask() — produce a boolean validity tensor (e.g. torch.isfinite, padding masks).

  4. compute_residual() — abstract; the only method every leaf must implement. Receives pred, target, and the validity mask produced by step 3.

  5. reduce() — collapse the residual tensor and validity mask into a scalar loss and populate per_sample_loss.

Loss authors subclass BaseLossFunction and override compute_residual() at a minimum. Normalization, masking, and reduction come free via the defaults, or can be overridden individually for domain-specific behaviour (e.g. per-atom energy division in normalize(), padding-aware force masking in mask(), graph-balanced force reduction in reduce()).

Leaves are weightless — weighting and scheduling live on ComposedLossFunction. Operator sugar (scalar * leaf, leaf + leaf, sum([...])) produces a composition; see ComposedLossFunction for semantics.

Parameters:

dtype_policy (DTypePolicy)

requires_eval_grad#

Whether this loss term requires autograd during evaluation. Losses based on derived outputs such as forces and stress should set this to True; direct scalar-output losses should set it to False. None means callers cannot infer the policy automatically.

Type:

bool | None

dtype_policy#

How forward handles prediction/target dtype mismatches before validation. strict preserves both tensors and raises on mismatch. The other policies cast one tensor to the other’s dtype before the leaf validates shapes and dtypes.

Type:

{“strict”, “prediction_to_target”, “target_to_prediction”}

per_sample_loss#

Detached per-graph loss tensor of shape (B,) left as a side effect of the most recent forward() call, or None when the loss does not naturally compute a per-graph view (or when forward has never been called). Intended for logging and diagnostics only — gradients flow through the scalar returned by forward(), not through this attribute.

Type:

torch.Tensor | None

align_dtypes(pred, target)[source]#

Return prediction and target tensors adjusted by dtype_policy.

strict preserves both tensors and leaves dtype mismatches to validate(). prediction_to_target and target_to_prediction cast only when needed and preserve the source tensor otherwise.

Parameters:
Return type:

tuple[Tensor, Tensor]

abstractmethod compute_residual(pred, target, valid)[source]#

Return the per-element residual tensor.

This is the only hook that must be overridden. The valid mask (from mask()) is provided so the leaf can zero out invalid positions before computing the residual (important for operations like vector_norm where masking after the reduction would be incorrect).

Parameters:
Return type:

Tensor

property dtype_policy: Literal['strict', 'prediction_to_target', 'target_to_prediction']#

Dtype alignment policy applied before validation.

forward(pred, target, **kwargs)[source]#

Template-method pipeline: validate → normalize → mask → residual → reduce.

Subclasses should not override this method. Override the individual hooks instead. Extra keyword arguments (batch, batch_idx, num_nodes_per_graph, etc.) are forwarded to every hook via **kwargs.

Parameters:
Return type:

Tensor

mask(pred, target, ctx, **kwargs)[source]#

Return a boolean validity mask for target.

The default implementation returns an all-True mask matching target’s shape. Override to exclude non-finite entries, padding, or any other invalid positions.

Parameters:
Return type:

Tensor

normalize(pred, target, **kwargs)[source]#

Pre-process prediction and target before residual computation.

Returns a (pred, target, ctx) triple. The default implementation is the identity — ctx is empty.

Override to inject per-atom energy division, or any other pre-processing that should be available to all loss authors as a composable step.

Parameters:
Return type:

tuple[Tensor, Tensor, ReductionContext]

reduce(residual, valid, ctx, **kwargs)[source]#

Collapse a residual tensor to a scalar loss.

The default implementation computes a validity-weighted mean: (residual * valid_float).sum() / valid_float.sum(), where valid_float incorporates optional ctx["weights"].

Override for domain-specific reductions (graph-balanced force reduction, RMSD, etc.). Implementations should also populate per_sample_loss with a detached (B,) tensor when a per-graph decomposition is available.

Parameters:
Return type:

Tensor

validate(pred, target)[source]#

Check shape and dtype compatibility of pred and target.

Default implementation calls assert_same_shape() with strict=True when prediction_key / target_key attributes are present on the instance.

Parameters:
Return type:

None

On this page
Show Source