Merge pull request #20822 from jakevdp:scipy-special-doc

PiperOrigin-RevId: 629129623

docs/conf.py

@@ -301,8 +301,8 @@ always_document_param_types = True
 
 # Tell sphinx autodoc how to render type aliases.
 autodoc_type_aliases = {
-    'ArrayLike': 'ArrayLike',
-    'DTypeLike': 'DTypeLike',
+    'ArrayLike': 'jax.typing.ArrayLike',
+    'DTypeLike': 'jax.typing.DTypeLike',
 }
 
 

jax/_src/ops/special.py

@@ -44,7 +44,7 @@ def logsumexp(a: ArrayLike, axis: Axis = None, b: ArrayLike | None = None,
               keepdims: bool = False, return_sign: bool = False, where: ArrayLike | None = None) -> Array | tuple[Array, Array]:
   r"""Log-sum-exp reduction.
 
-  Computes
+  JAX implementation of :func:`scipy.special.logsumexp`.
 
   .. math::
     \mathrm{logsumexp}(a) = \mathrm{log} \sum_j b \cdot \mathrm{exp}(a_{ij})