[flax] unfreeze initial cache in gpt models
# 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 