[flax] unfreeze initial cache in gpt models

[patil-suraj]

What does this PR do?

Fix flax generate for GPT models when the initial seq_len is 1.

The issue is the init_cache method of flax GPT2 returns the cache as a FrozenDict, but the model’s forward returns cache as a dict.

It works with seq_len > 1 because, when seq_len > 1, we call the body fun outside of the while loop -> body calls forward -> which returns cache as a dict.

then we iterate over body_fn, using lax.while_loop, and it works as the type signature of cache is similar.

It breaks for seq_len = 1 because, when it’s one we directly call the body_fn with lax.while_loop , so here the initial type of cache is FrozenDict but the forward in body_fn returns dict, which raises this error

body_fun output and input must have same type structure, got PyTreeDe...

cc @Narsil

[Narsil]

LGTM,

Thanks for looking into it!