jax.scipy.special.digamma returns NaN for signed zero and positive subnormal float32 inputs

### Description `jax.scipy.special.digamma` appears to return incorrect `NaN` results for `float32` signed zero and the smallest positive subnormal input. For `digamma(x)`, values at zero should follow the pole behavior. In particular, `digamma(+0.0)` should be `-inf`, and `digamma(-0.0)` should be `+inf` due to the signed-zero side of the pole. For a very small positive nonzero `float32` input, the result should also tend toward `-inf`. However, JAX returns `NaN` for these cases. ### Minimal reproducible example ``` #!/usr/bin/env python3 import os import sys import warnings warnings.filterwarnings("ignore") os.environ.setdefault("JAX_PLATFORMS", "cpu") try: import numpy as np import jax import jax.numpy as jnp import jax.scipy.special as jsp except ImportError as e: print(f"missing dep: {e}") sys.exit(2) nan = float("nan") inf = float("inf") x = jnp.asarray( np.asarray( [ [-0.0, 0.0, 1.401298464324817e-45], [-1.0, 1.0, 9.999999974752427e-07], ], dtype=np.float32, ) ) out = np.asarray(jsp.digamma(x)) expected = np.asarray( [ [inf, -inf, -inf], [nan, -0.5772156715393066, -1000000.5625], ], dtype=np.float32, ) print("jax:", out.tolist()) print("expected:", expected.tolist()) ``` ### Actual result ``` jax: [[nan, nan, nan], [nan, -0.5772153735160828, -1000000.5625]] expected: [[inf, -inf, -inf], [nan, -0.5772156715393066, -1000000.5625]] ``` ### Expected result ``` digamma(-0.0) should be +inf digamma(+0.0) should be -inf digamma(1.401298464324817e-45) should be -inf ``` The `NaN` result for `-1.0` is expected because negative integers are poles of the digamma function. The finite values for `1.0` and `9.999999974752427e-07` are also consistent with the expected result. The issue is specifically that JAX returns `NaN` for signed zero and the smallest positive subnormal `float32` input where infinities are expected. ### 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.