L2 norm of weights

[tychovdo]

Hi,

This week I started using equinox (great project!), and am slowly figuring out the best practices. I'm trying to do a simple L2 norm of the weights, which requires squaring parameter values and summing the result.

What would be the best way to implement this standard regularisation term with equinox? I suspect this should be done with some sort of reduce over the pytree, but can not figure it out how to do it.

Many thanks for answering this probably silly question!

Best, Tycho