flax

710d2023 - Add more assertions to dot_product_attention_weights