Florian

fix: ScaleModule and SumModule for DiagHessian.

3

Partially fixes #316 . ScaleModule is also used for `torch.nn.Identity`. Not sure if `hessian_is_zero()` should always return `True` for those two modules. Same with `accumulate_backpropagated_quantities()` which concats `dicts` instead of...

There are multiple layers which are specified as [being supported](https://docs.backpack.pt/en/master/supported-layers.html) for second order derivatives that actually do not work when trying to calculate the Hessian diagonal using `backpack-for-pytorch

Last remains of depreciated eqx.experimental statefulness behavior.

When running `jax.vmap`, e.g. to compute per-sample gradients, the `fin_steps` and `scale` attributes of `DynamicScale` might become arrays, leading to an error in the next step during training if not...