[Feature] Store SOAP condition matrices as the dtype of their parameters

[kylevedder]

After #333 SOAP functions correctly, but it has significant excess VRAM usage when training models with reduced precision weights (e.g. bfloat16).

This PR initializes and updates the condition matrices based on the precision of the parameters themselves rather than defaulting to float32 for everything.

Note that the only exception to this is the QR factorization to get the orthogonal Q is done in float32, regardless of the underlying matrix precision, as

• only float32 has CUDA kernel support as of PyTorch 2.5.1
• precision matters for making the matrix actually orthogonal (I attempted to hand-roll Newton iteration but it was causing NaNs during optimization)

[kylevedder]

should now comply with the formatting requirements :)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 100.00%. Comparing base (c7496b0) to head (b177ac2).
Report is 23 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##              main      #335    +/-   ##
==========================================
  Coverage   100.00%   100.00%            
==========================================
  Files          108       110     +2     
  Lines         8509      8731   +222     
==========================================
+ Hits          8509      8731   +222     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[kozistr]

Thanks for your fix! I left some minor reviews, and others look great to me :)

--- updated

I just committed them to merge this PR