flax

Fix vmap axes for input_positions in dot_product_attention_with_rope

#5537

Open