Make eager pmap take advantage of pmap cache
Make eager pmap take advantage of pmap cache

The current strategy of creating a `partial(primitive.bind, **params)` has the downside
of completely confusing the pmap cache and resulting in a new compilation for every single
primitive. Replacing it with a `HashableFunction` should fix it.

Also, pmap_test is now 2x faster!

PiperOrigin-RevId: 490749153

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