Repository: google/jax

Issue:
jnp.min or jnp.max returns nan if input array dtype is fp8
### Description

Below is a toy example to reproduce. Can be reproduced with CPU, TPU and GPU. Tested on colab.

```python
import jax
import jax.numpy as jnp
import jax.lax as lax

# Define a function to generate random arrays
def random_array(shape):
  return jax.random.uniform(jax.random.PRNGKey(0), shape, minval=-1, maxval=1
  
test = random_array((2,))
print(test, test.astype(jnp.float8_e4m3fn))
print(jnp.isnan(test))
print(jnp.min(test), jnp.max(test))
print(jnp.isnan(test.astype(jnp.float8_e4m3fn)))
print(jnp.min(test.astype(jnp.float8_e4m3fn)), jnp.max(test.astype(jnp.float8_e4m3fn)))

[0.21629536 0.8041241 ] [0.21875 0.8125]
[False False]
0.21629536 0.8041241
[False False]
nan nan
```

### What jax/jaxlib version are you using?

jax v0.4.22

### Which accelerator(s) are you using?

CPU/GPU/TPU

### Additional system info?

1.26.2 3.11.6 (stable, redacted, redacted) [Clang google3-trunk (af7a1453526a88a0e242baf156244aa4ae42ae4b)] uname_result(system='Linux', node='11a0354b588110fa-30029b6032f.borgtask.google.com', release='5.10.0-smp-1100.464.0.0', version='#1 [v5.10.0-1100.464.0.0] SMP @1700037056', machine='x86_64')

### NVIDIA GPU info

```
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| 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  Tesla V100-SXM2...  Off  | 00000000:B3:00.0 Off |                    0 |
| N/A   40C    P0    72W / 300W |  12449MiB / 16384MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     88841      C   .../mount/server/ml_notebook    12446MiB |
+-----------------------------------------------------------------------------+
```


Discussion:
Hi, thanks for the question! 8-bit floating point types are experimental and in general you should not expect them to be compatible with most operations in JAX (we focus on matrix multiplication and a few related operations that `float8*` were designed for).

I believe the particular behavior here is caused by the fact that `float8_e4m3fn` cannot represent infinities (the "`f`" in the type name stands for "finite only") and the `min` and `max` computation assume the presence of $\pm\infty$ as the identity for the reduction.

With that in mind, this working as expected for the time being – we could certainly modify the `reduce_min`/`reduce_max` lowering code to account for this special case, but it's only one of many examples where `float8` inputs will not be correctly handled in JAX, and it's not clear at the moment that addressing those cases more comprehensively is a goal for the project.

If you are doing general floating point operations, you should stick with non-experimental dtypes. For this particular application, you might try `float8_e5m2` which is able to represent infinite values.
So thinking more about this, we probably should make `min` and `max` do something sane in this case, and we need to do a better job of explicitly erroring for unsupported fp8 operations. I'm going to take a look

Resolving patch:
diff --git a/jax/_src/dtypes.py b/jax/_src/dtypes.py
--- a/jax/_src/dtypes.py
+++ b/jax/_src/dtypes.py
@@ -100,6 +100,13 @@ def type(self) -> type: ...
 _float8_e5m2_dtype: np.dtype = np.dtype(float8_e5m2)
 _float8_e5m2fnuz_dtype: np.dtype = np.dtype(float8_e5m2fnuz)
 
+def supports_inf(dtype: DTypeLike) -> bool:
+  """Return true if the dtype supports infinity, else return False."""
+  typ = np.dtype(dtype).type
+  if typ in {float8_e4m3b11fnuz, float8_e4m3fn, float8_e4m3fnuz, float8_e5m2fnuz}:
+    return False
+  return issubdtype(dtype, np.inexact)
+
 # bfloat16 support
 bfloat16: type[np.generic] = ml_dtypes.bfloat16
 _bfloat16_dtype: np.dtype = np.dtype(bfloat16)
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -1085,7 +1085,8 @@ def _get_prod_identity(dtype: DTypeLike) -> np.ndarray:
 
 def _get_max_identity(dtype: DTypeLike) -> np.ndarray:
   if dtypes.issubdtype(dtype, np.inexact):
-    return np.array(-np.inf, dtype)
+    return np.array(-np.inf if dtypes.supports_inf(dtype) else dtypes.finfo(dtype).min,
+                    dtype=dtype)
   elif dtypes.issubdtype(dtype, np.integer):
     return np.array(dtypes.iinfo(dtype).min, dtype)
   elif dtypes.issubdtype(dtype, np.bool_):
@@ -1095,7 +1096,8 @@ def _get_max_identity(dtype: DTypeLike) -> np.ndarray:
 
 def _get_min_identity(dtype: DTypeLike) -> np.ndarray:
   if dtypes.issubdtype(dtype, np.inexact):
-    return np.array(np.inf, dtype)
+    return np.array(np.inf if dtypes.supports_inf(dtype) else dtypes.finfo(dtype).max,
+                    dtype=dtype)
   elif dtypes.issubdtype(dtype, np.integer):
     return np.array(dtypes.iinfo(dtype).max, dtype)
   elif dtypes.issubdtype(dtype, np.bool_):
diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -167,6 +167,10 @@ def _reduction_init_val(a: ArrayLike, init_val: Any) -> np.ndarray:
   a_dtype = dtypes.canonicalize_dtype(dtypes.dtype(a))
   if a_dtype == 'bool':
     return np.array(init_val > 0, dtype=a_dtype)
+  if (np.isinf(init_val) and dtypes.issubdtype(a_dtype, np.floating)
+      and not dtypes.supports_inf(a_dtype)):
+    init_val = np.array(dtypes.finfo(a_dtype).min if np.isneginf(init_val)
+                        else dtypes.finfo(a_dtype).max, dtype=a_dtype)
   try:
     return np.array(init_val, dtype=a_dtype)
   except OverflowError:


Test patch:
diff --git a/jax/_src/public_test_util.py b/jax/_src/public_test_util.py
--- a/jax/_src/public_test_util.py
+++ b/jax/_src/public_test_util.py
@@ -93,8 +93,14 @@ def _assert_numpy_allclose(a, b, atol=None, rtol=None, err_msg=''):
     np.testing.assert_array_equal(a, b, err_msg=err_msg)
     return
 
-  custom_float_dtypes = [_dtypes.float8_e4m3b11fnuz, _dtypes.float8_e4m3fn,
-                         _dtypes.float8_e5m2, _dtypes.bfloat16]
+  custom_float_dtypes = [
+    _dtypes.float8_e4m3b11fnuz,
+    _dtypes.float8_e4m3fn,
+    _dtypes.float8_e4m3fnuz,
+    _dtypes.float8_e5m2,
+    _dtypes.float8_e5m2fnuz,
+    _dtypes.bfloat16,
+  ]
   def maybe_upcast(x):
     if x.dtype in custom_float_dtypes:
       return x.astype(np.float32)
diff --git a/tests/lax_numpy_reducers_test.py b/tests/lax_numpy_reducers_test.py
--- a/tests/lax_numpy_reducers_test.py
+++ b/tests/lax_numpy_reducers_test.py
@@ -46,6 +46,7 @@
 nonempty_shapes = scalar_shapes + nonempty_array_shapes
 all_shapes = scalar_shapes + array_shapes
 
+custom_float_dtypes = jtu.dtypes.custom_floats
 float_dtypes = jtu.dtypes.all_floating
 complex_dtypes = jtu.dtypes.complex
 int_dtypes = jtu.dtypes.all_integer
@@ -113,8 +114,8 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
     op_record("prod", 1, all_dtypes, all_shapes, jtu.rand_small_positive, []),
     op_record("sum", 1, all_dtypes, all_shapes, jtu.rand_default, [],
               tolerance={jnp.bfloat16: 2e-2}),
-    op_record("max", 1, all_dtypes, all_shapes, jtu.rand_default, []),
-    op_record("min", 1, all_dtypes, all_shapes, jtu.rand_default, []),
+    op_record("max", 1, all_dtypes + custom_float_dtypes, all_shapes, jtu.rand_default, []),
+    op_record("min", 1, all_dtypes + custom_float_dtypes, all_shapes, jtu.rand_default, []),
     op_record("nanprod", 1, inexact_dtypes, all_shapes, jtu.rand_small_positive, []),
     op_record("nansum", 1, inexact_dtypes, all_shapes, jtu.rand_default, [],
               tolerance={jnp.bfloat16: 3e-2}),