Using jax.numpy.complex128 in Jax reports an error, but jax.numpy.complex64 does not

[zhaodazhuang]

Here is a minimal working example, when I run this piece of code, the following error is issued

import jax.numpy as jnp

n = 100
a = jnp.empty((n),dtype=float)

b = jnp.zeros(n,dtype=complex);

b = jnp.complex64(a);

print(b)

 UserWarning: Explicitly requested dtype complex128 requested in asarray is not available, and will be truncated to dtype complex64. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.

When I change jax.numpy.complex128 to jax.numpy.complex64, there will be no error. What is the reason?

[jakevdp]

Thanks for the question! This is expected behavior – unless you explicitly enable 64-bit computation, JAX will issue a warning and convert your 64-bit values to 32-bit values. For more, see JAX Sharp Bits: Double Precision.

Try this instead:

import jax
jax.config.update('jax_enable_x64', True)

import jax.numpy as jnp

n = 100
a = jnp.empty((n),dtype=float)
b = jnp.complex128(a);

print(b)