code: def test_autodiff(self, shape, dtype, k, is_max_k):
      vals = np.arange(math.prod(shape), dtype=dtype)
      vals = self.rng().permutation(vals).reshape(shape)
      if is_max_k:
          fn = lambda vs: lax.approx_max_k(vs, k=k)[0]
      else:
          fn = lambda vs: lax.approx_min_k(vs, k=k)[0]
      jtu.check_grads(fn, (vals,), 2, ["fwd", "rev"], eps=1e-2)
  
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 3 for calling the AnnTest obj's underlying member method code with example usage: obj.test_autodiff(shape, dtype, k, is_max_k) without return types