def lambert_w(x):
  """Computes the Lambert W function of a `Tensor`.

  Computes the principal branch of the Lambert W function, i.e. the value w such
  that w * exp(w) = x for a a given x. For the principal branch, x must be real
  x >= -1 / e, and w >= -1.

  Args:
    x: A `Tensor` containing the values for which the principal branch of
      the Lambert W function is computed.

  Returns:
    A `Tensor` with the same shape and dtype as x containing the value of the
    Lambert W function.
  """
  dtype = x.dtype
  e = tf.constant(np.exp(1.0), dtype)
  inv_e = tf.constant(np.exp(-1.0), dtype)
  s = (np.exp(1) - 1.0) / (np.exp(2) - 1.0)
  slope = tf.constant(s, dtype)
  c = tf.constant(1 / np.exp(1) * (1 - s), dtype)
  log_s = tf.math.log(x)
  w_init = tf.where(
      x < inv_e,
      x,
      tf.where(x < e,
               slope * x + c,
               (log_s + (1.0 / log_s - 1.0) * tf.math.log(log_s))))

  def newton_update(count, w):
    expw = tf.math.exp(w)
    wexpw = w * expw
    return count + 1, w - (wexpw - x) / (expw + wexpw)

  count = tf.constant(0, tf.int32)
  num_iter = tf.constant(8)
  (unused_final_count, w) = tf.while_loop(
      lambda count, w: tf.less(count, num_iter),
      newton_update,
      [count, w_init])
  return w
Computes the Lambert W function of a `Tensor`.

  Computes the principal branch of the Lambert W function, i.e. the value w such
  that w * exp(w) = x for a a given x. For the principal branch, x must be real
  x >= -1 / e, and w >= -1.

  Args:
    x: A `Tensor` containing the values for which the principal branch of
      the Lambert W function is computed.

  Returns:
    A `Tensor` with the same shape and dtype as x containing the value of the
    Lambert W function.
  
tensorflow_transform/gaussianization.py