Add support for Opt in tf and flax
# What does this PR do?

Adds support for OPT in Flax and TF. 
Also clean Pytorch code a bit. 

## Who can review?
 @LysandreJik, @patrickvonplaten, @patil-suraj, @sgugger

Closing for a new pull request where history is fixed