Belt-and-braces check for skipping pmap variants when we only have a single device.
Belt-and-braces check for skipping pmap variants when we only have a single device.

PiperOrigin-RevId: 356561713
Change-Id: I6631edc7d5e956f746d2e661d5fc6dad2d7511df

deepmind/chex,deepmind/chex