Mixed precision for TPU with jax

[sh0416]

Hi, I usually work with GPU and use jmp for mixed precision to enable tensor core.

My question is..

Do we also cast the data type of array to float16 to accelerate tpu performance?

For now, I didn't cast it, so the performance is not that I expected..

When I cast it, It doesn't work...

Thanks,

[sh0416]

I think I found some useful link for me.

https://github.com/google/jax/blob/main/cloud_tpu_colabs/README.md