PR #1291 Expose `dot_product_attention_weights`

Commits

Expose `dot_product_attention_weights`

avital committed 4 years ago
Correct assertion string

avital committed 4 years ago
Add more assertions to dot_product_attention_weights

avital committed 4 years ago

flax Expose `dot_product_attention_weights` #1291 Merged