Mixed Precision

[kimihailv]

As I understand, there isn't ready mixed precision infrastructure in Flax as in PyTorch? I found the usage of dynamic_scale in ImageNet example, but I couldn't find any documentations. Does it exist? And what is the best way to train a flax model with mixed precision (float32 and bfloat16) on gpu?

[aaarrti]

Hi @kimihailv,
You could have a look into https://github.com/deepmind/jmp.

[llCurious]

Hi, @aaarrti . The link you provided is JMP library and is awesom. I wonder is there official support of mixed-precision in Flax? An example is in Haiku: https://dm-haiku.readthedocs.io/en/latest/api.html#module-haiku.mixed_precision