Ensure calculated MI values are non-negative.

PiperOrigin-RevId: 292439679
peytondmurray · Jan 30, 2020 · 70a08b4 · 70a08b4
1 parent 75e9524
commit 70a08b4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/RELEASE.md b/RELEASE.md
@@ -37,6 +37,7 @@
 *   Set num_missing and num_examples correctly in the presence of sparse
     features.
 *   Compute weighted feature stats for all weighted features declared in schema.
+*   Enforce that mutual information is non-negative.
 *   Depends on `tensorflow-metadata>=0.21.0,<0.22`.
 *   Depends on `pyarrow>=0.15` (removed the upper bound as it is determined by
     `tfx-bsl`).

diff --git a/tensorflow_data_validation/statistics/generators/sklearn_mutual_information.py b/tensorflow_data_validation/statistics/generators/sklearn_mutual_information.py
@@ -258,7 +258,7 @@ def _calculate_mi(self, df: pd.DataFrame, labels: np.ndarray,
     for i, (mi, shuffled_mi) in enumerate(
         zip(mi_per_feature, shuffled_mi_per_feature)):
       result[df.columns[i]] = {
-          MUTUAL_INFORMATION_KEY: mi,
+          MUTUAL_INFORMATION_KEY: mi.clip(min=0),
           ADJUSTED_MUTUAL_INFORMATION_KEY: mi - shuffled_mi
       }
     return result