code: def test_pmap(self, qy_shape, db_shape, dtype, k, recall):
      num_devices = jax.device_count()
      rng = jtu.rand_default(self.rng())
      qy = rng(qy_shape, dtype)
      db = rng(db_shape, dtype)
      db_size = db.shape[0]
      gt_scores = lax.dot_general(qy, db, (([1], [1]), ([], [])))
      _, gt_args = lax.top_k(
          -gt_scores, k
      )  # negate the score to get min-k
      db_per_device = db_size // num_devices
      sharded_db = db.reshape(num_devices, db_per_device, 128)
      db_offsets = (
          np.arange(num_devices, dtype=np.int32) * db_per_device
      )
  
      def parallel_topk(qy, db, db_offset):
          scores = lax.dot_general(qy, db, (([1], [1]), ([], [])))
          ann_vals, ann_args = lax.approx_min_k(
              scores,
              k=k,
              reduction_dimension=1,
              recall_target=recall,
              reduction_input_size_override=db_size,
              aggregate_to_topk=False,
          )
          return (ann_vals, ann_args + db_offset)
  
      # shape = qy_size, num_devices, approx_dp
      ann_vals, ann_args = jax.pmap(
          parallel_topk, in_axes=(None, 0, 0), out_axes=(1, 1)
      )(qy, sharded_db, db_offsets)
      # collapse num_devices and approx_dp
      ann_vals = lax.collapse(ann_vals, 1, 3)
      ann_args = lax.collapse(ann_args, 1, 3)
      ann_vals, ann_args = lax.sort_key_val(
          ann_vals, ann_args, dimension=1
      )
      ann_args = lax.slice_in_dim(
          ann_args, start_index=0, limit_index=k, axis=1
      )
      ann_recall = compute_recall(
          np.asarray(ann_args), np.asarray(gt_args)
      )
      self.assertGreater(ann_recall, recall)
  
file_path: jax/tests/ann_test.py
imports: from functools import partial
  import math
  from absl.testing import absltest
  import numpy
  import jax
  from jax import lax
  from jax._src import test_util
  from jax import config
label_desc: Use image node_id 4 for calling the AnnTest obj's underlying member method code with example usage: obj.test_pmap(qy_shape, db_shape, dtype, k, recall) and returns: ann_vals,  unknown