mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Minor improvements to doc for jax.nn.logsumexp.

Browse Source
This commit is contained in:
carlosgmartin 2025-04-13 15:17:11 -04:00
parent 773b323b26
commit 2336cd1695
1 changed files with 3 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
jax/_src/ops/special.py
View File

@ -47,16 +47,15 @@ def logsumexp(a: ArrayLike, axis: Axis = None, b: ArrayLike | None = None,
JAX implementation of :func:`scipy.special.logsumexp`.
.. math::
\mathrm{logsumexp}(a) = \mathrm{log} \sum_j b \cdot \mathrm{exp}(a_{ij})
\operatorname{logsumexp} a = \log \sum_i b_i \exp a_i
where the :math:`j` indices range over one or more dimensions to be reduced.
where the :math:`i` indices range over one or more dimensions to be reduced.
Args:
a: the input array
axis: int or sequence of ints, default=None. Axis along which the sum to be
computed. If None, the sum is computed along all the axes.
b: scaling factors for :math:`\mathrm{exp}(a)`. Must be broadcastable to the
shape of `a`.
b: scaling factors for the exponentials. Must be broadcastable to the shape of `a`.
keepdims: If ``True``, the axes that are reduced are left in the output as
dimensions of size 1.
return_sign: If ``True``, the output will be a ``(result, sign)`` pair,