integer overflow in `jnp.clip` with unsigned ints

[ev-br]

Description

Caught in data-apis/array-api-tests#410 :

when clip receives a uint64 array and python scalar max which is in range of uint64 but overflows int64,
it reports an internal overflow:

(array-api) ev-br@qgpu3:~$ JAX_ENABLE_X64=true ipython
Python 3.11.0 | packaged by conda-forge | (main, Jan 14 2023, 12:27:40) [GCC 11.3.0]
Type 'copyright', 'credits' or 'license' for more information
IPython 8.32.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import jax.numpy as jnp

In [2]: kw = {'min': 0, 'max': 9223372036854775808}

In [3]: jnp.clip(jnp.asarray(0, dtype=jnp.uint64), **kw)
WARNING:2026-01-17 20:43:18,604:jax._src.xla_bridge:850: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
---------------------------------------------------------------------------
OverflowError                             Traceback (most recent call last)
Cell In[3], line 1
----> 1 jnp.clip(jnp.asarray(0, dtype=jnp.uint64), **kw)

    [... skipping hidden 5 frame]

File ~/.conda/envs/array-api/lib/python3.11/site-packages/jax/_src/pjit.py:671, in _infer_input_type(fun, dbg_fn, explicit_args)
    669   dbg = dbg_fn()
    670   arg_path = f"argument path is {dbg.arg_names[i] if dbg.arg_names is not None else 'unknown'}"  # pytype: disable=name-error
--> 671   raise OverflowError(
    672     "An overflow was encountered while parsing an argument to a jitted "
    673     f"computation, whose {arg_path}."
    674   ) from None
    675 except TypeError:
    676   dbg = dbg_fn()

OverflowError: An overflow was encountered while parsing an argument to a jitted computation, whose argument path is max.

In [4]: jnp.iinfo(jnp.uint64).max - kw['max']
Out[4]: 9223372036854775807

In [5]: jnp.iinfo(jnp.int64).max - kw['max']
Out[5]: -1

System info (python version, jaxlib version, accelerator, etc.)

In [6]: import jax; jax.print_environment_info()
jax:    0.8.0
jaxlib: 0.8.0
numpy:  2.2.5
python: 3.11.0 | packaged by conda-forge | (main, Jan 14 2023, 12:27:40) [GCC 11.3.0]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='qgpu3', release='6.8.0-71-generic', version='#71-Ubuntu SMP PREEMPT_DYNAMIC Tue Jul 22 16:52:38 UTC 2025', machine='x86_64')
JAX_ENABLE_X64=true

$ nvidia-smi
Sat Jan 17 20:47:16 2026       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 560.35.03              Driver Version: 560.35.03      CUDA Version: 12.6     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 2060        Off |   00000000:01:00.0 Off |                  N/A |
| 31%   38C    P0             36W /  160W |       1MiB /   6144MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA GeForce RTX 2060        Off |   00000000:21:00.0 Off |                  N/A |
| 33%   41C    P0             37W /  160W |       1MiB /   6144MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+

[jakevdp]

Thanks for the report – this is working as expected. When JAX encounters a Python integer, it is always interpreted as a signed integer of the default width (32 bit or 64 bit, depending on the X64 flag state).

At one time we did have value-dependent dtype semantics, such that integers in the range (int64.max, uint64.max] would trigger compilation & tracing for uint64 dtypes, but this often led to unexpected JIT cache misses and brittle code (this is especially true when it comes to NumPy's curious choice of implicitly promoting int64 and uint64 to float64 – it meant that if you passed too large an integer value to otherwise working code, your code might start generating floats instead of ints, which can really cause problems!).

To remedy these issues, we removed value-dependent dtype semantics and added the error you're seeing: the intention is that if you want an integer value to be interpreted as an unsigned int, you need to convert it explicitly; e.g.

>>> kw = {'min': 0, 'max': np.uint64(9223372036854775808)}
>>> jnp.clip(jnp.asarray(0, dtype=jnp.uint64), **kw)
Array(0, dtype=uint64)

You could maybe argue that in the context of jnp.clip and other functions that immediately promote their inputs, perhaps the tracing machinery should do more to try and infer the user's intent. However, our experience with this kind of "smart" code is that it often leads to unexpected consequences, and as the zen of python puts it, "explicit is better than implicit".

[ev-br]

Thanks for the response Jake!

The interface between unbounded python ints and finite-range integers is bound to have some hiccups (and I totally agree that NumPy's int64->float64 promotion is... curious indeed). So I suppose this issue can be interpreted as

• a heads-up that array API testing of clip is bound to start failing after ENH: test binops with float_array, int_scalar combinations data-apis/array-api-tests#410 lands;
• a request to document this. This is an edge case of course, and the suggested workaround makes total sense; This hits users very rarely, but when it does, it's well in the "gnashing of teeth" territory of
  https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html#fin .

[mdhaber]

I just experienced this in the context of comparisons¹.

import jax.numpy as xp
xp.asarray(1, dtype=xp.uint64) == xp.iinfo(xp.uint64).max
# OverflowError: An overflow was encountered while parsing an argument to a jitted computation...

I would have expected the Python scalar to be promoted to the dtype of the array argument before performing the comparison per "Mixing arrays with Python scalars".

Details

It's possible that there is some wiggle room I'm not seeing, but it looks like "When JAX encounters a Python integer, it is always interpreted as a signed integer of the default width" is not compliant with the letter of the standard. There may be very good motivation for it - and maybe even good enough to change the wording of the standard². But yeah, in line #34455 (comment), I think it's at least worth a discussion if JAX doesn't consider this a bug.

Crossref cupy/cupy#9753 (same issue; CuPy seems to consider it a bug). Discovered in the context of testing MArray with JAX backend mdhaber/marray#154.

Footnotes

1. Update: also with where. ↩

2. Although re: https://github.com/jax-ml/jax/issues/34455#issuecomment-3768900196, I don't think avoiding "value-dependent semantics" should be the motivation, since that's a different thing. This is just a matter of following mixed array/Python scalar promotion rules; specifically, whether the standard's promotion rules are followed immediately / before following some other promotion rules. ↩

[jakevdp]

Thanks for the comment @mdhaber – I think this is a case where we're a bit constrained by JAX's current implementation. The way JAX is implemented, Python scalars must be converted to arrays at the function boundary. So when I write something like this:

@jit
def f(x, y):
  return x == y

f(jnp.uint64(1), jnp.iinfo(jnp.uint64).max)

the tracing machinery must decide the dtype of x and y before it knows how they are used. We could probably re-implement JAX so that that decision could be deferred in the case of Python scalars, but it would be a deep redesign of aspects of JAX's tracing machinery.

The current situation is maybe not as ergonomic as it could be, but there is no silent failure it does give the user an easy recourse: explicitly cast the inputs to the desired type.