3572) was rolled back because of internal test breakages. This PR re-adds the attention changes again and fixes internal tests

Commit

1 year ago

[#3572](https://github.com/google/flax/pull/3572) was rolled back because of internal test breakages. This PR re-adds the attention changes again and fixes internal tests PiperOrigin-RevId: 597881774

Author

chiamp

Committer

a-googler

Parents

ed44f521

flax 9cfe078a - [#3572](https://github.com/google/flax/pull/3572) was rolled back because of internal test breakages. This PR re-adds the attention changes again and fixes internal tests

flax
9cfe078a - [#3572](https://github.com/google/flax/pull/3572) was rolled back because of internal test breakages. This PR re-adds the attention changes again and fixes internal tests