JIT invariance problem for large Python integers as arguments to function

In [1]: import jax, jax.numpy as jnp

In [2]: jnp.multiply(2 ** 100, 3.)  # doesn't crash
   ...:
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Out[2]: DeviceArray(3.8029518e+30, dtype=float32)

In [3]: jax.jit(jnp.multiply)(2 ** 100, 3.)
...
OverflowError                             Traceback (most recent call last)
<ipython-input-3-2aed3283764c> in <module>
----> 1 jax.jit(jnp.multiply)(2 ** 100, 3.)

~/p/jax/jax/_src/dtypes.py in _scalar_type_to_dtype(typ, value)
    125   if typ is int and value is not None:
    126     if value < np.iinfo(dtype).min or value > np.iinfo(dtype).max:
--> 127       raise OverflowError(f"Python int {value} too large to convert to {dtype}")
    128   return dtype
    129

OverflowError: Python int 1267650600228229401496703205376 too large to convert to int32

[jakevdp]

If we’re going to pursue plan 2, then in this case I think we can leave this issue open for now and reduce churn. This seems like a pretty low-priority JIT invariance case, with an already good error message and a pretty clear fix (manually casting the input to the desired type)

[jakevdp]

I think the issue has been fixed:

In [1]: import jax, jax.numpy as jnp                                                                                                                                                                                                                                                                            

In [2]: jnp.multiply(2 ** 100, 3.)                                                                                                                                                                                                                                                                              
---------------------------------------------------------------------------
OverflowError                             Traceback (most recent call last)
<ipython-input-2-b332690c3005> in <module>
----> 1 jnp.multiply(2 ** 100, 3.)

    [... skipping hidden 6 frame]

~/github/google/jax/jax/_src/dtypes.py in _scalar_type_to_dtype(typ, value)
    149   if typ is int and value is not None:
    150     if value < np.iinfo(dtype).min or value > np.iinfo(dtype).max:
--> 151       raise OverflowError(f"Python int {value} too large to convert to {dtype}")
    152   return dtype
    153 

OverflowError: Python int 1267650600228229401496703205376 too large to convert to int32