Rename _wraps to implements

[jakevdp]

I've been wanting to change this for a while, because "wraps" is really a misnomer here: wraps implies that the decorated function somehow calls out to the wrapped function, but that is not the case for JAX implementations of NumPy functions. I think implements is more accurate for this case.

Also remove the leading underscore: it's misleading because we expect this function to be used outside the module that defines it (it's still defined in jax._src, however, so it's not a public API).

[mattjj]

Thank you!

[jecampagne]

Hi, does this @implement decorator would be in a future put outside _src with an API description, or would remain a sort of "private" deco? I ask as we use this deco for docstring description of functions that are our JAX implementation of a library.

[jakevdp]

No, I don't think we'll ever make @implements a public API. JAX is an accelerated linear algebra library, not a library of decorators to help generate docstrings.

I ask as we use this deco for docstring description of functions that are our JAX implementation of a library.

As with all JAX's internal APIs, I would not recommend using it outside JAX internal code, as it may be changed or removed in future releases without warning. In fact we're actively moving away from using @implements even internally (see #21461), and my goal would be to remove it entirely in the near future.