Changes to Goldens tests to support modules which change state on forward call
Changes to Goldens tests to support modules which change state on forward call

PiperOrigin-RevId: 251436267
Change-Id: I8712c567d307bf45c7a1ab508b897469562a9645

deepmind/sonnet,deepmind/sonnet