flax

51b3d8f1 - Merge pull request #5505 from mohsinm-dev:fix-attention-assert-to-valueerror