Enforce float16

[astanziola]

Hi,
I know that you can enforce 64bit precision by setting jax.config.update("jax_enable_x64", True), but is there an analogous for enforcing 16bit precision?

[jakevdp]

Thanks for the question! Just a quick clarification: jax_enable_x64 does not enforce 64-bit precision: it allows 64-bit to be used. To answer your question directly: no, there is no way to enforce 16-bit precision for all calculations. However, JAX's type promotion semantics have been designed specifically to make it easier to maintain whatever precision you desire by explicitly setting the dtype of values you create.

As a brief example, here's numpy's promotion behavior when adding a float and an int:

import numpy as np
x = np.arange(5, dtype='float16')
y = np.ones(5, dtype=int)
print((x + y).dtype)
# float64

and here's the equivalent in JAX:

import jax.numpy as jnp
x = jnp.arange(5, dtype='bfloat16')
y = jnp.ones(5, dtype=int)
print((x + y).dtype)
# bfloat16

[astanziola]

Thank you for your answer. That looks even better than I thought!