nvalchemi.training.ReductionContext

class nvalchemi.training.ReductionContext

Lightweight metadata bag flowing through the loss template pipeline.

A plain dict subclass used to pass metadata between BaseLossFunction.normalize(), mask(), and reduce(). Using a bare dict instead of TypedDict(total=False) keeps the type torch.compile-safe (Dynamo rejects TypedDict with optional keys).

Conventional keys

"weights"torch.Tensor

Per-sample weights for the final reduction. For energy losses with per_atom=True this carries atom counts (B, 1); for force losses it may carry per-atom or per-component weights.