FlaxGPTNeo
# What does this PR do?

This PR adds the Flax version of GPTNeo. For local attention, it uses the fix proposed by @finetuneanon in #11630.

Thanks a lot, @finetuneanon for proposing the solution, it's especially important in JAX/Flax where we can't have dynamic shapes.

Official GPTNeo flax checkpoints are up on the hub and slow tests are passing.