code: def testPmean(self):
      if extensions.tpu_devices():
          self.skipTest("pmean for TPU is not supported yet")
      devices = self._get_two_devices(require_same_type=True)
  
      def reduce_mean(f):
          return extensions.pmean(f)
  
      data = tf_np.asarray(tf.convert_to_tensor(value=[1, 3]))
      pmapped = extensions.pmap(reduce_mean, devices=devices)
      result = pmapped(data)
  
      self.assertAllClose(result[0], 2)
      self.assertAllClose(result[1], 2)
  
imports: import functools
  import itertools
  from absl import flags
  from absl.testing import parameterized
  from jax import lax
  import numpy
  import tensorflow.compat.v2
  from trax.tf_numpy import extensions
  import trax.tf_numpy.numpy
label_desc: Use image node_id 39 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testPmean() and returns: extensions