Rank2 tensor head

[lbluque]

This PR adds a rank 2 tensor head for equiformerV2.

It is based on the implementations in the stress_trainer branch:
https://github.com/FAIR-Chem/fairchem/blob/stress_trainer/ocpmodels/models/utils/outer_block.py

TODO

• Add unit tests

[codecov]

Codecov Report

Attention: Patch coverage is 56.25000% with 49 lines in your changes missing coverage. Please review.

Files	Patch %	Lines
...ore/models/equiformer_v2/prediction_heads/rank2.py	55.45%	49 Missing ⚠️

Files	Coverage Δ
...airchem/core/models/equiformer_v2/equiformer_v2.py	88.47% <100.00%> (ø)
...ore/models/equiformer_v2/prediction_heads/rank2.py	55.45% <55.45%> (ø)

... and 1 file with indirect coverage changes

[misko]

This looks great! No stress, LGTM! 🚀

[misko]

I ran comparisons between the original rank2 stress tensor code that @lbluque has and compared it to the Hydra equivalent model.
https://fairwandb.org/fairchem/mptrj/reports/Stress-head-MPTraj-original-vs-Hydra-version--Vmlldzo0NDQ4NQ

The results initially were different, however we traced the issue down to a difference in using amp in the head. This PR moves all heads to by default be used without AMP.

[lbluque]

should the default be to not use it?

[misko]

When i talked to @wood-b he mentioned that we should have it off by default, it shouldnt change much memory usage and from our experiments it only helps with performance metrics