Improve performance of dhconv contraction

[mcgibbon]

This PR changes the shape of the weight matrix in SpectralConvS2 from [group, in, out, lat] to [group, lat, out, in] so that the matrix multiplication dimensions [out, in] are the fast dimensions, with [in] being the fastest. Backwards compatibility is included for existing checkpoints.

Changes:

• No public API changes

• Tests added

[mcgibbon]

Commits are done such that there's verification the model produces identical behavior. I made a separate commit updating the regression target which only trivially changes the weight initialization (making it cleaner, but breaking regression).

Still need to write a test that ensures previous checkpoints have their behavior unchanged, but it should be implemented/working using the same method we use to support old checkpoints without a group dim.

Before and after (note the change to dhconv and the total):


Note that with 8 groups, we no longer see an improvement as the runtime is already low (15ms on T4).

[mcgibbon]

When reviewing, please check that the regression file updates are valid by looking per-commit. The weight shape refactor is done first in a way that maintains backwards compatibility and passes regression tests, with the regression targets only updated afterwards with a trivial change to the way the weights are initialized (which changes the values they get initialized to).

I added a checkpoint-load regression test in another PR, so this is no longer an issue.

[mcgibbon]

I added a regression test for the output from a loaded checkpoint being the same, using an input checkpoint generated from main. The test caught a bug, so it's working properly.