How to enable float64 globally?

[zgbkdlm]

Hi,

Can I ask how to enable float64 globally?

I tried to use jax.config.update("jax_enable_x64", True), but this does not work, and the resulting neural network parameters are still float32. I also tried to add dtype arguments in the model creation, but it does not work either.

class NN(nn.Module):

    dtype = jax.numpy.float64

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(features=10, dtype=jax.numpy.float64)(x)
        x = nn.gelu(x)
        x = nn.Dense(features=5, dtype=jax.numpy.float64)(x)
        x = nn.gelu(x)
        x = nn.Dense(features=1, dtype=jax.numpy.float64)(x)
        return x


model = NN()
params = model.init(key, dataset.xs)
print(params)

[soraros]

Flax defaults parameter dtype to flat32 which can be changed in your case by setting param_dtype to np.float64.

[jiaxi98]

Yeah, this method works, but is that possible to do this globally? Otherwise for a UNet implementation, I have to specify the param_dtype for every Conv blocks.