Possible to set float32 as defeault, but use float64 for some parts of calculation?

[Binbose]

I have a project that is built on top of another large library. I want to add an extension at the end of the calculation that requires some precise numerics in float64, but only for a small part of the calculation (essentially, I have to calculate the difference between a large and a very small vector and take the norm of this difference, which seems to be very unstable in float32 and then take the derivative of this). I know that I can enable float64 with jax.config.update('jax_enable_x64', True)
However, that sets the default to float64, requiring me to go into the library and change everything to explicitly use float32, which is very work-intensive in a large library. Is there a way to enable float64 so that I can explicitly request it but keep the default to float32 so that I don't have to change the rest of the code?

[jakevdp]

Unfortunately, that's not really possible. jax_enable_x64 is a global setting for the Python runtime. But with jax_enable_x64=True, there's nothing to stop you from doing float32 computations: just make sure to set the dtypes of your arrays explicitly. We've done a lot of work to ensure that for 32-bit inputs you'll get 32-bit computations and 32-bit outputs, even when X64 is enabled.