code: def patch_manual_special_functions():
      """Patches in the manually implemented special functions.
  
      Normally you'd just use functions coming from the array library you're using,
      but on some platforms (like TPU) the default implementations are
      insufficiently precise for certain tasks when running under 32 bits (64 bit
      implementations are typically okay).
  
      This patches in manual implementations of those functions which are provide
      higher precision at the cost of speed. The list of affected functions is:
  
      - `exp`
      - `log`
      - `expm1`
      - `log1p`
      - `logsumexp` (aka `reduce_logsumexp`)
      - `softplus`
  
      Yields:
        Nothing.
      """
      if JAX_MODE:
          old_expm1 = jax.numpy.expm1
          old_log1p = jax.numpy.log1p
          old_exp = jax.numpy.exp
          old_log = jax.numpy.log
          old_logsumexp = jax.scipy.special.logsumexp
          old_softplus = jax.nn.softplus
      else:
          old_expm1 = tf.math.expm1
          old_log1p = tf.math.log1p
          old_exp = tf.math.exp
          old_log = tf.math.log
          old_logsumexp = tf.math.reduce_logsumexp
          old_softplus = tf.math.softplus
  
      try:
          if JAX_MODE:
              jax.numpy.expm1 = expm1_pade_4_4
              jax.numpy.log1p = log1p_pade_4_4
              jax.numpy.exp = exp_pade_4_4
              jax.numpy.log = log_pade_4_4
              jax.scipy.special.logsumexp = reduce_logsumexp
              jax.nn.softplus = softplus
          else:
              tf.math.expm1 = expm1_pade_4_4
              tf.math.log1p = log1p_pade_4_4
              tf.math.exp = exp_pade_4_4
              tf.exp = exp_pade_4_4
              tf.math.log = log_pade_4_4
              tf.math.reduce_logsumexp = reduce_logsumexp
              tf.reduce_logsumexp = reduce_logsumexp
              tf.math.softplus = softplus
              tf.nn.softplus = softplus
          yield
      finally:
          if JAX_MODE:
              jax.numpy.expm1 = old_expm1
              jax.numpy.log1p = old_log1p
              jax.numpy.exp = old_exp
              jax.numpy.log = old_log
              jax.scipy.special.logsumexp = old_logsumexp
              jax.nn.softplus = old_softplus
          else:
              tf.math.expm1 = old_expm1
              tf.math.log1p = old_log1p
              tf.math.exp = old_exp
              tf.exp = old_exp
              tf.math.log = old_log
              tf.math.reduce_logsumexp = old_logsumexp
              tf.reduce_logsumexp = old_logsumexp
              tf.math.softplus = old_softplus
              tf.nn.softplus = old_softplus
  
imports: import contextlib
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import dtype_util
label_desc: Use image node_id 24 for calling a global function with example usage: patch_manual_special_functions() without return types