Flax beam search fix
# What does this PR do?

Makes it so you can pass `decoder_attention_mask` into `model.generate` for flax models when doing beam search. This is helpful for models like Whisper where there may be variable length decoder prefixes across a batch, so you'd have to define a `decoder_attention_mask`.

@sanchit-gandhi
@sgugger 
@frmccann97
cc @gante<|||||>_The documentation is not available anymore as the PR was closed or merged._