flax

b3785dc7 - Merge pull request #3984 from kaixih:rename_fp8_custom_dtype