Flax summarization script 
# What does this PR do?

Adds Flax summarization example.

Beam search `generate` works like a charm on TPU, really fast! This makes it very easy and fast to use the `generate` in the eval loop.