`jax.config.update("jax_enable_x64", True); jnp.array(..., dtype=int)` produces 32-bit instead of 64-bit result on Windows

[patrick-kidger]

So I realise Windows isn't properly supported. I'm raising this really just so that something is availble if people search for it, or as a bug to fix when/if Windows support ever happens.

This:

import jax
jax.config.update("jax_enable_x64", True)
jax.numpy.array(1, dtype=int)

produces an int64 array on Linux but an int32 array on Windows. I'm assuming the Linux behaviour is the correct one.

(I'm using jaxlib built from https://github.com/cloudhan/jax-windows-builder; in patrick-kidger/diffrax#67 a user reports the same behaviour when building jaxlib from source.)

[hawkinsp]

Not having access to a Windows box, are you sure that np.array(1, dtype=int) does not do the same?

[patrick-kidger]

NumPy has the same Windows->int32; Linux->int64 behaviour. e.g. numpy/numpy#17640 discusses this. It looks like dtype=int gives a C long which is 32-bit on Windows and 64-bit everywhere else. (And that maybe they'd like to change this behaviour at some point?)

JAX already handles dtypes differently to NumPy (a la the existence of JAX_ENABLE_X64) though of course.

[jakevdp]

JAX currently uses NumPy's dtype system, so when faced with dtype=int, it employs numpy's dtype constructor to convert it to a dtype object. I believe this is why its behavior is the same as NumPy for this case.

[dhirschfeld]

Isn't using dtype=int a code-smell on any platform? I think it's recommended to be explicit - e.g. dtype=np.int64.

i.e. code should never be using dtype=int so there's nothing to fix here... unless it's used internally in jax?

[hawkinsp]

JAX does not itself use dtype=int. But if you do use it, then it will mean the same thing to JAX that it means to NumPy. So I think this is working as intended: we're trying to match NumPy, and if you wanted something different then specify an explicit size?

[patrick-kidger]

FWIW my expectation/intention here -- without being aware of NumPy's behaviour -- was that it would respect the dtype-ness of the 64-bit flag. That is, it would match jax.numpy.array(1).dtype.

• I could see this being useful when needing to instead do e.g.

jax.numpy.array(<user-provided-float>, dtype=jax.numpy.array(1).dtype)

• I think having different behaviour cross-platform is a bad thing to do. (In particular as lots of pure-Python projects go through CI on only one platform.)

...but I don't actually have any strong feelings on this. This can just be closed if the resolution is "not a bug".

[Micky774]

Ultimately this is an issue with the semantics of ctypes (in particular, Windows vs Linux long), but within the scope of JAX, this behavior is consistent with expectations in the ecosystem, and in general dtype=int is an anti-pattern that should be avoided whenever possible. I'm going to go ahead and close this as intended behavior.