Skip to content

Commit

Permalink
Fix pyre errors in GradientShap (pytorch#1394)
Browse files Browse the repository at this point in the history 
Summary:

Initial work on fixing Pyre errors in GradientSHAP

Differential Revision: D64677343
  • Loading branch information
@facebook-github-bot
Vivek Miglani authored and facebook-github-bot committed Oct 22, 2024
1 parent eec1eb6 commit 1f2a97b
Showing 1 changed file with 25 additions and 57 deletions.
82 changes: 25 additions & 57 deletions captum/attr/_core/gradient_shap.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,13 @@

# pyre-strict
import typing
from typing import Any, Callable, Tuple, Union
from typing import Any, Callable, Literal, Tuple, Union

import numpy as np
import torch
from captum._utils.common import _is_tuple
from captum._utils.typing import (
BaselineType,
Literal,
TargetType,
Tensor,
TensorOrTupleOfTensorsGeneric,
Expand Down Expand Up @@ -57,8 +56,9 @@ class GradientShap(GradientAttribution):
samples and compute the expectation (smoothgrad).
"""

# pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
def __init__(self, forward_func: Callable, multiply_by_inputs: bool = True) -> None:
def __init__(
self, forward_func: Callable[..., Tensor], multiply_by_inputs: bool = True
) -> None:
r"""
Args:
Expand All @@ -82,8 +82,6 @@ def __init__(self, forward_func: Callable, multiply_by_inputs: bool = True) -> N
self._multiply_by_inputs = multiply_by_inputs

@typing.overload
# pyre-fixme[43]: The implementation of `attribute` does not accept all possible
# arguments of overload defined on line `84`.
def attribute(
self,
inputs: TensorOrTupleOfTensorsGeneric,
Expand All @@ -93,17 +91,12 @@ def attribute(
n_samples: int = 5,
stdevs: Union[float, Tuple[float, ...]] = 0.0,
target: TargetType = None,
# pyre-fixme[2]: Parameter annotation cannot be `Any`.
additional_forward_args: Any = None,
additional_forward_args: object = None,
*,
# pyre-fixme[31]: Expression `Literal[True]` is not a valid type.
# pyre-fixme[24]: Non-generic type `typing.Literal` cannot take parameters.
return_convergence_delta: Literal[True],
) -> Tuple[TensorOrTupleOfTensorsGeneric, Tensor]: ...

@typing.overload
# pyre-fixme[43]: The implementation of `attribute` does not accept all possible
# arguments of overload defined on line `99`.
def attribute(
self,
inputs: TensorOrTupleOfTensorsGeneric,
Expand All @@ -113,10 +106,7 @@ def attribute(
n_samples: int = 5,
stdevs: Union[float, Tuple[float, ...]] = 0.0,
target: TargetType = None,
additional_forward_args: Any = None,
# pyre-fixme[9]: return_convergence_delta has type `Literal[]`; used as `bool`.
# pyre-fixme[31]: Expression `Literal[False]` is not a valid type.
# pyre-fixme[24]: Non-generic type `typing.Literal` cannot take parameters.
additional_forward_args: object = None,
return_convergence_delta: Literal[False] = False,
) -> TensorOrTupleOfTensorsGeneric: ...

Expand All @@ -132,7 +122,7 @@ def attribute(
n_samples: int = 5,
stdevs: Union[float, Tuple[float, ...]] = 0.0,
target: TargetType = None,
additional_forward_args: Any = None,
additional_forward_args: object = None,
return_convergence_delta: bool = False,
) -> Union[
TensorOrTupleOfTensorsGeneric, Tuple[TensorOrTupleOfTensorsGeneric, Tensor]
Expand Down Expand Up @@ -265,20 +255,10 @@ def attribute(
"""
# since `baselines` is a distribution, we can generate it using a function
# rather than passing it as an input argument
# pyre-fixme[9]: baselines has type `Union[typing.Callable[...,
# Variable[TensorOrTupleOfTensorsGeneric <: [Tensor, typing.Tuple[Tensor,
# ...]]]], Variable[TensorOrTupleOfTensorsGeneric <: [Tensor,
# typing.Tuple[Tensor, ...]]]]`; used as `Tuple[Tensor, ...]`.
baselines = _format_callable_baseline(baselines, inputs)
# pyre-fixme[16]: Item `Callable` of `Union[(...) ->
# TensorOrTupleOfTensorsGeneric, TensorOrTupleOfTensorsGeneric]` has no
# attribute `__getitem__`.
assert isinstance(baselines[0], torch.Tensor), (
formatted_baselines = _format_callable_baseline(baselines, inputs)
assert isinstance(formatted_baselines[0], torch.Tensor), (
"Baselines distribution has to be provided in a form "
# pyre-fixme[16]: Item `Callable` of `Union[(...) ->
# TensorOrTupleOfTensorsGeneric, TensorOrTupleOfTensorsGeneric]` has no
# attribute `__getitem__`.
"of a torch.Tensor {}.".format(baselines[0])
"of a torch.Tensor {}.".format(formatted_baselines[0])
)

input_min_baseline_x_grad = InputBaselineXGradient(
Expand All @@ -296,7 +276,7 @@ def attribute(
nt_samples=n_samples,
stdevs=stdevs,
draw_baseline_from_distrib=True,
baselines=baselines,
baselines=formatted_baselines,
target=target,
additional_forward_args=additional_forward_args,
return_convergence_delta=return_convergence_delta,
Expand All @@ -322,8 +302,11 @@ def multiplies_by_inputs(self) -> bool:


class InputBaselineXGradient(GradientAttribution):
# pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
def __init__(self, forward_func: Callable, multiply_by_inputs: bool = True) -> None:
_multiply_by_inputs: bool

def __init__(
self, forward_func: Callable[..., Tensor], multiply_by_inputs: bool = True
) -> None:
r"""
Args:
Expand All @@ -345,37 +328,26 @@ def __init__(self, forward_func: Callable, multiply_by_inputs: bool = True) -> N
"""
GradientAttribution.__init__(self, forward_func)
# pyre-fixme[4]: Attribute must be annotated.
self._multiply_by_inputs = multiply_by_inputs

@typing.overload
# pyre-fixme[43]: The implementation of `attribute` does not accept all possible
# arguments of overload defined on line `318`.
def attribute(
self,
inputs: TensorOrTupleOfTensorsGeneric,
baselines: BaselineType = None,
target: TargetType = None,
# pyre-fixme[2]: Parameter annotation cannot be `Any`.
additional_forward_args: Any = None,
additional_forward_args: object = None,
*,
# pyre-fixme[31]: Expression `Literal[True]` is not a valid type.
# pyre-fixme[24]: Non-generic type `typing.Literal` cannot take parameters.
return_convergence_delta: Literal[True],
) -> Tuple[TensorOrTupleOfTensorsGeneric, Tensor]: ...

@typing.overload
# pyre-fixme[43]: The implementation of `attribute` does not accept all possible
# arguments of overload defined on line `329`.
def attribute(
self,
inputs: TensorOrTupleOfTensorsGeneric,
baselines: BaselineType = None,
target: TargetType = None,
additional_forward_args: Any = None,
# pyre-fixme[9]: return_convergence_delta has type `Literal[]`; used as `bool`.
# pyre-fixme[31]: Expression `Literal[False]` is not a valid type.
# pyre-fixme[24]: Non-generic type `typing.Literal` cannot take parameters.
additional_forward_args: object = None,
return_convergence_delta: Literal[False] = False,
) -> TensorOrTupleOfTensorsGeneric: ...

Expand All @@ -385,37 +357,33 @@ def attribute( # type: ignore
inputs: TensorOrTupleOfTensorsGeneric,
baselines: BaselineType = None,
target: TargetType = None,
additional_forward_args: Any = None,
additional_forward_args: object = None,
return_convergence_delta: bool = False,
) -> Union[
TensorOrTupleOfTensorsGeneric, Tuple[TensorOrTupleOfTensorsGeneric, Tensor]
]:
# Keeps track whether original input is a tuple or not before
# converting it into a tuple.
# 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)
inputs_tuple, baselines = _format_input_baseline(inputs, baselines)

rand_coefficient = torch.tensor(
np.random.uniform(0.0, 1.0, inputs[0].shape[0]),
device=inputs[0].device,
dtype=inputs[0].dtype,
np.random.uniform(0.0, 1.0, inputs_tuple[0].shape[0]),
device=inputs_tuple[0].device,
dtype=inputs_tuple[0].dtype,
)

input_baseline_scaled = tuple(
_scale_input(input, baseline, rand_coefficient)
for input, baseline in zip(inputs, baselines)
for input, baseline in zip(inputs_tuple, baselines)
)
grads = self.gradient_func(
self.forward_func, input_baseline_scaled, target, additional_forward_args
)

if self.multiplies_by_inputs:
input_baseline_diffs = tuple(
input - baseline for input, baseline in zip(inputs, baselines)
input - baseline for input, baseline in zip(inputs_tuple, baselines)
)
attributions = tuple(
input_baseline_diff * grad
Expand Down

0 comments on commit 1f2a97b

Please sign in to comment.