code: def testPsumStruct(self):
      devices = self._get_two_devices(require_same_type=True)
  
      def reduce_sum(a):
          a = extensions.psum(a)
          tf.nest.map_structure(
              lambda x: self.assertIsInstance(x, tf_np.ndarray), a
          )
          return a
  
      data = [tf_np.asarray([1, 3]), tf_np.asarray([2, 4], np.int64)]
      pmapped = extensions.pmap(reduce_sum, devices=devices)
      result = pmapped(data)
  
      self.assertIsInstance(result[0][0], tf_np.ndarray)
      self.assertIsInstance(result[0][1], tf_np.ndarray)
      self.assertIsInstance(result[1][0], tf_np.ndarray)
      self.assertIsInstance(result[1][1], tf_np.ndarray)
      self.assertAllClose(result[0][0], 4)
      self.assertAllClose(result[0][1], 4)
      self.assertAllClose(result[1][0], 6)
      self.assertAllClose(result[1][1], 6)
  
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 38 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testPsumStruct() and returns: a