Repository: google/jax

Issue:
jax.numpy.lcm int8 overflow difference with np
### Description

The following code in jax

```python
import jax.numpy as jnp
a = jnp.array([5], dtype=jnp.int8)
b = jnp.array([26], dtype=jnp.int8)
print(jnp.lcm(a, b))
```
produces `[126]` 

while in numpy (checked in 1.23.0 and 1.21.6) (and pytorch 1.11 and 1.12 ) it gives `[-126]` keeping true to the overflow for an `int8` dtype.

Is this difference of results between **`np`** and **`jnp`** desired here?

Surely, `jax.numpy` cannot follow NumPy API in every single output but this is something achievable in this case.
The cause is perhaps that jax uses `abs` at the end of lcm operation. I suppose it should suffice to apply `abs` on `x1` and `x2` as `gcd` does.

Note: Mathematically lcm(5,26)=130, so getting a negative value (as with numpy) can at least signal a warning but getting a positive (but incorrect) value close to the correct value can easily mislead.

int8 dtype was just picked for the sake of simpler example but I am guessing this won't be restricted to int8 dtype.

If you want then I can create a PR for this and in case I m totally mistakened please do enlighten me. Thank you very much.

### What jax/jaxlib version are you using?

jax 0.3.21, jaxlib v0.3.20. Also checked with jax master repo on GitHub.

### Which accelerator(s) are you using?

CPU

### Additional system info

Python 3.8.10, Ubuntu 20, Google Colab

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report. The reason for the difference is the `abs` call here: https://github.com/google/jax/blob/ec5b1c93d73da2fefb99940f629b41a5f3540de2/jax/_src/numpy/lax_numpy.py#L4192

I'm not sure why the `abs` is there, but I think in principle we could remove it in order to match numpy's behavior in this.
> Thanks for the report. The reason for the difference is the `abs` call here:
> 
> https://github.com/google/jax/blob/ec5b1c93d73da2fefb99940f629b41a5f3540de2/jax/_src/numpy/lax_numpy.py#L4192
> 
> I'm not sure why the `abs` is there, but I think in principle we could remove it in order to match numpy's behavior in this.

Exactly. I had checked and removed this abs call and it the results do match with np. Let me know if this is approved for for a PR alongside any other changes if required.
So long as it doesn't break any existing tests, that sounds good.

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -4367,11 +4367,12 @@ def gcd(x1: ArrayLike, x2: ArrayLike) -> Array:
 def lcm(x1: ArrayLike, x2: ArrayLike) -> Array:
   _check_arraylike("lcm", x1, x2)
   x1, x2 = _promote_dtypes(x1, x2)
+  x1, x2 = abs(x1), abs(x2)
   if not issubdtype(_dtype(x1), integer):
     raise ValueError("Arguments to jax.numpy.lcm must be integers.")
   d = gcd(x1, x2)
   return where(d == 0, _lax_const(d, 0),
-               abs(multiply(x1, floor_divide(x2, d))))
+               multiply(x1, floor_divide(x2, d)))
 
 
 @_wraps(np.extract)


Test patch:
diff --git a/tests/lax_numpy_operators_test.py b/tests/lax_numpy_operators_test.py
--- a/tests/lax_numpy_operators_test.py
+++ b/tests/lax_numpy_operators_test.py
@@ -282,6 +282,7 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
               all_shapes, jtu.rand_small_positive, []),
     op_record("gcd", 2, int_dtypes_no_uint64, all_shapes, jtu.rand_default, []),
     op_record("lcm", 2, int_dtypes_no_uint64, all_shapes, jtu.rand_default, []),
+    op_record("lcm", 2, [np.int8], all_shapes, jtu.rand_not_small, [])
 ]
 
 JAX_BITWISE_OP_RECORDS = [