code: def testPmap(self):
      pmap_funcs = [api._cpp_pmap]
  
      for pmap in pmap_funcs:
          f = pmap(lambda x: 0.0 / x)
          # For the Cpp pmap, the first execution always goes through Python.
          f(jnp.array([1.0]))
  
          with self.assertRaisesRegex(
              FloatingPointError,
              r"invalid value \(nan\) encountered in parallel computation",
          ):
              ans = f(jnp.array([0.0]))
              ans.block_until_ready()
  
          if jax.device_count() >= 2:
              with self.assertRaisesRegex(
                  FloatingPointError,
                  r"invalid value \(nan\) encountered in parallel computation",
              ):
                  ans = f(jnp.array([1.0, 0.0]))
                  ans.block_until_ready()
  
imports: from absl.testing import absltest
  import jax
  import numpy
  from unittest import SkipTest
  from jax._src import api
  from jax._src import test_util
  from jax import numpy
  from jax.experimental import pjit, maps
  from jax import config
label_desc: Use image node_id 11 for calling the DebugNaNsTest obj's underlying member method code with example usage: obj.testPmap() without return types