fix pmap test on GPU/TPU (#3682)
fix pmap test on GPU/TPU (#3682)

* fix pmap test on GPU/TPU

* AssertionError -> ValueError
google/jax,tensorflow/probability,google/jax,google/jax,google/jax,tensorflow/probability