Way to allow int64 without f64

[botev]

Currently, the only way to enable usage of i64 is using jax.enable_x64 which switches on the f64. Especially in modern times where schedules and things like that might depend on the token budget, which exceeds 2^32, it would be useful to have a way of enabling only int64, but not allowing for f64.

[jakevdp]

There is not currently any way to do this, but starting in JAX v0.8.0 the jax.enable_x64 decorator lets you enable X64 within particular contexts. So you could write something like this:

import jax
import jax.numpy as jnp

x = jnp.zeros(10)
with jax.enable_x64():
  y = jnp.arange(10)

print(x.dtype, y.dtype)  # float32 int64

[botev]

Yeah, that would indeed solve the problem I think.

[hawkinsp]

Feel free you use it right now, but you'll probably want to make sure your code is clean of accidental x64 types. Just set it to error and run your tests, fix anything that turns up.

[hawkinsp]

There's a design doc that I will try to put up on github that explains this flag. But basically it allows the explicit use of x64 types, but ways that are implicit (e.g., type promotion, converting classic numpy arrays to JAX) will still do what they used to do.

[botev]

So what happens if you then use an explicit x64 variable interacting with let's say an f32, which is outside the enable_x64 context manager?

[hawkinsp]

It turns out to work fine. x64 values obey their normal rules. All enable_x64 does is prevent you from forming one in the first place. But once you have one, they work normally.