Test multi-element models in nequip-framework.

[CompRhys]

Summary

• Brings forward Nequip framework changes from Big breaking API redesign #264.
• Decouples model changes from integration changes in Big breaking API redesign #264.
• Allows @falletta to look into allegro-pol in ts.
• Test multi-element nequip framework models through the model interface test factory output

currently blocked by mir-group/nequip#553

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format.
• Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.

We highly recommended installing the pre-commit hooks running in CI locally to speedup the development process. Simply run pip install pre-commit && pre-commit install to install the hooks which will check your code before each commit.

[abhijeetgangan]

Could you explain me why a separate PR is needed?

[CompRhys]

Only to elaborate the reasons already given above:

• big PRs that introduce multiple things simultaneously are more obscure than breaking apart into feature PRs. This addition is orthogonal to that api change and my impression from meeting was that you guys still wanted to change some other things about the high level integration and optimization apis. Okay there's this issue uncovered with compiling the foundation models on CPU so this may take longer to merge but otherwise this was to my mind a quick helpful thing to do given the next reason.
• We wanted to test something enabled by the nequip addition without needing to rewrite everything else we have using ts for the breaking api change.