feat(training): speedup multiscale loss

[cathalobrien]

Description

This PR introduces a few changes to various losses with the aim of speeding up the multiscale loss in the temporal downscaler.

The main change is grouping the N sparse matmuls in sparse_projector.py to 1 large sparse matmul. the multi scale loss was having to load the projection matrix into L2 memory with each call to sparse matmul.

Additionally the projection matrices are now converted to CSR format during multiscale loss init time. This further speeds up the loading of the projection matrices, as well as reducing peak memory usage.

After these changes, the individual CRPS loss computations begin to dominate. Compiling them gives another speedup (that's set in the config not in this PR)

some rough timings for the loss alone, using a standalone bm script.

version	time (s)	peak memory (GB)
Initial	3.77	49
grouped sparse matmul	3.30	49
+ CSR format	2.06	44
+ compile CRPS	1.08	45

The results roughly match what I see in the traces from full training runs. The end-to-end throughput of a full training run has gone from 0.22 it/s to 0.62 it/s (0.55 without compiling)

[japols]

Nice! Should we consider adding the CRPS loss as a default compile option in config.model.compile?

(probably need to revisit compile defaults when observations are merged, specifically dynamic=False)

[ssmmnn11]

I had this in a branch to allow for directly creating the matrix in csr ; which is relevant when creating large matrices. Can you check how this would fit here as well?

[mc4117]

just for context, we're still doing loss correctness tests on this branch and I will post the mlflows here once we have the comparisons