Remove implements decorator

jax-ml · Jun 25, 2024 · 0eb8a85 · 0eb8a85
1 parent 6067820
commit 0eb8a85
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
@@ -213,6 +213,7 @@ def force(x):
 """
 
 
+@partial(api.jit, static_argnames=('axis', 'dtype', 'keepdims', 'promote_integers'), inline=True)
 def _reduce_sum(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
                 out: None = None, keepdims: bool = False,
                 initial: ArrayLike | None = None, where: ArrayLike | None = None,
@@ -223,7 +224,7 @@ def _reduce_sum(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
                     initial=initial, where_=where, parallel_reduce=lax.psum,
                     promote_integers=promote_integers)
 
-@implements(np.sum, skip_params=['out'], extra_params=_PROMOTE_INTEGERS_DOC)
+
 def sum(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
         out: None = None, keepdims: bool = False, initial: ArrayLike | None = None,
         where: ArrayLike | None = None, promote_integers: bool = True) -> Array: