Disabling TF32 in JAX

[Honza9723]

Dear All,

I would like to ask, how to prevent JAX from running TF32 computations on A100 GPUs. I know that I can manually set precision = 'HIGHEST' each time I call jnp dot or matmul. However, given I want to prevent TF32 calculations all round my code (including inside neural network created using Haiku library), I would prefer to use some global flag. TensorFlow seems to have such an option:
https://www.tensorflow.org/api_docs/python/tf/config/experimental/enable_tensor_float_32_execution

Is there a way, how to globally disable TF32 in JAX? So far, I am using config.update("jax_enable_x64", True), however my goal is to run the code in a standard f32 precision, since f64 is a bit too expansive for the problem I am working on.

Best,
Jan

[jaro-sevcik]

Looking at jax/_src/config.py, the following config should do the job:

jax.config.update('jax_default_matmul_precision', 'highest')

[Honza9723]

Dear Jaroslav,

thank you!

Best,
Jan