code: def _log_pade_4_4_fwd(
      x,
  ):  # pylint: disable=missing-function-docstring
      x = tf.convert_to_tensor(x, dtype_hint=tf.float32)
      orig_x = x
  
      dtype = dtype_util.as_numpy_dtype(x.dtype)
  
      zero = np.zeros([], dtype)
      log2 = np.log(2).astype(dtype)
      one = np.ones([], dtype)
      two = np.array(2, dtype)
      inf = np.array(float("inf"), dtype)
      nan = np.array(float("nan"), dtype)
  
      # Scale x to [0.5, 1), extract exponent.
      if JAX_MODE:
          # Despite the bit arithmetic, benchmarks showed this to be slightly faster
          # than the fallback code below.
          x, e = jax.numpy.frexp(x)
      else:
          # TensorFlow does not expose frexp.
          e = tf.math.ceil(_real_log(x) / log2)
          x = x / two**e
  
      # We'll be using a Pade approximant for log1p(x), so move x closer to zero.
      x_is_small = x < log2
      es = e - one
      xs = two * x - one
      el = e
      xl = x - one
  
      e = tf.where(x_is_small, es, el)
      x = tf.where(x_is_small, xs, xl)
  
      coeffs_p = np.array([5 / 84, 13 / 21, 3 / 2, 1, 0], dtype)
      coeffs_q = np.array([1 / 70, 2 / 7, 9 / 7, 2, 1], dtype)
  
      res = _horner(x, coeffs_p) / _horner(x, coeffs_q)
      res = res + e * log2
  
      # Special points.
      res = tf.where(tf.equal(orig_x, one), tf.zeros_like(res), res)
      res = tf.where(tf.less(orig_x, zero), nan, res)
      res = tf.where(tf.equal(orig_x, zero), -inf, res)
      res = tf.where(tf.equal(orig_x, inf), inf, res)
      return res, orig_x
  
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 9 for calling a global function with example usage: _log_pade_4_4_fwd(x) and returns: res,  orig_x