fix multi-host pmap, disambiguate nrep
fix multi-host pmap, disambiguate nrep

When #1667 inlined a function into its caller, it mixed up two distinct
values referred to as `nrep` in the two functions: num_global_replicas
vs num_local_replicas. The result caused errors on multi-host setups.

Co-authored-by: Jonathan Heek <9b4127d7610427859afd8c31ac199b7fb82d3f75@google.com>

google/jax,google/jax,tensorflow/probability,tensorflow/probability,google/jax,google/jax