Deprecate X64 flag

This is the external tracking bug for deprecation of the X64 flag.

The final desired outcome here is the following:

• X64 flag no longer exists
• jax.numpy functions (e.g. jnp.arange, jnp.zeros, etc.) default to 32-bit, unless an explicit dtype is requested.
• Python scalars default to weak-typed 64-bit

The plan is roughly the following:

1. Change jax.numpy function defaults to 32-bit. This should only affect users of X64 mode, and will require adding explicit 64-bit dtypes to some code. At this point, enabling X64 will produce the desired long-term behavior.
2. Change the default X64 flag state to True. Users can still recover X32 behavior by explicitly setting it to False
3. Deprecate setting the X64 flag
4. Remove the X64 flag

[jakevdp]

The main rationale is that the modalism in the API complicates JAX both from a user perspective and a developer perspective.

For the user, there’s no way to reason about how a particular block of code will behave without global knowledge of the environment in which it is run (e.g. the value of the JAX_ENABLE_X64 environment variable at the time that JAX was imported).

Additionally, due to the way JAX tracing works, attempts to allow local overrides of the global state face fundamental challenges (e.g. https://github.com/google/jax/issues/5982).

Confusion about the x64 flag is one of the most common drivers of user confusion in our feedback forums.

Other similar frameworks (e.g. pytorch and tensorflow) do not have any such global modalism.

On the developer side, the global flag doubles the amount of CI testing we need to cover the package, and means that every line of code in the package must be written with this modalism in mind.

Of course, removing this would not come without tradeoffs: namely, some users would have to add more dtype specifications to their code in cases that they want 64-bit values.

These are some of the reasons we’ve decided to explore the possibility of deprecating the X64 flag. I’m curious to hear what thoughts you have – thanks

[connection-on-fiber-bundles]

Could you please shed some light on the rationale behind this change? It seems quite inconvenient for certain use cases, in which a global switch on 32 or 64 bit could be handy.