

TITLE
Flax summarization script 

COMMENTS
0

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
# 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.