Global way to automatically convert all floating-point computations to bfloat16

[carlosgmartin]

Given an existing program, is there any global way to automatically convert all of its floating-point computations to bfloat16? (Without having to manually litter the code with dtype=jnp.bfloat16 flags inside every constructor.)

If not, would it be worth considering adding such functionality as a feature? Perhaps via a context manager or configuration argument to jit.

[AlexGraefe]

Hi,

for this, you can use MPX (https://github.com/Data-Science-in-Mechanical-Engineering/mixed_precision_for_JAX) an update and extension of JMP that we released this month. MPX is a toolbox for mixed precision training and allows you to port your training pipeline to mixed precision with just a few changes.

In case you dont want to do mixed precision training, but just cast your function, the toolbox contains mpx.cast_function that can cast a function to a required datatype. MPX casts all inputs of the function (can be arbitrary PyTrees) to the specified datatype and by JAX's type promotion, all operations will be in this datatype (https://docs.jax.dev/en/latest/type_promotion.html).

Example:

import jax
import jax.numpy as jnp
import mpx

def my_func(x):
    x = x+1
    print(x.dtype)
    return x

my_func(jnp.zeros((42, )))   # prints float32
mpx.cast_function(my_func, dtype=jnp.bfloat16)(jnp.zeros((42, )))   # prints bfloat16

But you must make sure that the function does not contain a constant in a higher datatype (e.g. jnp.zeros) as those might make all subsequent operations to work on this higher datatype. In such cases, you need to manually cast them.

def my_func(x):
   my_const = jnp.zeros((42))   # change this to  jnp.zeros((42), dtype=jnp.bfloat16)
    x = x+my_const
    print(x.dtype)
    return x

Best,

Alex