class_docstr: Tests that inspect the HLO for the sharding annotations.
    
class_docstr_tok: ["Tests","that","inspect","the","HLO","for","the","sharding","annotations","."]
code: def test_xmap_collective_reduce(self):
      devices = np.reshape(self.devices, (1, 2))
      ashape = (16, 8, 5)
      a = np.arange(np.prod(ashape), dtype=np.float32).reshape(ashape)
      bshape = (2, 7)
      b = np.arange(np.prod(bshape), dtype=np.float32).reshape(bshape)
      f_jax = xmap(
          lambda a, b: (lax.psum(a * 2.0, "a"), b * 4.0),
          in_axes=(["a", "b", ...], {0: "c"}),
          out_axes=(["b", ...], {0: "c"}),
          axis_resources={"a": "x", "b": "y", "c": "x"},
      )
  
      @tf.function(autograph=False, jit_compile=True)
      def f_tf(a, b):
          f_converted = jax2tf.convert(f_jax, native_serialization=True)
          if jtu.test_device_matches(["tpu"]):
              res = tf.compat.v1.tpu.rewrite(
                  f_converted,
                  [tf.convert_to_tensor(a), tf.convert_to_tensor(b)],
                  device_assignment=self.device_assignment(
                      computation_shape=[1, 1, 1, 2]
                  ),
              )
              return (res[0], res[1])
          else:
              return f_converted(a, b)
  
      with Mesh(devices, ("x", "y")):
          res_jax = f_jax(a, b)
          self.assertAllClose(res_jax, ((a * 2.0).sum(0), b * 4.0))
          res_tf = f_tf(a, b)
          self.assertAllClose(res_tf, res_jax)
          self.check_sharding(
              jax2tf.convert(f_jax, native_serialization=True),
              [a, b],
              checks=[
                  (
                      r"f32\[16,8,5\].*custom_call_target.*Sharding.*sharding.*devices=\[1,2,1\]",
                      1,
                  ),
                  (
                      r"f32\[2,7\].*custom_call_target.*Sharding.*sharding.*replicated",
                      2,
                  ),
                  (
                      r"f32\[8,5\].*custom_call_target.*Sharding.*sharding.*devices=\[2,1\]",
                      1,
                  ),
              ],
          )
  
imports: from collections.abc import Sequence
  import contextlib
  from functools import partial
  import logging
  import math
  import os
  import re
  from typing import Any
  import unittest
  from absl.testing import absltest
  import jax
  from jax._src import compiler
  from jax._src import config
  from jax._src import maps
  from jax._src import test_util
  from jax._src import xla_bridge
  from jax import lax
  from jax.experimental import jax2tf
  from jax.experimental import pjit
  from jax.experimental.maps import xmap
  from jax.experimental.shard_map import shard_map
  from jax.sharding import NamedSharding
  from jax.sharding import Mesh
  from jax.sharding import PartitionSpec
  import jax.numpy
  import numpy
  import tensorflow
  from jax.experimental.jax2tf.tests import tf_test_util
label_desc: Use image node_id 15 for calling the ShardingTest obj's underlying member method code with example usage: obj.test_xmap_collective_reduce() and returns: res,  res, f_converted