code: def test_patching_jax(
      self, exp, log, expm1, log1p, logsumexp, softplus
  ):
      if not JAX_MODE or NUMPY_MODE:
          self.skipTest("This test is JAX-only")
  
      import jax  # pylint: disable=g-import-not-at-top
  
      exp_calls = 0
      expm1_calls = 0
      log_calls = 0
      log1p_calls = 0
      logsumexp_calls = 0
      softplus_calls = 0
      with manual_special_functions.patch_manual_special_functions():
          jax.numpy.exp(0.0)
          exp_calls += 1
          self.assertEqual(exp_calls, exp.call_count)
  
          jax.numpy.log(0.0)
          log_calls += 1
          self.assertEqual(log_calls, log.call_count)
  
          jax.numpy.expm1(0.0)
          expm1_calls += 1
          self.assertEqual(expm1_calls, expm1.call_count)
  
          jax.numpy.log1p(0.0)
          log1p_calls += 1
          self.assertEqual(log1p_calls, log1p.call_count)
  
          jax.scipy.special.logsumexp(0.0)
          logsumexp_calls += 1
          self.assertEqual(logsumexp_calls, logsumexp.call_count)
  
          jax.nn.softplus(0.0)
          softplus_calls += 1
          self.assertEqual(softplus_calls, softplus.call_count)
  
file_path: probability/tensorflow_probability/python/experimental/math/manual_special_functions_test.py
imports: import functools
  from absl import flags
  from absl.testing import parameterized
  import mock
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.math import manual_special_functions
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 6 for calling the ManualSpecialFunctionsTest obj's underlying member method code with example usage: obj.test_patching_jax(exp, log, expm1, log1p, logsumexp, softplus) without return types