def _hypergeometric_pmf(n, x_i, y_j):
  """Probablity for expectation computation under hypergeometric distribution.

  Args:
    n: The sum of weights for all values.
    x_i: The sum of weights for the first variable taking on value i
    y_j: The sum of weights for the second variable taking on value j

  Yields:
    The probability p_j at point n_j in the hypergeometric distribution.
  """
  start = int(round(max(0, x_i + y_j - n)))
  end = int(round(min(x_i, y_j)))
  # Use log factorial to preserve calculation precision.
  # Note: because the factorials are expensive to compute, we compute the
  # denominator incrementally, at the cost of some readability.
  numerator = (
      _logfactorial(x_i) + _logfactorial(y_j) + _logfactorial(n - x_i) +
      _logfactorial(n - y_j))
  denominator = (
      _logfactorial(n) + _logfactorial(start) + _logfactorial(x_i - start) +
      _logfactorial(y_j - start) + _logfactorial(n - x_i - y_j + start))
  for n_j in range(start, end + 1):
    p_j = math.exp(numerator - denominator)
    if n_j != end:
      denominator += (
          math.log(n_j + 1) - math.log(x_i - n_j) - math.log(y_j - n_j) +
          math.log(n - x_i - y_j + n_j + 1))
    yield n_j, p_j
Probablity for expectation computation under hypergeometric distribution.

  Args:
    n: The sum of weights for all values.
    x_i: The sum of weights for the first variable taking on value i
    y_j: The sum of weights for the second variable taking on value j

  Yields:
    The probability p_j at point n_j in the hypergeometric distribution.
  
tensorflow_transform/info_theory.py