flax

fe0b35d7 - Merge pull request #4686 from kaixih:support_fp8_einsum