Move variable declaration to `__call__` function.
Move variable declaration to `__call__` function.

Declaring flax variables outside the `@nn.compat` module causes issue in multi-device training (esp. when lifted transformation is involved). This change resolve this by moving the declarations within `__call__()`.

PiperOrigin-RevId: 387496222

google/edward2