code: def testWrongAxisName(self):
      devices = self._get_two_devices(require_same_type=True)
  
      def reduce_sum(f):
          return extensions.psum(f, axis_name="bar")
  
      data = tf_np.asarray(tf.convert_to_tensor(value=[1, 3]))
      with self.assertRaisesWithPredicateMatch(
          ValueError,
          r"axis_name (.*) is not equal to that of the surrounding",
      ):
          pmapped = extensions.pmap(
              reduce_sum, axis_name="foo", devices=devices
          )
          pmapped(data)
  
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 41 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testWrongAxisName() and returns: extensions