Add a WeightedSum model wrapper - #318
Open
zyxwwxyz wants to merge 2 commits into
Open
Conversation
CI's lint job runs `ruff format --diff`, which flagged a few lines in tests/weighted_sum.py that were not run through ruff format locally. Co-Authored-By: Claude Sonnet 5 <[email protected]>
zyxwwxyz
marked this pull request as ready for review
September 4, 2026 21:19
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
metatomic.torch.weighted_sum.WeightedSum, a wrapper that combines severalexisting outputs of a model into a new output computing their fixed linear
combination (for example, combining several
"energy/<head>"outputs, such aspredictions from different DFT functionals, into a single
"energy"output).All the heads entering a given weighted sum are requested from the wrapped
model in a single call and combined by directly operating on the raw
TensorMapvalues (mts.add/mts.multiply), without detaching or otherwiseinterrupting the autograd graph. This means a single
backward()call throughthe weighted-sum output produces forces and stresses that are exactly the
weighted sum of the forces and stresses of the individual heads, instead of
requiring one evaluation (and one set of gradients) per head.
The wrapper follows the same
wrap()-based pattern asHeatFluxandDFTD3.An optional
normalize_coefficientsflag onwrap()rescales the givencoefficients to sum to one (this also works with negative coefficients, as
long as they don't sum to exactly zero).
This is motivated by checkpoints exposing several "proper" heads (e.g. one per
training-data functional) that should be combined into physically meaningful
targets - for example an ensemble/consensus energy, without paying for a
separate forward+backward evaluation per combination.
Contributor (creator of pull-request) checklist
Reviewer checklist