np.cint/np.int32 type confusion

[cloudhan]

import numpy as np
import jax.numpy as jnp

xs = np.array([1000000000000, 2, 3, 4])
jxs = jnp.array([1000000000000, 2, 3, 4])

print(xs)
print(jxs)

Outputs:

[1000000000000             2             3             4]
[-727379968          2          3          4]

[cloudhan]

OK, export JAX_ENABLE_X64=1 solves this, but no warning is raised.

[cloudhan]

Here is another issue :

import numpy as np
import jax.numpy as jnp

xs = np.array([1000000000000, 2, 3, 4])
jxs = jnp.array([1000000000000, 2, 3, 4])

for x in xs:
    print(x, x.dtype, type(x))

for x in jxs:
    print(x, x.dtype, type(x))

On Linux output:

1000000000000 int64 <class 'numpy.int64'>
2 int64 <class 'numpy.int64'>
3 int64 <class 'numpy.int64'>
4 int64 <class 'numpy.int64'>
1000000000000 int64 <class 'numpy.longlong'>
2 int64 <class 'numpy.longlong'>
3 int64 <class 'numpy.longlong'>
4 int64 <class 'numpy.longlong'>

On Windows output:

1000000000000 int64 <class 'numpy.int64'>
2 int64 <class 'numpy.int64'>
3 int64 <class 'numpy.int64'>
4 int64 <class 'numpy.int64'>
1000000000000 int64 <class 'numpy.int64'>
2 int64 <class 'numpy.int64'>
3 int64 <class 'numpy.int64'>
4 int64 <class 'numpy.int64'>

jax iterator outputs platform dependent element type
In this case, int64 became numpy.longlong on Linux but not on Windows.

I also observe int32 became numpy.cint on Windows, but not on Linux.

[hawkinsp]

I think this is an issue that other folks have hit before. There's a difference between "native" and "standard" sizes in NumPy that manifests only in Windows:
https://docs.python.org/3/library/struct.html#byte-order-size-and-alignment
pybind/pybind11#1806
pybind/pybind11#1908

My guess is we need to change some of the format descriptors returned here:
https://github.com/tensorflow/tensorflow/blob/516f4a121c5c1f9bb712b774109c6ad5283d27ad/tensorflow/compiler/xla/python/types.cc#L104

to have = prefixes.

[cloudhan]

This does not specific to windows. I updated my previous comment.