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

Fix incorrect wrapped docstring of jax.scipy.special.gamma

Browse Source 
Fixes the docstring `jax.scipy.special.gamma`, which was wrapping `scipy.special.gammaln` by mistake. Also adds a note that the function currently only accepts real inputs.
This commit is contained in:
Nicolas Tessore 2023-06-01 20:13:37 +01:00
parent 5c2070c204
commit a835cafdad
No known key found for this signature in database
GPG Key ID: 6110AD9819080235
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jax/_src/scipy/special.py
View File

@ -42,7 +42,8 @@ def gammaln(x: ArrayLike) -> Array:
return lax.lgamma(x)
@_wraps(osp_special.gammaln, module='scipy.special')
@_wraps(osp_special.gamma, module='scipy.special', lax_description="""\
The JAX version only accepts real-valued inputs.""")
def gamma(x: ArrayLike) -> Array:
x, = promote_args_inexact("gamma", x)
return lax.exp(lax.lgamma(x))