flax

738078c6 - Merge pull request #3435 from kaixih:update_fp8_support