jax
5a5bbe33 - Fix non-termination in `jnp.gcd` and `jnp.lcm` for minimum signed integers.

Commit
4 days ago
Fix non-termination in `jnp.gcd` and `jnp.lcm` for minimum signed integers. For a signed integer dtype, taking the absolute value of the minimum representable integer (`iinfo(dtype).min`) overflows back to the negative minimum value in two's complement arithmetic. In `jnp.gcd` (and `jnp.lcm`, which delegates to `jnp.gcd`), passing this negative value into the Euclidean `while_loop` produced negative remainders and triggered the post-modulo `x1 < x2` swap whenever the remainder reached zero, swapping zero into `x1` and the negative divisor into `x2` so the loop never terminated (https://github.com/jax-ml/jax/issues/40840). Cast `abs(x1)` and `abs(x2)` to the corresponding unsigned integer dtype before computing the greatest common divisor and least common multiple, then convert the result back to the promoted input dtype while preserving weak typing. Also remove the redundant post-modulo comparison swap in the `jnp.gcd` loop body. Fixes #40840. PiperOrigin-RevId: 1018491603
Author
Jake VanderPlas
Parents
Loading