de-duplicate util.memoize and util.cache
de-duplicate util.memoize and util.cache

The only difference between the two was that
jax.config.jax_check_tracer_leaks disables the caching under util.cache
but not under util.memoize.

We could add that as an option on the same function if it turns out to
be important, but it seems unnecessary. Moreover there are only two
callers (in dtypes.py and in batching.py).

Co-authored-by: Skye Wanderman-Milne <4c878c1ae364587078acae22009e8542fed717b6@google.com>

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