diff --git a/captum/_utils/gradient.py b/captum/_utils/gradient.py index 69502b7443..1e2b827ab4 100644 --- a/captum/_utils/gradient.py +++ b/captum/_utils/gradient.py @@ -28,6 +28,7 @@ from captum._utils.sample_gradient import SampleGradientWrapper from captum._utils.typing import ( ModuleOrModuleList, + SliceIntType, TargetType, TensorOrTupleOfTensorsGeneric, ) @@ -775,8 +776,11 @@ def compute_layer_gradients_and_eval( def construct_neuron_grad_fn( layer: Module, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable], + neuron_selector: Union[ + int, + Tuple[Union[int, SliceIntType], ...], + Callable[[Union[Tensor, Tuple[Tensor, ...]]], Tensor], + ], device_ids: Union[None, List[int]] = None, attribute_to_neuron_input: bool = False, # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. diff --git a/captum/_utils/typing.py b/captum/_utils/typing.py index 10a2385611..80cf22d451 100644 --- a/captum/_utils/typing.py +++ b/captum/_utils/typing.py @@ -41,6 +41,13 @@ TensorLikeList5D, ] +try: + # Subscripted slice syntax is not supported in previous Python versions, + # falling back to slice type. + SliceIntType = slice[int, int, int] +except TypeError: + # pyre-fixme[24]: Generic type `slice` expects 3 type parameters. + SliceIntType = slice # type: ignore # Necessary for Python >=3.7 and <3.9! if TYPE_CHECKING: diff --git a/captum/attr/_core/neuron/neuron_conductance.py b/captum/attr/_core/neuron/neuron_conductance.py index 359736d0c9..6c8020f93e 100644 --- a/captum/attr/_core/neuron/neuron_conductance.py +++ b/captum/attr/_core/neuron/neuron_conductance.py @@ -14,7 +14,12 @@ _verify_select_neuron, ) from captum._utils.gradient import compute_layer_gradients_and_eval -from captum._utils.typing import BaselineType, TargetType, TensorOrTupleOfTensorsGeneric +from captum._utils.typing import ( + BaselineType, + SliceIntType, + TargetType, + TensorOrTupleOfTensorsGeneric, +) from captum.attr._utils.approximation_methods import approximation_parameters from captum.attr._utils.attribution import GradientAttribution, NeuronAttribution from captum.attr._utils.batching import _batch_attribution @@ -39,8 +44,7 @@ class NeuronConductance(NeuronAttribution, GradientAttribution): def __init__( self, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - forward_func: Callable, + forward_func: Callable[..., Tensor], layer: Module, device_ids: Union[None, List[int]] = None, multiply_by_inputs: bool = True, @@ -94,8 +98,11 @@ def __init__( def attribute( self, inputs: TensorOrTupleOfTensorsGeneric, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - neuron_selector: Union[int, Tuple[int, ...], Callable], + neuron_selector: Union[ + int, + Tuple[Union[int, SliceIntType], ...], + Callable[[Union[Tensor, Tuple[Tensor, ...]]], Tensor], + ], baselines: BaselineType = None, target: TargetType = None, additional_forward_args: Optional[object] = None, @@ -285,28 +292,24 @@ def attribute( " results.", stacklevel=1, ) - # pyre-fixme[6]: For 1st argument expected `Tensor` but got - # `TensorOrTupleOfTensorsGeneric`. is_inputs_tuple = _is_tuple(inputs) - # pyre-fixme[9]: inputs has type `TensorOrTupleOfTensorsGeneric`; used as - # `Tuple[Tensor, ...]`. - inputs, baselines = _format_input_baseline(inputs, baselines) - # pyre-fixme[6]: For 1st argument expected `Tuple[Tensor, ...]` but got - # `TensorOrTupleOfTensorsGeneric`. - _validate_input(inputs, baselines, n_steps, method) + formatted_inputs, formatted_baselines = _format_input_baseline( + inputs, baselines + ) + _validate_input(formatted_inputs, formatted_baselines, n_steps, method) - num_examples = inputs[0].shape[0] + num_examples = formatted_inputs[0].shape[0] if internal_batch_size is not None: - num_examples = inputs[0].shape[0] + num_examples = formatted_inputs[0].shape[0] attrs = _batch_attribution( self, num_examples, internal_batch_size, n_steps, - inputs=inputs, - baselines=baselines, + inputs=formatted_inputs, + baselines=formatted_baselines, neuron_selector=neuron_selector, target=target, additional_forward_args=additional_forward_args, @@ -315,11 +318,9 @@ def attribute( ) else: attrs = self._attribute( - # pyre-fixme[6]: For 1st argument expected `Tuple[Tensor, ...]` but - # got `TensorOrTupleOfTensorsGeneric`. - inputs=inputs, + inputs=formatted_inputs, neuron_selector=neuron_selector, - baselines=baselines, + baselines=formatted_baselines, target=target, additional_forward_args=additional_forward_args, n_steps=n_steps, @@ -334,8 +335,11 @@ def attribute( def _attribute( self, inputs: Tuple[Tensor, ...], - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - neuron_selector: Union[int, Tuple[int, ...], Callable], + neuron_selector: Union[ + int, + Tuple[Union[int, SliceIntType], ...], + Callable[[Union[Tensor, Tuple[Tensor, ...]]], Tensor], + ], baselines: Tuple[Union[Tensor, int, float], ...], target: TargetType = None, additional_forward_args: Optional[object] = None, @@ -409,8 +413,9 @@ def _attribute( # Aggregates across all steps for each tensor in the input tuple total_grads = tuple( - # pyre-fixme[6]: For 4th argument expected `Tuple[int, ...]` but got `Size`. - _reshape_and_sum(scaled_grad, n_steps, num_examples, input_grad.shape[1:]) + _reshape_and_sum( + scaled_grad, n_steps, num_examples, tuple(input_grad.shape[1:]) + ) for (scaled_grad, input_grad) in zip(scaled_grads, input_grads) ) diff --git a/captum/attr/_core/neuron/neuron_deep_lift.py b/captum/attr/_core/neuron/neuron_deep_lift.py index f4648b43b5..e7e3f2a77e 100644 --- a/captum/attr/_core/neuron/neuron_deep_lift.py +++ b/captum/attr/_core/neuron/neuron_deep_lift.py @@ -4,7 +4,11 @@ from typing import Callable, cast, Optional, Tuple, Union from captum._utils.gradient import construct_neuron_grad_fn -from captum._utils.typing import BaselineType, TensorOrTupleOfTensorsGeneric +from captum._utils.typing import ( + BaselineType, + SliceIntType, + TensorOrTupleOfTensorsGeneric, +) from captum.attr._core.deep_lift import DeepLift, DeepLiftShap from captum.attr._utils.attribution import GradientAttribution, NeuronAttribution from captum.log import log_usage @@ -79,8 +83,11 @@ def __init__( def attribute( self, inputs: TensorOrTupleOfTensorsGeneric, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable], + neuron_selector: Union[ + int, + Tuple[Union[int, SliceIntType], ...], + Callable[[Union[Tensor, Tuple[Tensor, ...]]], Tensor], + ], baselines: BaselineType = None, additional_forward_args: Optional[object] = None, attribute_to_neuron_input: bool = False, @@ -309,8 +316,11 @@ def __init__( def attribute( self, inputs: TensorOrTupleOfTensorsGeneric, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable], + neuron_selector: Union[ + int, + Tuple[Union[int, SliceIntType], ...], + Callable[[Union[Tensor, Tuple[Tensor, ...]]], Tensor], + ], baselines: Union[ TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric] ], diff --git a/captum/attr/_core/neuron/neuron_feature_ablation.py b/captum/attr/_core/neuron/neuron_feature_ablation.py index c72cf806a1..d391481ed4 100644 --- a/captum/attr/_core/neuron/neuron_feature_ablation.py +++ b/captum/attr/_core/neuron/neuron_feature_ablation.py @@ -6,7 +6,11 @@ import torch from captum._utils.common import _verify_select_neuron from captum._utils.gradient import _forward_layer_eval -from captum._utils.typing import BaselineType, TensorOrTupleOfTensorsGeneric +from captum._utils.typing import ( + BaselineType, + SliceIntType, + TensorOrTupleOfTensorsGeneric, +) from captum.attr._core.feature_ablation import FeatureAblation from captum.attr._utils.attribution import NeuronAttribution, PerturbationAttribution from captum.log import log_usage @@ -31,8 +35,7 @@ class NeuronFeatureAblation(NeuronAttribution, PerturbationAttribution): def __init__( self, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - forward_func: Callable, + forward_func: Callable[..., Union[int, float, Tensor]], layer: Module, device_ids: Union[None, List[int]] = None, ) -> None: @@ -61,8 +64,11 @@ def __init__( def attribute( self, inputs: TensorOrTupleOfTensorsGeneric, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable], + neuron_selector: Union[ + int, + Tuple[Union[int, SliceIntType], ...], + Callable[[Union[Tensor, Tuple[Tensor, ...]]], Tensor], + ], baselines: BaselineType = None, additional_forward_args: Optional[object] = None, feature_mask: Union[None, TensorOrTupleOfTensorsGeneric] = None, @@ -250,8 +256,7 @@ def attribute( >>> feature_mask=feature_mask) """ - # pyre-fixme[3]: Return type must be annotated. - def neuron_forward_func(*args: Any): + def neuron_forward_func(*args: Any) -> Tensor: with torch.no_grad(): layer_eval = _forward_layer_eval( self.forward_func, diff --git a/captum/attr/_core/neuron/neuron_gradient.py b/captum/attr/_core/neuron/neuron_gradient.py index 0e74382d33..b806c1f4c2 100644 --- a/captum/attr/_core/neuron/neuron_gradient.py +++ b/captum/attr/_core/neuron/neuron_gradient.py @@ -14,9 +14,10 @@ apply_gradient_requirements, undo_gradient_requirements, ) -from captum._utils.typing import TensorOrTupleOfTensorsGeneric +from captum._utils.typing import SliceIntType, TensorOrTupleOfTensorsGeneric from captum.attr._utils.attribution import GradientAttribution, NeuronAttribution from captum.log import log_usage +from torch import Tensor from torch.nn import Module @@ -28,8 +29,7 @@ class NeuronGradient(NeuronAttribution, GradientAttribution): def __init__( self, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - forward_func: Callable, + forward_func: Callable[..., Union[int, float, Tensor]], layer: Module, device_ids: Union[None, List[int]] = None, ) -> None: @@ -60,8 +60,11 @@ def __init__( def attribute( self, inputs: TensorOrTupleOfTensorsGeneric, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable], + neuron_selector: Union[ + int, + Tuple[Union[int, SliceIntType], ...], + Callable[[Union[Tensor, Tuple[Tensor, ...]]], Tensor], + ], additional_forward_args: Optional[object] = None, attribute_to_neuron_input: bool = False, ) -> TensorOrTupleOfTensorsGeneric: @@ -162,18 +165,12 @@ def attribute( >>> # index (4,1,2). >>> attribution = neuron_ig.attribute(input, (4,1,2)) """ - # pyre-fixme[6]: For 1st argument expected `Tensor` but got - # `TensorOrTupleOfTensorsGeneric`. is_inputs_tuple = _is_tuple(inputs) - # pyre-fixme[9]: inputs has type `TensorOrTupleOfTensorsGeneric`; used as - # `Tuple[Tensor, ...]`. - inputs = _format_tensor_into_tuples(inputs) + inputs_tuple = _format_tensor_into_tuples(inputs) additional_forward_args = _format_additional_forward_args( additional_forward_args ) - # pyre-fixme[6]: For 1st argument expected `Tuple[Tensor, ...]` but got - # `TensorOrTupleOfTensorsGeneric`. - gradient_mask = apply_gradient_requirements(inputs) + gradient_mask = apply_gradient_requirements(inputs_tuple) _, input_grads = _forward_layer_eval_with_neuron_grads( self.forward_func, @@ -185,9 +182,9 @@ def attribute( attribute_to_layer_input=attribute_to_neuron_input, ) - # pyre-fixme[6]: For 1st argument expected `Tuple[Tensor, ...]` but got - # `TensorOrTupleOfTensorsGeneric`. - undo_gradient_requirements(inputs, gradient_mask) - # pyre-fixme[7]: Expected `TensorOrTupleOfTensorsGeneric` but got - # `Tuple[Tensor, ...]`. + undo_gradient_requirements(inputs_tuple, gradient_mask) + + # pyre-fixme[7]: Expected `Variable[TensorOrTupleOfTensorsGeneric <: + # [Tensor, typing.Tuple[Tensor, ...]]]` but got `Union[Tensor, + # typing.Tuple[Tensor, ...]]`. return _format_output(is_inputs_tuple, input_grads) diff --git a/captum/attr/_core/neuron/neuron_gradient_shap.py b/captum/attr/_core/neuron/neuron_gradient_shap.py index 18b6507231..b0b82084f5 100644 --- a/captum/attr/_core/neuron/neuron_gradient_shap.py +++ b/captum/attr/_core/neuron/neuron_gradient_shap.py @@ -4,10 +4,11 @@ from typing import Callable, List, Optional, Tuple, Union from captum._utils.gradient import construct_neuron_grad_fn -from captum._utils.typing import TensorOrTupleOfTensorsGeneric +from captum._utils.typing import SliceIntType, TensorOrTupleOfTensorsGeneric from captum.attr._core.gradient_shap import GradientShap from captum.attr._utils.attribution import GradientAttribution, NeuronAttribution from captum.log import log_usage +from torch import Tensor from torch.nn import Module @@ -50,8 +51,7 @@ class NeuronGradientShap(NeuronAttribution, GradientAttribution): def __init__( self, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - forward_func: Callable, + forward_func: Callable[..., Union[int, float, Tensor]], layer: Module, device_ids: Union[None, List[int]] = None, multiply_by_inputs: bool = True, @@ -97,8 +97,11 @@ def __init__( def attribute( self, inputs: TensorOrTupleOfTensorsGeneric, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable], + neuron_selector: Union[ + int, + Tuple[Union[int, SliceIntType], ...], + Callable[[Union[Tensor, Tuple[Tensor, ...]]], Tensor], + ], baselines: Union[ TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric] ], diff --git a/captum/attr/_core/neuron/neuron_integrated_gradients.py b/captum/attr/_core/neuron/neuron_integrated_gradients.py index 8e56221d77..0e4504bee9 100644 --- a/captum/attr/_core/neuron/neuron_integrated_gradients.py +++ b/captum/attr/_core/neuron/neuron_integrated_gradients.py @@ -4,7 +4,7 @@ from typing import Callable, List, Optional, Tuple, Union from captum._utils.gradient import construct_neuron_grad_fn -from captum._utils.typing import TensorOrTupleOfTensorsGeneric +from captum._utils.typing import SliceIntType, TensorOrTupleOfTensorsGeneric from captum.attr._core.integrated_gradients import IntegratedGradients from captum.attr._utils.attribution import GradientAttribution, NeuronAttribution from captum.log import log_usage @@ -27,8 +27,7 @@ class NeuronIntegratedGradients(NeuronAttribution, GradientAttribution): def __init__( self, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - forward_func: Callable, + forward_func: Callable[..., Tensor], layer: Module, device_ids: Union[None, List[int]] = None, multiply_by_inputs: bool = True, @@ -76,8 +75,11 @@ def __init__( def attribute( self, inputs: TensorOrTupleOfTensorsGeneric, - # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters. - neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable], + neuron_selector: Union[ + int, + Tuple[Union[int, SliceIntType], ...], + Callable[[Union[Tensor, Tuple[Tensor, ...]]], Tensor], + ], baselines: Union[None, Tensor, Tuple[Tensor, ...]] = None, additional_forward_args: Optional[object] = None, n_steps: int = 50,