Refactors acme.jax.networks.rescaling to convert hk modules to pure functions.
Refactors acme.jax.networks.rescaling to convert hk modules to pure functions.

Resolves https://github.com/deepmind/acme/issues/151

PiperOrigin-RevId: 401749906
Change-Id: I752bdd5cf2d12f6d93b3fc8418b48c98a06720be

deepmind/acme,deepmind/acme