code: def testPmapSimpleModel(self):
      devices = self._get_two_devices(require_same_type=True)
      n_devices = len(devices)
  
      (
          params,
          params_true,
          inputs,
          targets,
      ) = generate_params_inputs_targets()
  
      def _train_and_reduce(params, inputs, targets, learning_rate=0.1):
          new_w, new_b = train_step(
              params, inputs, targets, learning_rate
          )
  
          return (
              extensions.psum(new_w) / n_devices,
              extensions.psum(new_b) / n_devices,
          )
  
      train_step_pmapped = extensions.pmap(
          _train_and_reduce, devices=devices
      )
  
      def replicate(x, num_devices=2):
          return tf_np.broadcast_to(x, (num_devices,) + x.shape)
  
      params = tf.nest.map_structure(replicate, params)
  
      def reshape(x, num_devices=2):
          x_shape = list(x.shape)
          batch_size = x_shape[0]
          batch_size_per_device = batch_size // num_devices
  
          # New shape.
          new_shape_prefix = [num_devices, batch_size_per_device]
          return tf_np.reshape(x, new_shape_prefix + x_shape[1:])
  
      inputs = tf.nest.map_structure(reshape, inputs)
      targets = tf.nest.map_structure(reshape, targets)
  
      for _ in range(50):
          params = train_step_pmapped(params, inputs, targets)
  
      # PMAP returns sharded tensors.
  
      # Since the inputs are identical, the returned tensors should be identical
      self.assertAllClose(params[0][0], params[0][1])
      self.assertAllClose(params[1][0], params[1][1])
  
      # This is not trained super well, but it usually gets "close".
      self.assertAllClose(params[0][0], params_true[0], atol=1e-1)
      self.assertAllClose(params[1][0], params_true[1], atol=1e-1)
  
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 36 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testPmapSimpleModel() and returns: unknown,  unknown, tf_np, tf_np