Support arbitrary model outputs

[mshuaibii]

Part of #520. Provides support for arbitrary model head predictions by leveraging a model's node/edge embeddings.

Specifically:

• The custom_head flag in the config controls whether we want the model to create the head or the model architecture was created specifically for that target. If False, the desired outputs are expected to be hard-coded as part of the model. This is mainly to support existing models energy/force predictions, without changing those underlying models.
• If custom_head=True, for the model makes the respective prediction depending on whether it isn't an equivariant property (by defining irrep_dim) or a scalar/list of scalars (shape).
• For equivariant properties, we preserve equivariance through a multiplication of messages and spherical harmonics of the edges (sum[m_st*Y^irrep_dim(r_st)]).

Task tracking:

• Resolve failing test following OCP Trainer merge.
• Invariant prediction
• Equivariant predictions
  • scalar-edge multiplication (GemNet-OC)
  • eSCN spherical harmonics
• Support for arbitrary head aggregation (add, mean, etc.)
• Add unit tests for arbitrary outputs

Baselines

• OC20-2M: GemNet-OC
  • Custom energy
  • Custom forces
  • Custom both
• OC20-2M: EqV2
• Displays: EqV2 DT prediction

[misko]

@wood-b @mshuaibii @janiceblue Is this PR ready to merge? or is there something still outstanding?

[misko]

Are these changes relevant to arbitrary model support? I didnt see calculator mentioned elsewhere

[janiceblue]

Probably was an issue with merge conflicts, as OCPCalculator is in src/fairchem/core/__init__.py already, will remove