code: def test_vmap_after(self):
      batch = 4
      qy_size = 128
      db_size = 1024
      feature_dim = 32
      k = 10
      rng = jtu.rand_default(self.rng())
      qy = rng([qy_size, feature_dim, batch], np.float32)
      db = rng([db_size, feature_dim, batch], np.float32)
      recall = 0.95
  
      # Create ground truth
      gt_scores = lax.dot_general(qy, db, (([1], [1]), ([2], [2])))
      _, gt_args = lax.top_k(gt_scores, k)
      gt_args = lax.transpose(gt_args, [2, 0, 1])
      gt_args = lax.reshape(gt_args, [qy_size * batch, k])
  
      # test target
      def approx_max_k(qy, db):
          scores = qy @ db.transpose()
          return lax.approx_max_k(scores, k)
  
      _, ann_args = jax.vmap(approx_max_k, (2, 2))(qy, db)
      ann_args = lax.transpose(ann_args, [2, 0, 1])
      ann_args = lax.reshape(ann_args, [qy_size * batch, k])
      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 6 for calling the AnnTest obj's underlying member method code with example usage: obj.test_vmap_after() and returns: lax