interpreto 0.4.9__py3-none-any.whl

interpreto/__init__.py

@@ -0,0 +1,80 @@
+# MIT License
+#
+# Copyright (c) 2025 IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL and FOR are research programs operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL and FOR are research programs operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+# =====================================================================================
+
+from .attributions import (
+    GradientShap,
+    IntegratedGradients,
+    KernelShap,
+    Lime,
+    Occlusion,
+    Saliency,
+    SmoothGrad,
+    Sobol,
+    SquareGrad,
+    VarGrad,
+)
+from .commons import (
+    Granularity,
+)
+from .model_wrapping import ModelWithSplitPoints
+from .visualizations.attributions import (
+    AttributionVisualization,
+)
+from .visualizations.concepts import (
+    ConceptHighlightVisualization,
+)
+
+
+def get_version() -> str:
+    """Returns the current version of the Interpreto library."""
+    import pkg_resources  # noqa: PLC0415  # ruff: disable=import-outside-toplevel
+
+    try:
+        return pkg_resources.get_distribution("interpreto").version
+    except pkg_resources.DistributionNotFound:
+        return "unknown"
+
+
+__all__ = [
+    "AttributionVisualization",
+    "ConceptHighlightVisualization",
+    "GradientShap",
+    "Granularity",
+    "IntegratedGradients",
+    "KernelShap",
+    "Lime",
+    "ModelWithSplitPoints",
+    "Occlusion",
+    "SquareGrad",
+    "Saliency",
+    "SmoothGrad",
+    "Sobol",
+    "VarGrad",
+    "get_version",
+]

interpreto/attributions/__init__.py

@@ -0,0 +1,50 @@
+# MIT License
+#
+# Copyright (c) 2025 IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL and FOR are research programs operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from .base import InferenceModes
+from .methods import (
+    GradientShap,
+    IntegratedGradients,
+    KernelShap,
+    Lime,
+    Occlusion,
+    Saliency,
+    SmoothGrad,
+    Sobol,
+    SquareGrad,
+    VarGrad,
+)
+
+__all__ = [
+    "GradientShap",
+    "IntegratedGradients",
+    "KernelShap",
+    "Lime",
+    "Occlusion",
+    "Saliency",
+    "SmoothGrad",
+    "Sobol",
+    "SquareGrad",
+    "VarGrad",
+]

interpreto/attributions/aggregations/__init__.py

@@ -0,0 +1,36 @@
+# MIT License
+#
+# Copyright (c) 2025 IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL and FOR are research programs operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from .base import (
+    Aggregator,
+    MaskwiseMeanAggregator,
+    MeanAggregator,
+    OcclusionAggregator,
+    SquaredMeanAggregator,
+    SumAggregator,
+    TrapezoidalMeanAggregator,
+    VarianceAggregator,
+)
+from .linear_regression_aggregation import LinearRegressionAggregator
+from .sobol_aggregation import SobolAggregator

interpreto/attributions/aggregations/base.py

@@ -0,0 +1,159 @@
+# MIT License
+#
+# Copyright (c) 2025 IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL and FOR are research programs operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""
+Aggregations used at the end of an attribution method
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+from beartype import beartype
+from jaxtyping import Float, jaxtyped
+from torch import Tensor
+
+
+class Aggregator:
+    """
+    Abstract class for aggregation made at the end of attribution methods
+    """
+
+    def aggregate(self, results: torch.Tensor, mask) -> torch.Tensor:
+        """
+        Get results from multiple "Inference wrappers", aggregate results and gives an explanation
+        """
+        return results.squeeze(0)
+
+    def __call__(self, results: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
+        return self.aggregate(results, mask)
+
+
+class TorchAggregator(Aggregator):
+    """
+    Basic aggregator using built-in torch methods to perform aggregation
+    """
+
+    _method: Callable[[torch.Tensor, int], torch.Tensor]
+
+    @jaxtyped(typechecker=beartype)
+    def aggregate(
+        self,
+        results: Float[Tensor, "p t l"],
+        mask: torch.Tensor | None = None,
+    ) -> Float[Tensor, "t l"]:
+        return self._method(results, dim=0)  # type: ignore
+
+
+class MeanAggregator(TorchAggregator):
+    """
+    Mean of attributions
+    """
+
+    _method = torch.mean
+
+
+class SquaredMeanAggregator(TorchAggregator):
+    """
+    Mean of squares of attributions
+    """
+
+    @staticmethod
+    def _method(x: torch.Tensor, dim: int = 0) -> torch.Tensor:  # type: ignore
+        return torch.mean(torch.square(x), dim=dim)
+
+
+class SumAggregator(TorchAggregator):
+    """
+    Sum of attributions
+    """
+
+    _method = torch.sum
+
+
+class VarianceAggregator(TorchAggregator):
+    """
+    Variance of attributions
+    """
+
+    _method = torch.var
+
+
+class MaskwiseMeanAggregator(Aggregator):
+    """Average scores for each target weighted by their perturbation mask.
+
+    This aggregator assumes that ``results`` contains one score per
+    perturbation and target and that ``mask`` indicates which input elements were
+    perturbed. The returned tensor has one importance score per target and input
+    element.
+    """
+
+    @jaxtyped(typechecker=beartype)
+    def aggregate(
+        self,
+        results: Float[Tensor, "p t"],
+        mask: Float[Tensor, "p l"],
+    ) -> Float[Tensor, "t l"]:
+        return torch.einsum("pt,pl->tl", results, mask) / mask.sum(dim=0)
+
+
+class OcclusionAggregator(Aggregator):
+    """Aggregate occlusion scores using a reference prediction.
+
+    The first row in ``results`` must correspond to the unperturbed prediction.
+    Subsequent rows represent the prediction for each perturbation. The returned
+    tensor contains the mean difference between the reference and each
+    perturbation weighted by the corresponding mask.
+    """
+
+    @jaxtyped(typechecker=beartype)
+    def aggregate(
+        self,
+        results: Float[Tensor, "p t"],
+        mask: Float[Tensor, "p l"],
+    ) -> Float[Tensor, "t l"]:
+        # first prediction is reference, unmodified input
+        scores = results[..., 0, :] - results[..., 1:, :]
+        mask = mask[..., 1:, :]
+        return torch.einsum("pt,pl->tl", scores, mask) / mask.sum(dim=0)
+
+
+class TrapezoidalMeanAggregator(Aggregator):
+    """
+    Weighted mean using the trapezoidal rule.
+
+    This aggregator performs a weighted average across the first axis of the input tensor,
+    assigning a weight of 0.5 to the first and last elements, and 1.0 to all intermediate ones.
+
+    This reflects trapezoidal numerical integration over uniformly spaced samples.
+    """
+
+    @jaxtyped(typechecker=beartype)
+    def aggregate(
+        self,
+        results: Float[Tensor, "p t l"],
+        mask: torch.Tensor | None = None,
+    ) -> Float[Tensor, "t l"]:
+        return torch.trapezoid(results, dim=0) / (results.shape[0] - 1)

interpreto/attributions/aggregations/linear_regression_aggregation.py

@@ -0,0 +1,194 @@
+# MIT License
+#
+# Copyright (c) 2025 IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL and FOR are research programs operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""
+Aggregator for LIME and KernelSHAP
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from enum import Enum
+
+import torch
+from beartype import beartype
+from jaxtyping import Float, jaxtyped
+from torch import Tensor
+from torch.nn import functional as F
+
+from interpreto.attributions.aggregations.base import Aggregator
+
+DistancesFromMaskProtocol = Callable[[Float[Tensor, "p l"]], Float[Tensor, "p"]]
+SimilarityKernelProtocol = Callable[[Float[Tensor, "p"], float], Float[Tensor, "p"]]
+
+
+def hamming_distance(mask: torch.Tensor) -> torch.Tensor:
+    """
+    Compute the Hamming distance between a mask and the original input
+    """
+    return torch.sum(mask, dim=1).float()
+
+
+def euclidean_distance(mask: torch.Tensor) -> torch.Tensor:
+    """
+    Compute the Euclidean distance between a mask and the original input
+    """
+    return torch.norm(mask, dim=1)
+
+
+def cosine_distance(mask: torch.Tensor) -> torch.Tensor:
+    """
+    Compute the cosine distance between a mask and the original input
+    """
+    return 1 - F.cosine_similarity(torch.ones_like(mask), mask, dim=1)
+
+
+class DistancesFromMask(Enum):
+    """
+    Enumeration of available distance functions.
+    """
+
+    HAMMING = staticmethod(hamming_distance)
+    EUCLIDEAN = staticmethod(euclidean_distance)
+    COSINE = staticmethod(cosine_distance)
+
+
+def exponential_kernel(distances: torch.Tensor, kernel_width: float) -> torch.Tensor:
+    """
+    Compute the exponential kernel for a given distance matrix
+    """
+    return torch.exp(-(distances**2) / (kernel_width**2))
+
+
+def ones_kernel(distances: torch.Tensor, kernel_width: float) -> torch.Tensor:
+    """
+    Compute the ones kernel for a given distance matrix
+    """
+    return torch.ones_like(distances)
+
+
+class Kernels(Enum):
+    """
+    Enumeration of available kernels.
+    """
+
+    EXPONENTIAL = staticmethod(exponential_kernel)
+    ONES = staticmethod(ones_kernel)
+
+
+def default_kernel_width_fn(mask: torch.Tensor) -> float:
+    """
+    Compute the default kernel width for a given mask
+    """
+    return (mask.shape[1] ** 0.5) * 0.75  # kernel parameter inspired by LIME
+
+
+class LinearRegressionAggregator(Aggregator):
+    """
+    Aggregator for masked perturbations using a linear model
+    """
+
+    def __init__(
+        self,
+        distance_function: DistancesFromMaskProtocol | None = None,  # noqa: UP036
+        similarity_kernel: SimilarityKernelProtocol = exponential_kernel,
+        kernel_width: float | Callable | None = None,
+    ):
+        """
+        Initialize the aggregator.
+
+        Args:
+            distance_function (DistancesFromMaskProtocol): distance function used to compute the similarity between perturbations and the original input from the mask.
+            similarity_kernel (Callable): similarity kernel used to compute the similarity between perturbations and the original input
+            kernel_width (float | Callable): kernel width used to compute the similarity between perturbations and the original input
+        """
+        self.distance_function = distance_function
+        self.similarity_kernel = similarity_kernel
+        self.kernel_width = kernel_width if kernel_width is not None else default_kernel_width_fn
+
+    @jaxtyped(typechecker=beartype)
+    def aggregate(
+        self,
+        results: Float[Tensor, "p t"],
+        mask: Float[Tensor, "p l"],
+    ) -> Float[Tensor, "t l"]:
+        """
+        Aggregate the results of the perturbations using a linear model.
+        We use the closed form solution for the weighted least squares problem.
+
+        Args:
+            results (torch.Tensor): The results of the perturbations. Shape: (p, t).
+                p the number of perturbations
+                t the number of target classes/tokens
+            mask (torch.Tensor): The mask used for the perturbations. Shape: (p, l).
+                l the length of the sequence
+        Returns:
+            torch.Tensor: The aggregated results. Shape: (t, l).
+        """
+        # Simplify typing
+        p, l = mask.shape
+        t = results.shape[1]
+
+        if results.shape[0] != p:
+            raise ValueError(
+                "The number of perturbations must match the number of results."
+                + f"Got {results.shape} perturbations and {results.shape} results."
+            )
+
+        # Define the weights for the linear model
+        if self.distance_function is not None:  # LIME
+            # Compute distance between perturbations and original input using the mask
+            distances: Float[Tensor, p] = self.distance_function(mask)
+
+            # Compute the similarities between perturbations and original input using the distances
+            if isinstance(self.kernel_width, Callable):
+                kernel_width: float = self.kernel_width(mask)
+            else:
+                assert isinstance(self.kernel_width, float) or isinstance(self.kernel_width, int)
+                kernel_width: float = self.kernel_width
+            weights: Float[Tensor, p] = self.similarity_kernel(distances, kernel_width)
+        else:  # Kernel SHAP
+            weights: Float[Tensor, p] = self.similarity_kernel(mask[:, 0], 1.0)
+
+        # We take the opposite of the mask
+        # Otherwise the linear treat the perturbation of a feature as passing from 0 to 1 for this feature
+        # We want exactly the opposite
+        mask = 1 - mask
+
+        # Compute closed form solution for the linear model
+        # Add a constant column for the bias term
+        X: Float[Tensor, p, l + 1] = torch.cat([torch.ones(results.shape[0], 1, device=mask.device), mask], dim=1)
+
+        # formula from wikipedia: https://en.wikipedia.org/wiki/Weighted_least_squares
+        # \theta =(X^T w X)^{-1} (X^T w y)
+        XT_W: Float[Tensor, l + 1, p] = X.T * weights
+        epsilon: Float[Tensor, l + 1, l + 1] = (
+            torch.eye(XT_W.shape[0], device=XT_W.device) * 1e-6
+        )  # To prevent inversion errors
+        theta: Float[Tensor, l + 1, t] = torch.inverse(XT_W @ X + epsilon) @ (XT_W @ results)
+
+        # exclude the bias term
+        token_importance: Float[Tensor, t, l] = theta[1:].T
+
+        return token_importance

interpreto/attributions/aggregations/sobol_aggregation.py

@@ -0,0 +1,116 @@
+# MIT License
+#
+# Copyright (c) 2025 IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL and FOR are research programs operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+from enum import Enum
+
+import torch
+from beartype import beartype
+from jaxtyping import Float, jaxtyped
+
+from interpreto.attributions.aggregations.base import Aggregator
+
+
+class SobolIndicesOrders(Enum):
+    """
+    Enumeration of available Sobol indices orders.
+    """
+
+    FIRST_ORDER = "first order"
+    TOTAL_ORDER = "total order"
+
+
+class SobolAggregator(Aggregator):
+    """
+    Aggregates Sobol indices from model outputs.
+    """
+
+    def __init__(
+        self,
+        n_token_perturbations: int,
+        sobol_indices_order: SobolIndicesOrders = SobolIndicesOrders.FIRST_ORDER,
+    ):
+        """
+        Initialize the aggregator.
+
+        Args:
+            n_token_perturbations (int): Number of token perturbations.
+            sobol_indices_order (SobolIndicesOrders): Sobol indices order, either `FIRST_ORDER` or `TOTAL_ORDER`.
+        """
+        self.n_token_perturbations = n_token_perturbations
+        self.sobol_indices_order = sobol_indices_order.value
+
+    @jaxtyped(typechecker=beartype)
+    def aggregate(self, results: Float[torch.Tensor, "p t"], mask: torch.Tensor) -> Float[torch.Tensor, "t l"]:  # noqa: UP037
+        """
+        Compute the Sobol indices from the model outputs perturbed inputs.
+
+        Args:
+            results (torch.Tensor): The model outputs on perturbed inputs. Shape: (p, t) with p = (l + 2) * k
+                - l is the number of elements perturbed.
+                - k is the number of perturbations per element.
+                - t is the targets dimensions, i.e. the number of classes or number of generated tokens.
+            mask (torch.Tensor): Ignored.
+
+        Returns:
+            token_importance (torch.Tensor): The Sobol attribution indices for each token. Shape: (l,)
+        """
+        # simplify typing
+        k = self.n_token_perturbations
+        l = (results.shape[0] // k) - 2
+        t = results.shape[1]
+
+        # Extract initial matrices
+        fA: Float[torch.Tensor, k, t] = results[:k]
+        fB: Float[torch.Tensor, k, t] = results[k : 2 * k]
+        fC: Float[torch.Tensor, l, k, t] = results[2 * k :].view(l, k, t)
+
+        var_fA: Float[torch.Tensor, 1, t] = torch.var(fA, dim=0, keepdim=True)
+
+        token_importance: Float[torch.Tensor, l, t]
+        if self.sobol_indices_order == SobolIndicesOrders.FIRST_ORDER.value:
+            # (var(f(A)) - 0.5*mean(f(B) - f(C))^2) / var(f(A))
+            fBfC_diff: Float[torch.Tensor, l, k, t] = fB.unsqueeze(0) - fC
+            fBfC_mean_square_diff: Float[torch.Tensor, l, t] = torch.mean(fBfC_diff**2, dim=1) / 2
+            token_importance: Float[torch.Tensor, l, t] = (var_fA - fBfC_mean_square_diff) / (var_fA + 1e-6)
+        elif self.sobol_indices_order == SobolIndicesOrders.TOTAL_ORDER.value:
+            # mean(f(A) - f(C))^2) / 2var(f(A))
+            fAfC_diff: Float[torch.Tensor, l, k, t] = fA.unsqueeze(0) - fC
+            fAfC_mean_square_diff: Float[torch.Tensor, l, t] = torch.mean(fAfC_diff**2, dim=1)
+            token_importance: Float[torch.Tensor, l, t] = fAfC_mean_square_diff / (2 * var_fA + 1e-6)
+        else:
+            raise ValueError(f"Unknown Sobol indices order: {self.sobol_indices_order}")
+
+        return token_importance.T
+
+        # # Compute token-wise variance on the initial mask
+        # initial_results: Float[torch.Tensor, t, k] = results[:k].T
+        # initial_var: Float[torch.Tensor, t, 1] = torch.var(initial_results, dim=1, keepdim=True)
+
+        # # Compute token-wise sobol attribution indices
+        # token_results: Float[torch.Tensor, t, l, k] = results[k:].view(t, l, k)
+        # difference: Float[torch.Tensor, t, l, k] = token_results - initial_results.unsqueeze(1)
+        # token_importance: Float[torch.Tensor, t, l] = torch.mean(difference**2, dim=-1) / (initial_var + 1e-6)
+        # return token_importance