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

Merge pull request #26862 from jakevdp:logsumexp-docs

Browse Source 
PiperOrigin-RevId: 733080943
This commit is contained in:
jax authors 2025-03-03 15:24:10 -08:00
commit f9f47217df
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -53,8 +53,8 @@ def logsumexp(a: ArrayLike, axis: Axis = None, b: ArrayLike | None = None,
Args:
a: the input array
axis: the axis or axes over which to reduce. May be either ``None``, an
int, or a tuple of ints.
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`.
keepdims: If ``True``, the axes that are reduced are left in the output as