flax

cf7d89a3 - Merge pull request #3322 from wenscarl:shuw/add_fp8_support