Adapting existing codebase to include a small function that uses 64-bit mode

[adam-hartshorne]

I have an existing codebase which is run in float32. However, I need to incorporate a method which includes taking the inverse of a matrix that can become unstable when done in 32-bit. Thus I wish to cast the input of the function to 64-bit, take the inverse and cast it back.

i.e.

def func(x):
    x = x.astype(jnp.float64)
    y = do stuff...
    return y.astype(jnp.float32)

Obviously, I need to enable x64 using

jax.config.update("jax_enable_x64", True)

However, this will also default everything else to x64 unless otherwise explicitly typed to 32-bit. Now in theory, I could go through my entire code base and insert the casts, but this will be convoluted process. Is there an easier way to say despite being in 64-bit mode everything should default to 32-bit unless I specifically cast to 64-bit?

[adam-hartshorne]

To answer my own question, and pose another one,

def func(x):
    with jax.experimental.enable_x64():
        x = x.astype(jnp.float64)
        y = do stuff...
        return y.astype(jnp.float32)

This is under the experimental section of the api. Now some "experimental" in JAX are really far from experimental e.g. the ode solver, is this ready for prime time or it is known to still be buggy?

[jakevdp]

This context manager has known bugs and I wouldn't suggest relying on it.

Unfortunately, we don't have any great alternative aside from enabling X64 globally and being careful to explicitly use 32-bit types elsewhere.