code: def testScipySpecialFun(
      self,
      op,
      rng_factory,
      shapes,
      dtypes,
      test_autodiff,
      nondiff_argnums,
  ):
      scipy_op = getattr(osp_special, op)
      lax_op = getattr(lsp_special, op)
      rng = rng_factory(self.rng())
      args_maker = self._GetArgsMaker(rng, shapes, dtypes)
      args = args_maker()
      self.assertAllClose(
          scipy_op(*args),
          lax_op(*args),
          atol=1e-3,
          rtol=1e-3,
          check_dtypes=False,
      )
      self._CompileAndCheck(lax_op, args_maker, rtol=1e-4)
  
      if test_autodiff:
  
          def partial_lax_op(*vals):
              list_args = list(vals)
              for i in nondiff_argnums:
                  list_args.insert(i, args[i])
              return lax_op(*list_args)
  
          assert list(nondiff_argnums) == sorted(set(nondiff_argnums))
          diff_args = [
              x for i, x in enumerate(args) if i not in nondiff_argnums
          ]
          jtu.check_grads(
              partial_lax_op,
              diff_args,
              order=1,
              atol=0.1 if jtu.test_device_matches(["tpu"]) else 1e-3,
              rtol=0.1,
              eps=1e-3,
          )
  
file_path: jax/tests/lax_scipy_special_functions_test.py
imports: import collections
  import functools
  import itertools
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  import scipy.special
  import jax
  from jax._src import test_util
  from jax.scipy import special
  from jax import config
label_desc: Use image node_id 2 for calling the LaxScipySpcialFunctionsTest obj's underlying member method code with example usage: obj.testScipySpecialFun(op, rng_factory, shapes, dtypes, test_autodiff, nondiff_argnums) and returns: lax_op