jax.nn.selu gradient returns zero for large negative float64 input

### Description `jax.nn.selu` appears to return an incorrect zero gradient for a large negative finite `float64` input. For the negative branch of SELU, the derivative is proportional to `exp(x)`. At `x = -700.0`, the true derivative is approximately `1.7334290832552396e-304`, which is finite and representable in `float64`. However, JAX returns `0.0`. ### Minimal reproducible example ``` #!/usr/bin/env python3 import os import sys import math os.environ.setdefault("CUDA_VISIBLE_DEVICES", "-1") os.environ.setdefault("JAX_PLATFORMS", "cpu") os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3") try: import jax import jax.numpy as jnp except ImportError as e: print(f"missing dep: {e}") sys.exit(2) jax.config.update("jax_enable_x64", True) x = -700.0 expected = 1.7334290832552396e-304 out = float(jax.grad(lambda z: jax.nn.selu(z))(jnp.asarray(x, dtype=jnp.float64))) print("jax:", out) print("expected:", expected) if not (math.isfinite(out) and abs(out - expected) <= 1e-6 * abs(expected)): sys.exit(0) sys.exit(1) ``` ### Actual result ``` jax: 0.0 expected: 1.7334290832552396e-304 ``` ### Expected result ``` 1.7334290832552396e-304 ``` For the negative branch of SELU, the derivative should be a positive scaled `exp(x)` value. At `x = -700.0`, this value is still finite and representable in `float64`, so the backward pass should not flush it to zero. ### System info (python version, jaxlib version, accelerator, etc.) jax: 0.10.1 jaxlib: 0.10.1 numpy: 2.4.6 python: 3.13.13 platform: Ubuntu 24.04.3 LTS, Linux x86_64 accelerator: CPU backend Note: GPU is present, but CUDA-enabled jaxlib is not installed, so JAX falls back to CPU.