Repository: google/jax
Issue: jax.numpy.reciprocal/remainder docstring

In https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.reciprocal.html#jax.numpy.reciprocal, the link to the numpy implementation actually self-references the jax version.

Instead, it should point to https://numpy.org/doc/stable/reference/generated/numpy.reciprocal.html

The same applies for remainder: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.remainder.html#jax.numpy.remainder

Changed files:
jax/_src/numpy/lax_numpy.py, jax/_src/numpy/util.py, jax/_src/numpy/ufuncs.py, jax/numpy/__init__.py, jax/_src/scipy/special.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index ec45d646d..d423ed214 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -302,12 +302,12 @@ def load(*args, **kwargs):
 
 ### implementations of numpy functions in terms of lax
 
-@_wraps(np.fmin)
+@_wraps(np.fmin, module='numpy')
 @jit
 def fmin(x1, x2):
   return where((x1 < x2) | isnan(x2), x1, x2)
 
-@_wraps(np.fmax)
+@_wraps(np.fmax, module='numpy')
 @jit
 def fmax(x1, x2):
   return where((x1 > x2) | isnan(x2), x1, x2)
@@ -346,7 +346,7 @@ def trapz(y, x=None, dx=1.0, axis: int = -1):
   return 0.5 * (dx * (y[..., 1:] + y[..., :-1])).sum(-1)
 
 
-@_wraps(np.trunc)
+@_wraps(np.trunc, module='numpy')
 @jit
 def trunc(x):
   _check_arraylike('trunc', x)
@@ -2760,7 +2760,7 @@ def dot(a, b, *, precision=None):  # pylint: disable=missing-docstring
   return lax.dot_general(a, b, (contract_dims, batch_dims), precision)
 
 
-@_wraps(np.matmul, lax_description=_PRECISION_DOC)
+@_wraps(np.matmul, module='numpy', lax_description=_PRECISION_DOC)
 @partial(jit, static_argnames=('precision',), inline=True)
 def matmul(a, b, *, precision=None):  # pylint: disable=missing-docstring
   _check_arraylike("matmul", a, b)
@@ -4087,7 +4087,7 @@ def _gcd_body_fn(xs):
             where(x2 != 0, lax.rem(x1, x2), _lax_const(x2, 0)))
   return (where(x1 < x2, x2, x1), where(x1 < x2, x1, x2))
 
-@_wraps(np.gcd)
+@_wraps(np.gcd, module='numpy')
 @jit
 def gcd(x1, x2):
   _check_arraylike("gcd", x1, x2)
@@ -4100,7 +4100,7 @@ def gcd(x1, x2):
   return gcd
 
 
-@_wraps(np.lcm)
+@_wraps(np.lcm, module='numpy')
 @jit
 def lcm(x1, x2):
   _check_arraylike("lcm", x1, x2)
@@ -4580,8 +4580,8 @@ _NOT_IMPLEMENTED_DESC = """
 *** This function is not yet implemented by jax.numpy, and will raise NotImplementedError ***
 """
 
-def _not_implemented(fun):
-  @_wraps(fun, update_doc=False, lax_description=_NOT_IMPLEMENTED_DESC)
+def _not_implemented(fun, module=None):
+  @_wraps(fun, module=module, update_doc=False, lax_description=_NOT_IMPLEMENTED_DESC)
   def wrapped(*args, **kwargs):
     msg = "Numpy function {} not yet implemented"
     raise NotImplementedError(msg.format(fun))
diff --git a/jax/_src/numpy/ufuncs.py b/jax/_src/numpy/ufuncs.py
index abba9e3ea..bd4a1ca45 100644
--- a/jax/_src/numpy/ufuncs.py
+++ b/jax/_src/numpy/ufuncs.py
@@ -57,9 +57,9 @@ def _one_to_one_unop(numpy_fn, lax_fn, promote_to_inexact=False, lax_doc=False):
   fn = jit(fn, inline=True)
   if lax_doc:
     doc = dedent('\\n\\n'.join(lax_fn.__doc__.split('\\n\\n')[1:])).strip()
-    return _wraps(numpy_fn, lax_description=doc)(fn)
+    return _wraps(numpy_fn, lax_description=doc, module='numpy')(fn)
   else:
-    return _wraps(numpy_fn)(fn)
+    return _wraps(numpy_fn, module='numpy')(fn)
 
 
 def _one_to_one_binop(numpy_fn, lax_fn, promote_to_inexact=False, lax_doc=False):
@@ -70,9 +70,9 @@ def _one_to_one_binop(numpy_fn, lax_fn, promote_to_inexact=False, lax_doc=False)
   fn = jit(fn, inline=True)
   if lax_doc:
     doc = dedent('\\n\\n'.join(lax_fn.__doc__.split('\\n\\n')[1:])).strip()
-    return _wraps(numpy_fn, lax_description=doc)(fn)
+    return _wraps(numpy_fn, lax_description=doc, module='numpy')(fn)
   else:
-    return _wraps(numpy_fn)(fn)
+    return _wraps(numpy_fn, module='numpy')(fn)
 
 
 def _maybe_bool_binop(numpy_fn, lax_fn, bool_lax_fn, lax_doc=False):
@@ -82,9 +82,9 @@ def _maybe_bool_binop(numpy_fn, lax_fn, bool_lax_fn, lax_doc=False):
   fn = jit(fn, inline=True)
   if lax_doc:
     doc = dedent('\\n\\n'.join(lax_fn.__doc__.split('\\n\\n')[1:])).strip()
-    return _wraps(numpy_fn, lax_description=doc)(fn)
+    return _wraps(numpy_fn, lax_description=doc, module='numpy')(fn)
   else:
-    return _wraps(numpy_fn)(fn)
+    return _wraps(numpy_fn, module='numpy')(fn)
 
 
 def _comparison_op(numpy_fn, lax_fn):
@@ -100,11 +100,11 @@ def _comparison_op(numpy_fn, lax_fn):
       return lax.select(lax.eq(rx, ry), lax_fn(lax.imag(x1), lax.imag(x2)),
                         lax_fn(rx, ry))
     return lax_fn(x1, x2)
-  return _wraps(numpy_fn)(fn)
+  return _wraps(numpy_fn, module='numpy')(fn)
 
 
 def _logical_op(np_op, bitwise_op):
-  @_wraps(np_op, update_doc=False)
+  @_wraps(np_op, update_doc=False, module='numpy')
   @partial(jit, inline=True)
   def op(*args):
     zero = lambda x: lax.full_like(x, shape=(), fill_value=0)
@@ -165,7 +165,7 @@ logical_or = _logical_op(np.logical_or, lax.bitwise_or)
 logical_xor = _logical_op(np.logical_xor, lax.bitwise_xor)
 
 
-@_wraps(np.arccosh)
+@_wraps(np.arccosh, module='numpy')
 @jit
 def arccosh(x):
   # Note: arccosh is multi-valued for complex input, and lax.acosh uses a different
@@ -176,7 +176,7 @@ def arccosh(x):
   return out
 
 
-@_wraps(np.right_shift)
+@_wraps(np.right_shift, module='numpy')
 @partial(jit, inline=True)
 def right_shift(x1, x2):
   x1, x2 = _promote_args(np.right_shift.__name__, x1, x2)
@@ -185,16 +185,16 @@ def right_shift(x1, x2):
   return lax_fn(x1, x2)
 
 
-@_wraps(np.absolute)
+@_wraps(np.absolute, module='numpy')
 @partial(jit, inline=True)
 def absolute(x):
   _check_arraylike('absolute', x)
   dt = dtypes.dtype(x)
   return x if dt == np.bool_ or dtypes.issubdtype(dt, np.unsignedinteger) else lax.abs(x)
-abs = _wraps(np.abs)(absolute)
+abs = _wraps(np.abs, module='numpy')(absolute)
 
 
-@_wraps(np.rint)
+@_wraps(np.rint, module='numpy')
 @jit
 def rint(x):
   _check_arraylike('rint', x)
@@ -206,7 +206,7 @@ def rint(x):
   return lax.round(x, lax.RoundingMethod.TO_NEAREST_EVEN)
 
 
-@_wraps(np.sign)
+@_wraps(np.sign, module='numpy')
 @jit
 def sign(x):
   _check_arraylike('sign', x)
@@ -218,7 +218,7 @@ def sign(x):
   return lax.sign(x)
 
 
-@_wraps(np.copysign)
+@_wraps(np.copysign, module='numpy')
 @jit
 def copysign(x1, x2):
   x1, x2 = _promote_args_inexact("copysign", x1, x2)
@@ -227,7 +227,7 @@ def copysign(x1, x2):
   return _where(signbit(x2).astype(bool), -lax.abs(x1), lax.abs(x1))
 
 
-@_wraps(np.true_divide)
+@_wraps(np.true_divide, module='numpy')
 @partial(jit, inline=True)
 def true_divide(x1, x2):
   x1, x2 = _promote_args_inexact("true_divide", x1, x2)
@@ -236,7 +236,7 @@ def true_divide(x1, x2):
 divide = true_divide
 
 
-@_wraps(np.floor_divide)
+@_wraps(np.floor_divide, module='numpy')
 @jit
 def floor_divide(x1, x2):
   x1, x2 = _promote_args("floor_divide", x1, x2)
@@ -261,7 +261,7 @@ def floor_divide(x1, x2):
     return _float_divmod(x1, x2)[0]
 
 
-@_wraps(np.divmod)
+@_wraps(np.divmod, module='numpy')
 @jit
 def divmod(x1, x2):
   x1, x2 = _promote_args("divmod", x1, x2)
@@ -305,7 +305,7 @@ def _power(x1, x2):
   return acc
 
 
-@_wraps(np.power)
+@_wraps(np.power, module='numpy')
 def power(x1, x2):
   # Special case for concrete integer scalars: use binary exponentiation.
   # Using lax.pow may be imprecise for floating-point values; the goal of this
@@ -322,7 +322,7 @@ def power(x1, x2):
 
 
 @custom_jvp
-@_wraps(np.logaddexp)
+@_wraps(np.logaddexp, module='numpy')
 @jit
 def logaddexp(x1, x2):
   x1, x2 = _promote_args_inexact("logaddexp", x1, x2)
@@ -360,7 +360,7 @@ def _logaddexp_jvp(primals, tangents):
 
 
 @custom_jvp
-@_wraps(np.logaddexp2)
+@_wraps(np.logaddexp2, module='numpy')
 @jit
 def logaddexp2(x1, x2):
   x1, x2 = _promote_args_inexact("logaddexp2", x1, x2)
@@ -388,28 +388,28 @@ def _logaddexp2_jvp(primals, tangents):
   return primal_out, tangent_out
 
 
-@_wraps(np.log2)
+@_wraps(np.log2, module='numpy')
 @partial(jit, inline=True)
 def log2(x):
   x, = _promote_args_inexact("log2", x)
   return lax.div(lax.log(x), lax.log(_constant_like(x, 2)))
 
 
-@_wraps(np.log10)
+@_wraps(np.log10, module='numpy')
 @partial(jit, inline=True)
 def log10(x):
   x, = _promote_args_inexact("log10", x)
   return lax.div(lax.log(x), lax.log(_constant_like(x, 10)))
 
 
-@_wraps(np.exp2)
+@_wraps(np.exp2, module='numpy')
 @partial(jit, inline=True)
 def exp2(x):
   x, = _promote_args_inexact("exp2", x)
   return lax.exp(lax.mul(lax.log(_constant_like(x, 2)), x))
 
 
-@_wraps(np.signbit)
+@_wraps(np.signbit, module='numpy')
 @jit
 def signbit(x):
   x, = _promote_args("signbit", x)
@@ -446,7 +446,7 @@ def _normalize_float(x):
   return lax.bitcast_convert_type(x1, int_type), x2
 
 
-@_wraps(np.ldexp)
+@_wraps(np.ldexp, module='numpy')
 @jit
 def ldexp(x1, x2):
   _check_arraylike("ldexp", x1, x2)
@@ -495,7 +495,7 @@ def ldexp(x1, x2):
   return _where(isinf(x1) | isnan(x1) | (x1 == 0), x1, x)
 
 
-@_wraps(np.frexp)
+@_wraps(np.frexp, module='numpy')
 @jit
 def frexp(x):
   _check_arraylike("frexp", x)
@@ -519,7 +519,7 @@ def frexp(x):
   return _where(cond, x, x1), lax.convert_element_type(x2, np.int32)
 
 
-@_wraps(np.remainder)
+@_wraps(np.remainder, module='numpy')
 @jit
 def remainder(x1, x2):
   x1, x2 = _promote_args("remainder", x1, x2)
@@ -529,10 +529,10 @@ def remainder(x1, x2):
   do_plus = lax.bitwise_and(
       lax.ne(lax.lt(trunc_mod, zero), lax.lt(x2, zero)), trunc_mod_not_zero)
   return lax.select(do_plus, lax.add(trunc_mod, x2), trunc_mod)
-mod = _wraps(np.mod)(remainder)
+mod = _wraps(np.mod, module='numpy')(remainder)
 
 
-@_wraps(np.fmod)
+@_wraps(np.fmod, module='numpy')
 @jit
 def fmod(x1, x2):
   _check_arraylike("fmod", x1, x2)
@@ -541,21 +541,21 @@ def fmod(x1, x2):
   return lax.rem(*_promote_args("fmod", x1, x2))
 
 
-@_wraps(np.square)
+@_wraps(np.square, module='numpy')
 @partial(jit, inline=True)
 def square(x):
   _check_arraylike("square", x)
   return lax.integer_pow(x, 2)
 
 
-@_wraps(np.deg2rad)
+@_wraps(np.deg2rad, module='numpy')
 @partial(jit, inline=True)
 def deg2rad(x):
   x, = _promote_args_inexact("deg2rad", x)
   return lax.mul(x, _lax_const(x, np.pi / 180))
 
 
-@_wraps(np.rad2deg)
+@_wraps(np.rad2deg, module='numpy')
 @partial(jit, inline=True)
 def rad2deg(x):
   x, = _promote_args_inexact("rad2deg", x)
@@ -566,7 +566,7 @@ degrees = rad2deg
 radians = deg2rad
 
 
-@_wraps(np.conjugate)
+@_wraps(np.conjugate, module='numpy')
 @partial(jit, inline=True)
 def conjugate(x):
   _check_arraylike("conjugate", x)
@@ -587,7 +587,7 @@ def real(val):
   _check_arraylike("real", val)
   return lax.real(val) if np.iscomplexobj(val) else val
 
-@_wraps(np.modf, skip_params=['out'])
+@_wraps(np.modf, module='numpy', skip_params=['out'])
 @jit
 def modf(x, out=None):
   _check_arraylike("modf", x)
@@ -598,7 +598,7 @@ def modf(x, out=None):
   return x - whole, whole
 
 
-@_wraps(np.isfinite)
+@_wraps(np.isfinite, module='numpy')
 @jit
 def isfinite(x):
   _check_arraylike("isfinite", x)
@@ -611,7 +611,7 @@ def isfinite(x):
     return lax.full_like(x, True, dtype=np.bool_)
 
 
-@_wraps(np.isinf)
+@_wraps(np.isinf, module='numpy')
 @jit
 def isinf(x):
   _check_arraylike("isinf", x)
@@ -649,14 +649,14 @@ isneginf = _wraps(np.isneginf, skip_params=['out'])(
 )
 
 
-@_wraps(np.isnan)
+@_wraps(np.isnan, module='numpy')
 @jit
 def isnan(x):
   _check_arraylike("isnan", x)
   return lax.ne(x, x)
 
 
-@_wraps(np.heaviside)
+@_wraps(np.heaviside, module='numpy')
 @jit
 def heaviside(x1, x2):
   _check_arraylike("heaviside", x1, x2)
@@ -666,7 +666,7 @@ def heaviside(x1, x2):
                 _where(lax.gt(x1, zero), _lax_const(x1, 1), x2))
 
 
-@_wraps(np.hypot)
+@_wraps(np.hypot, module='numpy')
 @jit
 def hypot(x1, x2):
   _check_arraylike("hypot", x1, x2)
@@ -677,7 +677,7 @@ def hypot(x1, x2):
   return lax.select(x1 == 0, x1, x1 * lax.sqrt(1 + lax.square(lax.div(x2, lax.select(x1 == 0, lax_internal._ones(x1), x1)))))
 
 
-@_wraps(np.reciprocal)
+@_wraps(np.reciprocal, module='numpy')
 @partial(jit, inline=True)
 def reciprocal(x):
   _check_arraylike("reciprocal", x)
diff --git a/jax/_src/numpy/util.py b/jax/_src/numpy/util.py
index 4b0223bcd..e6b395bd6 100644
--- a/jax/_src/numpy/util.py
+++ b/jax/_src/numpy/util.py
@@ -124,6 +124,7 @@ def _wraps(
     sections: Sequence[str] = ('Parameters', 'Returns', 'References'),
     skip_params: Sequence[str] = (),
     extra_params: Optional[str] = None,
+    module: Optional[str] = None,
 ) -> Callable[[_T], _T]:
   """Specialized version of functools.wraps for wrapping numpy functions.
 
@@ -147,13 +148,20 @@ def _wraps(
     extra_params: an optional string containing additional parameter descriptions.
       When ``update_doc=True``, these will be added to the list of parameter
       descriptions in the updated doc.
+    module: an optional string specifying the module from which the wrapped function
+      is imported. This is useful for objects such as ufuncs, where the module cannot
+      be determined from the wrapped function itself.
   """
   def wrap(op):
     docstr = getattr(fun, "__doc__", None)
+    name = getattr(fun, "__name__", getattr(op, "__name__", str(op)))
     try:
-      name = f"{fun.__module__}.{fun.__name__}"
+      mod = module or fun.__module__
     except AttributeError:
-      name = getattr(fun, "__name__", getattr(op, "__name__", str(op)))
+      if config.jax_enable_checks:
+        raise ValueError(f"function {fun} defines no __module__; pass module keyword to _wraps.")
+    else:
+      name = f"{mod}.{name}"
     if docstr:
       try:
         parsed = _parse_numpydoc(docstr)
diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
index 5cc73fb49..7524721d5 100644
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -29,25 +29,25 @@ from jax._src.numpy.lax_numpy import _reduction_dims, _promote_args_inexact
 from jax._src.numpy.util import _wraps
 
 
-@_wraps(osp_special.gammaln)
+@_wraps(osp_special.gammaln, module='scipy.special')
 def gammaln(x):
   x, = _promote_args_inexact("gammaln", x)
   return lax.lgamma(x)
 
 
-@_wraps(osp_special.betaln)
+@_wraps(osp_special.betaln, module='scipy.special')
 def betaln(x, y):
   x, y = _promote_args_inexact("betaln", x, y)
   return lax.lgamma(x) + lax.lgamma(y) - lax.lgamma(x + y)
 
 
-@_wraps(osp_special.betainc)
+@_wraps(osp_special.betainc, module='scipy.special')
 def betainc(a, b, x):
   a, b, x = _promote_args_inexact("betainc", a, b, x)
   return lax.betainc(a, b, x)
 
 
-@_wraps(osp_special.digamma, lax_description="""\\
+@_wraps(osp_special.digamma, module='scipy.special', lax_description="""\\
 The JAX version only accepts real-valued inputs.""")
 def digamma(x):
   x, = _promote_args_inexact("digamma", x)
@@ -57,38 +57,38 @@ ad.defjvp(
     lambda g, x: lax.mul(g, polygamma(1, x)))  # type: ignore[has-type]
 
 
-@_wraps(osp_special.gammainc, update_doc=False)
+@_wraps(osp_special.gammainc, module='scipy.special', update_doc=False)
 def gammainc(a, x):
   a, x = _promote_args_inexact("gammainc", a, x)
   return lax.igamma(a, x)
 
 
-@_wraps(osp_special.gammaincc, update_doc=False)
+@_wraps(osp_special.gammaincc, module='scipy.special', update_doc=False)
 def gammaincc(a, x):
   a, x = _promote_args_inexact("gammaincc", a, x)
   return lax.igammac(a, x)
 
 
-@_wraps(osp_special.erf)
+@_wraps(osp_special.erf, module='scipy.special')
 def erf(x):
   x, = _promote_args_inexact("erf", x)
   return lax.erf(x)
 
 
-@_wraps(osp_special.erfc, update_doc=False)
+@_wraps(osp_special.erfc, module='scipy.special', update_doc=False)
 def erfc(x):
   x, = _promote_args_inexact("erfc", x)
   return lax.erfc(x)
 
 
-@_wraps(osp_special.erfinv)
+@_wraps(osp_special.erfinv, module='scipy.special')
 def erfinv(x):
   x, = _promote_args_inexact("erfinv", x)
   return lax.erf_inv(x)
 
 
 @api.custom_jvp
-@_wraps(osp_special.logit, update_doc=False)
+@_wraps(osp_special.logit, module='scipy.special', update_doc=False)
 def logit(x):
   x, = _promote_args_inexact("logit", x)
   return lax.log(lax.div(x, lax.sub(_lax_const(x, 1), x)))
@@ -97,7 +97,7 @@ logit.defjvps(
 
 
 @api.custom_jvp
-@_wraps(osp_special.expit, update_doc=False)
+@_wraps(osp_special.expit, module='scipy.special', update_doc=False)
 def expit(x):
   x, = _promote_args_inexact("expit", x)
   one = _lax_const(x, 1)
@@ -105,7 +105,7 @@ def expit(x):
 expit.defjvps(lambda g, ans, x: g * ans * (_lax_const(ans, 1) - ans))
 
 
-@_wraps(osp_special.logsumexp)
+@_wraps(osp_special.logsumexp, module='scipy.special')
 def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
   if b is not None:
     a, b = _promote_args_inexact("logsumexp", a, b)
@@ -145,7 +145,7 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
   return out
 
 
-@_wraps(osp_special.xlogy)
+@_wraps(osp_special.xlogy, module='scipy.special')
 def xlogy(x, y):
   x, y = _promote_args_inexact("xlogy", x, y)
   x_ok = x != 0.
@@ -154,7 +154,7 @@ def xlogy(x, y):
   return jnp.where(x_ok, lax.mul(safe_x, lax.log(safe_y)), jnp.zeros_like(x))
 
 
-@_wraps(osp_special.xlog1py, update_doc=False)
+@_wraps(osp_special.xlog1py, module='scipy.special', update_doc=False)
 def xlog1py(x, y):
   x, y = _promote_args_inexact("xlog1py", x, y)
   x_ok = x != 0.
@@ -163,7 +163,7 @@ def xlog1py(x, y):
   return jnp.where(x_ok, lax.mul(safe_x, lax.log1p(safe_y)), jnp.zeros_like(x))
 
 
-@_wraps(osp_special.entr)
+@_wraps(osp_special.entr, module='scipy.special')
 def entr(x):
   x, = _promote_args_inexact("entr", x)
   return lax.select(lax.lt(x, _lax_const(x, 0)),
@@ -208,7 +208,7 @@ _BERNOULLI_COEFS = [
 ]
 
 
-@_wraps(osp_special.zeta)
+@_wraps(osp_special.zeta, module='scipy.special')
 def zeta(x, q=None):
   assert q is not None, "Riemann zeta function is not implemented yet."
   # Reference: Johansson, Fredrik.
@@ -237,7 +237,7 @@ def zeta(x, q=None):
   return S + I + T
 
 
-@_wraps(osp_special.polygamma, update_doc=False)
+@_wraps(osp_special.polygamma, module='scipy.special', update_doc=False)
 def polygamma(n, x):
   assert jnp.issubdtype(lax.dtype(n), jnp.integer)
   n, x = _promote_args_inexact("polygamma", n, x)
@@ -657,22 +657,22 @@ def _norm_logpdf(x):
   log_normalizer = _lax_const(x, _norm_logpdf_constant)
   return lax.sub(lax.mul(neg_half, lax.square(x)), log_normalizer)
 
-@_wraps(osp_special.i0e)
+@_wraps(osp_special.i0e, module='scipy.special')
 def i0e(x):
   x, = _promote_args_inexact("i0e", x)
   return lax.bessel_i0e(x)
 
-@_wraps(osp_special.i0)
+@_wraps(osp_special.i0, module='scipy.special')
 def i0(x):
   x, = _promote_args_inexact("i0", x)
   return lax.mul(lax.exp(lax.abs(x)), lax.bessel_i0e(x))
 
-@_wraps(osp_special.i1e)
+@_wraps(osp_special.i1e, module='scipy.special')
 def i1e(x):
   x, = _promote_args_inexact("i1e", x)
   return lax.bessel_i1e(x)
 
-@_wraps(osp_special.i1)
+@_wraps(osp_special.i1, module='scipy.special')
 def i1(x):
   x, = _promote_args_inexact("i1", x)
   return lax.mul(lax.exp(lax.abs(x)), lax.bessel_i1e(x))
@@ -1307,7 +1307,7 @@ def _expi_pos(x):
   )
 
 
-@_wraps(osp_special.expi)
+@_wraps(osp_special.expi, module='scipy.special')
 @api.custom_jvp
 @jit
 def expi(x):
@@ -1426,7 +1426,7 @@ def _expn3(n, x):
   return (ans + one) * jnp.exp(-x) / xk
 
 
-@_wraps(osp_special.expn)
+@_wraps(osp_special.expn, module='scipy.special')
 @partial(api.custom_jvp, nondiff_argnums=(0,))
 @jnp.vectorize
 @jit
@@ -1466,7 +1466,7 @@ def expn_jvp(n, primals, tangents):
   )
 
 
-@_wraps(osp_special.exp1)
-def exp1(x):
+@_wraps(osp_special.exp1, module="scipy.special")
+def exp1(x, module='scipy.special'):
   (x,) = _promote_args_inexact("exp1", x)
   return expn(1, x)
diff --git a/jax/numpy/__init__.py b/jax/numpy/__init__.py
index dc2028821..e74c1165c 100644
--- a/jax/numpy/__init__.py
+++ b/jax/numpy/__init__.py
@@ -423,7 +423,7 @@ def _init():
   for name, func in util.get_module_functions(np).items():
     if name not in globals():
       _NOT_IMPLEMENTED.append(name)
-      globals()[name] = lax_numpy._not_implemented(func)
+      globals()[name] = lax_numpy._not_implemented(func, module='numpy')
 
 _init()
 del _init