Is it possible to get access to uint64 without changing the default behavior of jax.numpy?

[hr0nix]

From what I understand, it seems impossible to get access to uint64 dtype without changing the behavior of jax.numpy (i.e. the dtype of the result when dtype is not provided), or am I wrong?

In my case, I just need uint64 in a single place in my code (accumulating large numbers such as token counts), but I want to preserve the behaviour of the rest of the code, mainly because I don't trust external libraries to always specify dtypes explicitly, and I want to ensure my computations don't use 64-bit types.

Is there a way to achieve what I need? If not, what's the reason for the current semantics of jax_enable_x64? My use case seems (at least to me) perfectly reasonable.

[jakevdp]

No, there is no way to allow creation of 64-bit data types without allowing them globally via the JAX_ENABLE_X64 flag.

[davisyoshida]

That's only true for code running under the JIT right? Otherwise doesn't the context manager let you toggle the behavior? You could JIT your code with x64 enabled and JIT the 3rd party code without for example.

[jakevdp]

jax.experimental.enable_x64 remains experimental because there are surprising corner cases where it doesn't work correctly. You can try it, but understand you're using it at your own risk.

[davisyoshida]

Ah interesting, thanks!