flax

changed `return_weights` to `sow_weights` for attention layer

#3550

Merged