def calculate_partial_mutual_information(n_ij, x_i, y_j, n):
  """Calculates Mutual Information for x=i, y=j from sample counts.

  The standard formulation of mutual information is:
  MI(X,Y) = Sum_i,j {p_ij * log2(p_ij / p_i * p_j)}
  We are operating over counts (p_ij = n_ij / n), so this is transformed into
  MI(X,Y) = Sum_i,j {n_ij * (log2(n_ij) + log2(n) - log2(x_i) - log2(y_j))} / n
  This function returns the argument to the summation, the mutual information
  for a particular pair of values x_i, y_j (the caller is expected to divide
  the summation by n to compute the final mutual information result).

  Args:
    n_ij: The co-occurrence of x=i and y=j
    x_i: The frequency of x=i.
    y_j: The frequency of y=j.
    n: The total # observations

  Returns:
    Mutual information for the cell x=i, y=j.
  """
  if n_ij == 0 or x_i == 0 or y_j == 0:
    return 0
  return n_ij * ((log2(n_ij) + log2(n)) -
                 (log2(x_i) + log2(y_j)))
Calculates Mutual Information for x=i, y=j from sample counts.

  The standard formulation of mutual information is:
  MI(X,Y) = Sum_i,j {p_ij * log2(p_ij / p_i * p_j)}
  We are operating over counts (p_ij = n_ij / n), so this is transformed into
  MI(X,Y) = Sum_i,j {n_ij * (log2(n_ij) + log2(n) - log2(x_i) - log2(y_j))} / n
  This function returns the argument to the summation, the mutual information
  for a particular pair of values x_i, y_j (the caller is expected to divide
  the summation by n to compute the final mutual information result).

  Args:
    n_ij: The co-occurrence of x=i and y=j
    x_i: The frequency of x=i.
    y_j: The frequency of y=j.
    n: The total # observations

  Returns:
    Mutual information for the cell x=i, y=j.
  
tensorflow_transform/info_theory.py