Remove workaround for lack of trainable ReplicaLocalVariables.
Remove workaround for lack of trainable ReplicaLocalVariables.

PiperOrigin-RevId: 252795306
Change-Id: I8dfb0e3239fd4f6f1b3f21f20ccac4d796ac8aed

deepmind/sonnet,deepmind/sonnet