Is there any convenient way to do spectrum normalization? (Sonnet version: 2.0.0)

Dear Deepminder:

For generative adversarial network, there introduces a normalization called spectrum normalization (ref: https://arxiv.org/abs/1802.05957).

However I find according to the knowledge of me, it is a little hard for me to incorporate it into Sonnet.

This is because spectral normalization changes the weights. But I checked snt.Conv2D, seems the weight is intractable due to internal management.

Old version sonnet seems to have a solution, (https://github.com/deepmind/sonnet/issues/99) but I can't find "custom_getter" and other Solutions in Sonnet version: 2.0.0.

So could I be given some suggestion regarding such a situation? Thanks a lot.