[Flax] Add docstrings & model outputs (#11498)
* add attentions & hidden states
* add model outputs + docs
* finish docs
* finish tests
* finish impl
* del @
* finish
* finish
* correct test
* apply sylvains suggestions
* Update src/transformers/models/bert/modeling_flax_bert.py
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
* simplify more
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
