code: def testPmap(self):
      devices = self._get_two_devices()
  
      @functools.partial(extensions.pmap, devices=devices)
      def return_three(f):
          return f, f + 1.0, f + 2.0
  
      result = return_three(tf.ones((2, 20)))
      # The function returned 3 items, so we got 3 items back.
      self.assertLen(result, 3)
  
      # Each of the items should be a ShardedNdarray that when converted to tensor
      # should produce a tensor of shape (2, 20)
      converted = tf.nest.map_structure(tf.convert_to_tensor, result)
  
      self.assertLen(result, 3)
  
      self.assertAllEqual(converted[0].shape, converted[1].shape)
      self.assertAllEqual(converted[0].shape, converted[2].shape)
  
      self.assertAllEqual(converted[0], tf.ones((2, 20)))
      self.assertAllEqual(converted[1], 1 + tf.ones((2, 20)))
      self.assertAllEqual(converted[2], 2 + tf.ones((2, 20)))
  
      @functools.partial(extensions.pmap, devices=devices)
      def return_one(f):
          return f + 2.0
  
      result = return_one(tf.ones((2, 20)))
  
      # Only a single item is returned, so we can convert it directly.
      converted = tf.convert_to_tensor(value=result)
      self.assertAllEqual(converted, 2 + tf.ones((2, 20)))
  
      @functools.partial(extensions.pmap, devices=devices)
      def return_list(f):
          return [f + 2.0]
  
      result = return_list(tf.ones((2, 20)))
  
      # A singleton list is returned.
      self.assertLen(result, 1)
      converted = tf.convert_to_tensor(value=result[0])
      self.assertAllEqual(converted, 2 + tf.ones((2, 20)))
  
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 34 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testPmap() and returns: f,  unknown,  unknown, unknown, list