flax
2414d1cc - Merge pull request #3550 from chiamp:attention

Commit

2 years ago

Merge pull request #3550 from chiamp:attention PiperOrigin-RevId: 590974333

References

#3550 - changed `return_weights` to `sow_weights` for attention layer

Author

a-googler

a-googler

Parents

Loading