[pmap] Push more into _cached_shard_map.

NOTE: We aren't using all the fields of CachedShardMap yet. Subsequent CLs will make use of them.

Improving the `jax.jit(jax.shard_map)` implementation of `jax.pmap`.

PiperOrigin-RevId: 861952422

Author: danielsuo

tensorflow_probability/python/distributions/jax_transformation_test.py

@@ -75,6 +75,7 @@
 PMAP_SAMPLE_BLOCKLIST = (
     'Bates',
     'BatchReshape',  # http://b/163171224
+    'IncrementLogProb',  # pmap_shmap_merge sharding conflict
     'MixtureSameFamily',  # Too slow: http://b/170871051
     'NegativeBinomial',  # Too slow: http://b/170871051
     'ZeroInflatedNegativeBinomial',  # Too slow: http://b/170871051