Last part of #364.
Different types of targets makes it necessary to have different types of heads in models. In this PR, heads are made more flexible. A change to the handling of internal features is also needed, as the last-layer features can now be per-target (or, equivalently, per-head). The LLPR module is adapted accordingly.
Finally, the last common representation of the model (before the heads are applied) is exposed as the features output (a standard output in metatensor.torch.atomistic).
Coming soon in this PR: features and last-layer features for PET.
Contributor (creator of pull-request) checklist
[x] Tests updated (for new features and bugfixes)?
[x] Documentation updated (for new features)?
[x] Issue referenced (for PRs that solve an issue)?
Last part of #364. Different types of targets makes it necessary to have different types of heads in models. In this PR, heads are made more flexible. A change to the handling of internal features is also needed, as the last-layer features can now be per-target (or, equivalently, per-head). The LLPR module is adapted accordingly. Finally, the last common representation of the model (before the heads are applied) is exposed as the
featuresoutput (a standard output inmetatensor.torch.atomistic).Coming soon in this PR: features and last-layer features for PET.
Contributor (creator of pull-request) checklist
📚 Documentation preview 📚: https://metatrain--398.org.readthedocs.build/en/398/