Implement ProductBaselines (#1212)

Summary:
Pull Request resolved: #1212

Implement a Callable Baselines class that returns a sample from the Cartesian product of the inputs' available baselines.

Reviewed By: vivekmig

Differential Revision: D51582979

fbshipit-source-id: ec7c833a5572b6a15c5cc8acb3fb9b1bcf439065

Author: aobo-y

captum/attr/_utils/baselines.py

@@ -0,0 +1,62 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+import random
+from typing import Any, Dict, List, Tuple, Union
+
+
+class ProductBaselines:
+    """
+    A Callable Baselines class that returns a sample from the Cartesian product of
+    the inputs' available baselines.
+
+    Args:
+        baseline_values (List or Dict): A list or dict of lists containing
+            the possible values for each feature. If a dict is provided, the keys
+            can a string of the feature name and the values is a list of available
+            baselines. The keys can also be a tuple of strings to group
+            multiple features whose baselines are not independent to each other.
+            If the key is a tuple, the value must be a list of tuples of
+            the corresponding values.
+    """
+
+    def __init__(
+        self,
+        baseline_values: Union[
+            List[List[Any]],
+            Dict[Union[str, Tuple[str, ...]], List[Any]],
+        ],
+    ):
+        if isinstance(baseline_values, dict):
+            dict_keys = list(baseline_values.keys())
+            baseline_values = [baseline_values[k] for k in dict_keys]
+        else:
+            dict_keys = []
+
+        self.dict_keys = dict_keys
+        self.baseline_values = baseline_values
+
+    def sample(self) -> Union[List[Any], Dict[str, Any]]:
+        baselines = [
+            random.choice(baseline_list) for baseline_list in self.baseline_values
+        ]
+
+        if not self.dict_keys:
+            return baselines
+
+        dict_baselines = {}
+        for key, val in zip(self.dict_keys, baselines):
+            if not isinstance(key, tuple):
+                key, val = (key,), (val,)
+
+            for k, v in zip(key, val):
+                dict_baselines[k] = v
+
+        return dict_baselines
+
+    def __call__(self) -> Union[List[Any], Dict[str, Any]]:
+        """
+        Returns:
+
+            baselines (List or Dict): A sample from the Cartesian product of
+                the inputs' available baselines
+        """
+        return self.sample()

tests/attr/test_baselines.py

@@ -0,0 +1,63 @@
+from typing import cast, Dict, List, Tuple, Union
+
+from captum.attr._utils.baselines import ProductBaselines
+
+# from parameterized import parameterized
+from tests.helpers.basic import BaseTest
+
+
+class TestProductBaselines(BaseTest):
+    def test_list(self) -> None:
+        baseline_values = [
+            [1, 2, 3],
+            [4, 5, 6, 7],
+            [8, 9],
+        ]
+
+        baselines = ProductBaselines(baseline_values)
+
+        baseline_sample = baselines()
+
+        self.assertIsInstance(baseline_sample, list)
+        for sample_val, vals in zip(baseline_sample, baseline_values):
+            self.assertIn(sample_val, vals)
+
+    def test_dict(self) -> None:
+        baseline_values = {
+            "f1": [1, 2, 3],
+            "f2": [4, 5, 6, 7],
+            "f3": [8, 9],
+        }
+
+        baselines = ProductBaselines(
+            cast(Dict[Union[str, Tuple[str, ...]], List[int]], baseline_values)
+        )
+
+        baseline_sample = baselines()
+
+        self.assertIsInstance(baseline_sample, dict)
+        baseline_sample = cast(dict, baseline_sample)
+
+        for sample_key, sample_val in baseline_sample.items():
+            self.assertIn(sample_val, baseline_values[sample_key])
+
+    def test_dict_tuple_key(self) -> None:
+        baseline_values: Dict[Union[str, Tuple[str, ...]], List] = {
+            ("f1", "f2"): [(1, "1"), (2, "2"), (3, "3")],
+            "f3": [4, 5],
+        }
+
+        baselines = ProductBaselines(baseline_values)
+
+        baseline_sample = baselines()
+
+        self.assertIsInstance(baseline_sample, dict)
+        baseline_sample = cast(dict, baseline_sample)
+
+        self.assertEqual(len(baseline_sample), 3)
+
+        self.assertIn(
+            (baseline_sample["f1"], baseline_sample["f2"]),
+            baseline_values[("f1", "f2")],
+        )
+        self.assertIn(baseline_sample["f3"], baseline_values["f3"])