code: def test_non_fused_min_k(
      self, k, row_size, num_rows, aggregate_to_topk
  ):
      # Use the new rng api
      row = np.arange(row_size, dtype=np.float32)
      db = np.stack(
          list(self._rng.permutation(row) for _ in range(num_rows))
      )
  
      @function(jit_compile=True)
      def ann(db, k=10):
          return nn_ops.approx_min_k(
              db, k, aggregate_to_topk=aggregate_to_topk
          )
  
      with ops.device("/device:TPU:0"):
          db_op = variables.Variable(db)
          result = ann(db_op, k)[1]
  
      gt = np.argsort(db)[:, :k]
      ann_recall = self.compute_recall(result.numpy(), gt)
      self.assertGreaterEqual(ann_recall, 0.95)
  
imports: import itertools
  from absl.testing import parameterized
  import numpy
  from tensorflow.python.eager import backprop
  from tensorflow.python.eager import test
  from tensorflow.python.eager.def_function import function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.ops import variables
label_desc: Use image node_id 4 for calling the ApproxTopkTest obj's underlying member method code with example usage: obj.test_non_fused_min_k(k, row_size, num_rows, aggregate_to_topk) and returns: nn_ops