ENH _wraps() in jax._src.numpy.util is now returns a generic function
ENH _wraps() in jax._src.numpy.util is now returns a generic function

This frees type checkers from the need to explicitly infer the return type
of _wraps() at each call site.

google/jax,google/jax,google/jax,google/jax