JIT invariance problem for large Python integers as arguments to function

[hawkinsp]

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]

Looks like the problem is in _promote_dtypes:

In [1]: from jax._src.numpy.lax_numpy import _promote_dtypes                                                                 

In [2]: _promote_dtypes(2 ** 100, 3.)                                                                                        
Out[2]: [DeviceArray(1.2676506e+30, dtype=float32), DeviceArray(3., dtype=float32)]

[hawkinsp]

Also note this comes from #6165 which has a test in api_test.py.

[jakevdp]

Root cause I think is a jit invariance issue in lax.convert_element_type:

In [1]: from jax import jit, lax                                                                                             

In [2]: lax.convert_element_type(2 ** 100, 'float32')                                                                        
Out[2]: DeviceArray(1.2676506e+30, dtype=float32)

In [3]: jit(lax.convert_element_type, static_argnums=1)(2 ** 100, 'float32') 
<snip>
OverflowError: Python int 1267650600228229401496703205376 too large to convert to int32

[hawkinsp]

Matt's thought on this was:
basically, when we device_put values for executing a jitted computation, we should look at the corresponding aval in the jaxpr
(rather than looking just at the type of the value we're putting, and the current value of the x64 flag)

[jakevdp]

Looking at #6165 now... I wish I'd seen that earlier. I think I previously took out the specialization on python scalar input that #6165 added back in, specifically to prevent this kind of issue

[mattjj]

Yes but that removal (in #6014, right?) caused a regression in downstream user code, specifically in a Haiku test IIRC.

I think we can keep this behavior and have jit invariance without too much trouble. But I may be wrong! If that's wrong, we'll have to change the downstream user code instead. In either case, we'll end up with jit invariance here.

[jakevdp]

I think we have three possible approaches to this jit invariance issue in lax.convert_element_type:

1. revert fix convert_element_type on large Python int inputs #6165 so non-jitted code fails like jitted code
2. entirely re-design how constants are lowered during tracing so that they are converted to the type eventually used in the jaxpr, rather than to the type corresponding to the variable the user passed
3. Accept and document this jit invariance breakage

Thoughts?

[mattjj]

I want to try 2! After all, we already want it for other reasons: to make the x64 context manager work correctly with jit. (Btw see the jax-dev chat for more context on this.)

[mattjj]

The backup plan would be 1, but that'd require changing some downstream user code. (It might be a minor change though; hard to say without investigating it.)

[mattjj]

I think this was the test that started failing after #6014, which #6165 was trying to fix. Fixing the test wouldn't be hard, and maybe the error message would be pretty clear for users who had code written this way... so 1 is a pretty viable option too. I'm not sure how much real user code relies on this now. (I really hate fixing downstream code, which is part of why I want to try approach 2 first!)

[mattjj]

@jakevdp do you want to pursue plan 1 now (with a global presubmit check)? Then if that works out easily, i.e. there isn't much downstream code to fix, we can unblock jit invariance right away. If it doesn't work easily, at least we'll have more motivation for plan 2.

I plan to revise constant handling stuff for the x64 context manager, and I'll see if plan 2 can be folded into that easily. But this way we might not have to block on it.

WDYT?

[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