Typing fix

Differential Revision: D58703398

Author: [email protected]

captum/attr/_core/noise_tunnel.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 from enum import Enum
-from typing import Any, cast, List, Tuple, Union
+from typing import Any, cast, List, Optional, Tuple, Union
 
 import torch
 from captum._utils.common import (
@@ -80,7 +80,7 @@ def attribute(
         inputs: Union[Tensor, Tuple[Tensor, ...]],
         nt_type: str = "smoothgrad",
         nt_samples: int = 5,
-        nt_samples_batch_size: int = None,
+        nt_samples_batch_size: Optional[int] = None,
         stdevs: Union[float, Tuple[float, ...]] = 1.0,
         draw_baseline_from_distrib: bool = False,
         **kwargs: Any,

captum/attr/_utils/interpretable_input.py

@@ -1,3 +1,4 @@
+from abc import ABC, abstractmethod
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
@@ -45,7 +46,7 @@ def _scatter_itp_attr_by_mask(
     return attr
 
 
-class InterpretableInput:
+class InterpretableInput(ABC):
     """
     InterpretableInput is an adapter for different kinds of model inputs to
     work in Captum's attribution methods. Generally, attribution methods of Captum
@@ -94,6 +95,7 @@ class to create other types of customized input.
     is only allowed in certain attribution classes like LLMAttribution for now.)
     """
 
+    @abstractmethod
     def to_tensor(self) -> Tensor:
         """
         Return the interpretable representation of this input as a tensor
@@ -104,6 +106,7 @@ def to_tensor(self) -> Tensor:
         """
         pass
 
+    @abstractmethod
     def to_model_input(self, itp_tensor: Optional[Tensor] = None) -> Any:
         """
         Get the (perturbed) input in the format required by the model

captum/attr/_utils/visualization.py

@@ -3,12 +3,14 @@
 from enum import Enum
 from typing import Any, Iterable, List, Optional, Tuple, Union
 
+import matplotlib
+
 import numpy as np
 from matplotlib import cm, colors, pyplot as plt
+from matplotlib.axes import Axes
 from matplotlib.collections import LineCollection
-from matplotlib.colors import LinearSegmentedColormap
+from matplotlib.colors import Colormap, LinearSegmentedColormap
 from matplotlib.figure import Figure
-from matplotlib.pyplot import axis, figure
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 from numpy import ndarray
 
@@ -51,7 +53,8 @@ def _normalize_scale(attr: ndarray, scale_factor: float):
         warnings.warn(
             "Attempting to normalize by value approximately 0, visualized results"
             "may be misleading. This likely means that attribution values are all"
-            "close to 0."
+            "close to 0.",
+            stacklevel=2,
         )
     attr_norm = attr / scale_factor
     return np.clip(attr_norm, -1, 1)
@@ -80,37 +83,39 @@ def _normalize_attr(
 
     # Choose appropriate signed values and rescale, removing given outlier percentage.
     if VisualizeSign[sign] == VisualizeSign.all:
-        threshold = _cumulative_sum_threshold(np.abs(attr_combined), 100 - outlier_perc)
+        threshold = _cumulative_sum_threshold(
+            np.abs(attr_combined), 100.0 - outlier_perc
+        )
     elif VisualizeSign[sign] == VisualizeSign.positive:
         attr_combined = (attr_combined > 0) * attr_combined
-        threshold = _cumulative_sum_threshold(attr_combined, 100 - outlier_perc)
+        threshold = _cumulative_sum_threshold(attr_combined, 100.0 - outlier_perc)
     elif VisualizeSign[sign] == VisualizeSign.negative:
         attr_combined = (attr_combined < 0) * attr_combined
         threshold = -1 * _cumulative_sum_threshold(
-            np.abs(attr_combined), 100 - outlier_perc
+            np.abs(attr_combined), 100.0 - outlier_perc
         )
     elif VisualizeSign[sign] == VisualizeSign.absolute_value:
         attr_combined = np.abs(attr_combined)
-        threshold = _cumulative_sum_threshold(attr_combined, 100 - outlier_perc)
+        threshold = _cumulative_sum_threshold(attr_combined, 100.0 - outlier_perc)
     else:
         raise AssertionError("Visualize Sign type is not valid.")
     return _normalize_scale(attr_combined, threshold)
 
 
 def visualize_image_attr(
     attr: ndarray,
-    original_image: Union[None, ndarray] = None,
+    original_image: Optional[ndarray] = None,
     method: str = "heat_map",
     sign: str = "absolute_value",
-    plt_fig_axis: Union[None, Tuple[figure, axis]] = None,
+    plt_fig_axis: Optional[Tuple[Figure, Axes]] = None,
     outlier_perc: Union[int, float] = 2,
-    cmap: Union[None, str] = None,
+    cmap: Optional[Union[str, Colormap]] = None,
     alpha_overlay: float = 0.5,
     show_colorbar: bool = False,
-    title: Union[None, str] = None,
+    title: Optional[str] = None,
     fig_size: Tuple[int, int] = (6, 6),
     use_pyplot: bool = True,
-):
+) -> Tuple[Figure, Axes]:
     r"""
     Visualizes attribution for a given image by normalizing attribution values
     of the desired sign (positive, negative, absolute value, or all) and displaying
@@ -231,7 +236,8 @@ def visualize_image_attr(
             plt_fig, plt_axis = plt.subplots(figsize=fig_size)
         else:
             plt_fig = Figure(figsize=fig_size)
-            plt_axis = plt_fig.subplots()
+            plt_axis = plt_fig.subplots()  # type: ignore
+            # Figure.subplots returns Axes or array of Axes
 
     if original_image is not None:
         if np.max(original_image) <= 1.0:
@@ -264,8 +270,8 @@ def visualize_image_attr(
 
         # Set default colormap and bounds based on sign.
         if VisualizeSign[sign] == VisualizeSign.all:
-            default_cmap = LinearSegmentedColormap.from_list(
-                "RdWhGn", ["red", "white", "green"]
+            default_cmap: Union[str, LinearSegmentedColormap] = (
+                LinearSegmentedColormap.from_list("RdWhGn", ["red", "white", "green"])
             )
             vmin, vmax = -1, 1
         elif VisualizeSign[sign] == VisualizeSign.positive:
@@ -345,7 +351,7 @@ def visualize_image_attr_multiple(
     original_image: Union[None, ndarray],
     methods: List[str],
     signs: List[str],
-    titles: Union[None, List[str]] = None,
+    titles: Optional[List[str]] = None,
     fig_size: Tuple[int, int] = (8, 6),
     use_pyplot: bool = True,
     **kwargs: Any,
@@ -423,17 +429,22 @@ def visualize_image_attr_multiple(
         plt_fig = Figure(figsize=fig_size)
     plt_axis = plt_fig.subplots(1, len(methods))
 
+    plt_axis_list: List[Axes] = []
     # When visualizing one
     if len(methods) == 1:
-        plt_axis = [plt_axis]
+        plt_axis_list = [plt_axis]  # type: ignore
+        # Figure.subplots returns Axes or array of Axes
+    else:
+        plt_axis_list = plt_axis  # type: ignore
+        # Figure.subplots returns Axes or array of Axes
 
     for i in range(len(methods)):
         visualize_image_attr(
             attr,
             original_image=original_image,
             method=methods[i],
             sign=signs[i],
-            plt_fig_axis=(plt_fig, plt_axis[i]),
+            plt_fig_axis=(plt_fig, plt_axis_list[i]),
             use_pyplot=False,
             title=titles[i] if titles else None,
             **kwargs,
@@ -452,12 +463,12 @@ def visualize_timeseries_attr(
     sign: str = "absolute_value",
     channel_labels: Optional[List[str]] = None,
     channels_last: bool = True,
-    plt_fig_axis: Union[None, Tuple[figure, axis]] = None,
+    plt_fig_axis: Optional[Tuple[Figure, Union[Axes, List[Axes]]]] = None,
     outlier_perc: Union[int, float] = 2,
-    cmap: Union[None, str] = None,
+    cmap: Optional[Union[str, Colormap]] = None,
     alpha_overlay: float = 0.7,
     show_colorbar: bool = False,
-    title: Union[None, str] = None,
+    title: Optional[str] = None,
     fig_size: Tuple[int, int] = (6, 6),
     use_pyplot: bool = True,
     **pyplot_kwargs,
@@ -596,7 +607,8 @@ def visualize_timeseries_attr(
     if num_channels > timeseries_length:
         warnings.warn(
             "Number of channels ({}) greater than time series length ({}), "
-            "please verify input format".format(num_channels, timeseries_length)
+            "please verify input format".format(num_channels, timeseries_length),
+            stacklevel=2,
         )
 
     num_subplots = num_channels
@@ -624,17 +636,20 @@ def visualize_timeseries_attr(
             )
         else:
             plt_fig = Figure(figsize=fig_size)
-            plt_axis = plt_fig.subplots(nrows=num_subplots, sharex=True)
+            plt_axis = plt_fig.subplots(nrows=num_subplots, sharex=True)  # type: ignore
+            # Figure.subplots returns Axes or array of Axes
 
     if not isinstance(plt_axis, ndarray):
-        plt_axis = np.array([plt_axis])
+        plt_axis_list = np.array([plt_axis])
+    else:
+        plt_axis_list = plt_axis
 
     norm_attr = _normalize_attr(attr, sign, outlier_perc, reduction_axis=None)
 
     # Set default colormap and bounds based on sign.
     if VisualizeSign[sign] == VisualizeSign.all:
-        default_cmap = LinearSegmentedColormap.from_list(
-            "RdWhGn", ["red", "white", "green"]
+        default_cmap: Union[str, LinearSegmentedColormap] = (
+            LinearSegmentedColormap.from_list("RdWhGn", ["red", "white", "green"])
         )
         vmin, vmax = -1, 1
     elif VisualizeSign[sign] == VisualizeSign.positive:
@@ -653,7 +668,6 @@ def visualize_timeseries_attr(
     cm_norm = colors.Normalize(vmin, vmax)
 
     def _plot_attrs_as_axvspan(attr_vals, x_vals, ax):
-
         half_col_width = (x_values[1] - x_values[0]) / 2.0
         for icol, col_center in enumerate(x_vals):
             left = col_center - half_col_width
@@ -670,52 +684,47 @@ def _plot_attrs_as_axvspan(attr_vals, x_vals, ax):
         TimeseriesVisualizationMethod[method]
         == TimeseriesVisualizationMethod.overlay_individual
     ):
-
         for chan in range(num_channels):
-
-            plt_axis[chan].plot(x_values, data[chan, :], **pyplot_kwargs)
+            plt_axis_list[chan].plot(x_values, data[chan, :], **pyplot_kwargs)
             if channel_labels is not None:
-                plt_axis[chan].set_ylabel(channel_labels[chan])
+                plt_axis_list[chan].set_ylabel(channel_labels[chan])
 
-            _plot_attrs_as_axvspan(norm_attr[chan], x_values, plt_axis[chan])
+            _plot_attrs_as_axvspan(norm_attr[chan], x_values, plt_axis_list[chan])
 
         plt.subplots_adjust(hspace=0)
 
     elif (
         TimeseriesVisualizationMethod[method]
         == TimeseriesVisualizationMethod.overlay_combined
     ):
-
         # Dark colors are better in this case
-        cycler = plt.cycler("color", cm.Dark2.colors)
-        plt_axis[0].set_prop_cycle(cycler)
+        cycler = plt.cycler("color", matplotlib.colormaps["Dark2"])  # type: ignore
+        plt_axis_list[0].set_prop_cycle(cycler)
 
         for chan in range(num_channels):
             label = channel_labels[chan] if channel_labels else None
-            plt_axis[0].plot(x_values, data[chan, :], label=label, **pyplot_kwargs)
+            plt_axis_list[0].plot(x_values, data[chan, :], label=label, **pyplot_kwargs)
 
-        _plot_attrs_as_axvspan(norm_attr, x_values, plt_axis[0])
+        _plot_attrs_as_axvspan(norm_attr, x_values, plt_axis_list[0])
 
-        plt_axis[0].legend(loc="best")
+        plt_axis_list[0].legend(loc="best")
 
     elif (
         TimeseriesVisualizationMethod[method]
         == TimeseriesVisualizationMethod.colored_graph
     ):
-
         for chan in range(num_channels):
-
             points = np.array([x_values, data[chan, :]]).T.reshape(-1, 1, 2)
             segments = np.concatenate([points[:-1], points[1:]], axis=1)
 
             lc = LineCollection(segments, cmap=cmap, norm=cm_norm, **pyplot_kwargs)
             lc.set_array(norm_attr[chan, :])
-            plt_axis[chan].add_collection(lc)
-            plt_axis[chan].set_ylim(
+            plt_axis_list[chan].add_collection(lc)
+            plt_axis_list[chan].set_ylim(
                 1.2 * np.min(data[chan, :]), 1.2 * np.max(data[chan, :])
             )
             if channel_labels is not None:
-                plt_axis[chan].set_ylabel(channel_labels[chan])
+                plt_axis_list[chan].set_ylabel(channel_labels[chan])
 
         plt.subplots_adjust(hspace=0)
 
@@ -725,7 +734,7 @@ def _plot_attrs_as_axvspan(attr_vals, x_vals, ax):
     plt.xlim([x_values[0], x_values[-1]])
 
     if show_colorbar:
-        axis_separator = make_axes_locatable(plt_axis[-1])
+        axis_separator = make_axes_locatable(plt_axis_list[-1])
         colorbar_axis = axis_separator.append_axes("bottom", size="5%", pad=0.4)
         colorbar_alpha = alpha_overlay
         if (
@@ -740,7 +749,7 @@ def _plot_attrs_as_axvspan(attr_vals, x_vals, ax):
             alpha=colorbar_alpha,
         )
     if title:
-        plt_axis[0].set_title(title)
+        plt_axis_list[0].set_title(title)
 
     if use_pyplot:
         plt.show()

captum/insights/attr_vis/features.py

@@ -8,12 +8,14 @@
 from captum._utils.common import safe_div
 from captum.attr._utils import visualization as viz
 from captum.insights.attr_vis._utils.transforms import format_transforms
+from matplotlib.figure import Figure
 from torch import Tensor
 
+
 FeatureOutput = namedtuple("FeatureOutput", "name base modified type contribution")
 
 
-def _convert_figure_base64(fig) -> str:
+def _convert_figure_base64(fig: Figure) -> str:
     buff = BytesIO()
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")

tests/insights/test_contribution.py

@@ -1,14 +1,15 @@
 #!/usr/bin/env python3
 
 import unittest
-from typing import Callable, List, Union
+from typing import Any, Callable, Generator, List, Tuple, Union
 
 import torch
 import torch.nn as nn
 from captum.insights import AttributionVisualizer, Batch
 from captum.insights.attr_vis.app import FilterConfig
 from captum.insights.attr_vis.features import BaseFeature, FeatureOutput, ImageFeature
 from tests.helpers import BaseTest
+from torch import Tensor
 
 
 class RealFeature(BaseFeature):
@@ -26,7 +27,8 @@ def __init__(
             visualization_transform=None,
         )
 
-    def visualization_type(self) -> str:
+    @staticmethod
+    def visualization_type() -> str:
         return "real"
 
     def visualize(self, attribution, data, contribution_frac) -> FeatureOutput:
@@ -103,7 +105,7 @@ def _labelled_img_data(
     height: int = 8,
     depth: int = 3,
     num_labels: int = 10,
-):
+) -> Generator[Tuple[Tensor, Tensor], Any, Any]:
     for _ in range(num_samples):
         yield torch.empty(depth, height, width).uniform_(0, 1), torch.randint(
             num_labels, (1,)