Add support for feature extraction on list outputs in addition to np.…

…ndarray and empty lists. PiperOrigin-RevId: 329956465
dsnoor · Sep 3, 2020 · 9c07f41 · 9c07f41
1 parent b1e47da
commit 9c07f41
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tensorflow_model_analysis/extractors/feature_extractor.py b/tensorflow_model_analysis/extractors/feature_extractor.py
@@ -87,8 +87,9 @@ def _AugmentExtracts(data: Dict[Text, Any], prefix: Text, excludes: List[bytes],
       extracts[col_name] = types.MaterializedColumn(
           name=col_name, value=val.values)
 
-    elif isinstance(val, np.ndarray):
-      val = val[0]  # only support first dim for now.
+    elif isinstance(val, np.ndarray) or isinstance(val, list):
+      # Only support first dim for now
+      val = val[0] if len(val) > 0 else []  # pylint: disable=g-explicit-length-test
       extracts[col_name] = types.MaterializedColumn(name=col_name, value=val)
 
     else: