[flax generate] allow passing params to encode
# What does this PR do?

Allows passing the user-provided `params` to the `encode` method for seq-2-seq generation. This is required to be able to `pjit` the `generate` method as we need to explicitly pass the sharded parameters. 