Feature request: spectral normalization

[rdilip]

Are spectral normalization layers something that anyone's working on at the moment? They're present in Haiku, so there's a template to work off of, and they'd be very useful (e.g., to enable things like https://www.tensorflow.org/tutorials/understanding/sngp).

[cgarciae]

@chiamp recently added #2881, maybe this could also be added?

[dnajera27]

I am also interested in leveraging spectral normalization. I was wondering if there had been any work done to implement SNGP since this was posted. Thanks!

[chiamp]

Spectral norm is now available after #3335. Check out the API docs.