flax

8ed5af59 - Fix linen wmt example cache init handling bug.

Fix linen wmt example cache init handling bug. AddPositionEmbs was initializing a cache index in the Encoder which does not use the cache, resulting in an error with our stricter mutability checks. PiperOrigin-RevId: 339313349