From 3012de17270783878388437b515fc95efcbdb6a1 Mon Sep 17 00:00:00 2001
From: Zach Carmichael <craymichael@meta.com>
Date: Fri, 28 Aug 2026 10:47:40 -0700
Subject: [PATCH 1/6] Validate perturbation baselines before attribution
 (#1920)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary:

Summary

Feature Ablation and Shapley formatted baselines but did not validate tuple arity or tensor shape before running the model. `_tensorize_baseline` also paired inputs and baselines with `zip`, silently discarding extra baselines.

Problem

For two inputs, a one-element baseline tuple reached the model with a missing argument, while a three-element tuple silently ignored its final baseline. A baseline pool with shape `[2, 1]` for a three-row input was neither a per-example baseline nor a singleton baseline and failed later through opaque broadcasting errors.

Fix

* Validate baselines in both synchronous and future Feature Ablation and Shapley entry points.
* Defensively reject arity mismatches in `_tensorize_baseline`.
* Require Shapley tensor baselines to match the input or use a singleton leading dimension with matching trailing dimensions.
* Preserve Feature Ablation’s documented support for any tensor shape that broadcasts exactly to the input, including `[F]` and 0-D tensors.

Test Plan

Before: the new regressions failed through late `IndexError`, missing-forward-argument, and broadcast errors; the extra-baseline case did not raise at all.

After:
* `buck test fbcode//pytorch/captum/tests/attr:test_common fbcode//pytorch/captum/tests/attr:test_feature_ablation fbcode//pytorch/captum/tests/attr:test_shapley` — Pass 131, Fail 0.
* `arc lint -a` on changed Python files — no source lint issues; focused autodeps updates applied.
* `arc lint -a --engine extra --take CITRINEAGENT` on changed Captum implementation files — no issues.
* `arc pyre check-owning-targets` on changed files — no type errors.

Differential Revision: D117601314
---
 captum/_utils/common.py               | 17 ++++++++++++-
 captum/attr/_core/feature_ablation.py |  3 +++
 captum/attr/_core/shapley_value.py    |  3 +++
 captum/attr/_utils/common.py          |  4 +++
 tests/attr/test_common.py             |  6 +++++
 tests/attr/test_feature_ablation.py   | 36 +++++++++++++++++++++++++++
 tests/attr/test_shapley.py            | 35 ++++++++++++++++++++++++++
 7 files changed, 103 insertions(+), 1 deletion(-)

diff --git a/captum/_utils/common.py b/captum/_utils/common.py
index cab42bf2cd..703e7a5c28 100644
--- a/captum/_utils/common.py
+++ b/captum/_utils/common.py
@@ -117,6 +117,7 @@ def _validate_input(
     inputs: Tuple[Tensor, ...],
     baselines: Tuple[Union[Tensor, int, float], ...],
     draw_baseline_from_distrib: bool = False,
+    allow_broadcastable_baselines: bool = False,
 ) -> None:
     assert len(inputs) == len(baselines), (
         "Input and baseline must have the same "
@@ -137,10 +138,24 @@ def _validate_input(
                 " Found baseline: {} and input: {} ".format(baseline, input)
             )
         else:
+            baseline_is_broadcastable = False
+            if allow_broadcastable_baselines and isinstance(baseline, Tensor):
+                try:
+                    baseline_is_broadcastable = (
+                        torch.broadcast_shapes(input.shape, baseline.shape)
+                        == input.shape
+                    )
+                except RuntimeError:
+                    pass
             assert (
                 isinstance(baseline, (int, float))
                 or input.shape == baseline.shape
-                or baseline.shape[0] == 1
+                or (
+                    baseline.dim() > 0
+                    and baseline.shape[0] == 1
+                    and input.shape[1:] == baseline.shape[1:]
+                )
+                or baseline_is_broadcastable
             ), (
                 "Baseline can be provided as a tensor for just one input and"
                 " broadcasted to the batch or input and baseline must have the"
diff --git a/captum/attr/_core/feature_ablation.py b/captum/attr/_core/feature_ablation.py
index 505d4cb260..3111a2ab72 100644
--- a/captum/attr/_core/feature_ablation.py
+++ b/captum/attr/_core/feature_ablation.py
@@ -22,6 +22,7 @@
     _is_tuple,
     _maybe_expand_parameters,
     _run_forward,
+    _validate_input,
 )
 from captum._utils.exceptions import FeatureAblationFutureError
 from captum._utils.progress import NullProgress, progress, Progress
@@ -501,6 +502,7 @@ def attribute(
         is_inputs_tuple = _is_tuple(inputs)
 
         formatted_inputs, baselines = _format_input_baseline(inputs, baselines)
+        _validate_input(formatted_inputs, baselines, allow_broadcastable_baselines=True)
         formatted_additional_forward_args = _format_additional_forward_args(
             additional_forward_args
         )
@@ -834,6 +836,7 @@ def attribute_future(
         # converting it into a tuple.
         is_inputs_tuple = _is_tuple(inputs)
         formatted_inputs, baselines = _format_input_baseline(inputs, baselines)
+        _validate_input(formatted_inputs, baselines, allow_broadcastable_baselines=True)
         formatted_additional_forward_args = _format_additional_forward_args(
             additional_forward_args
         )
diff --git a/captum/attr/_core/shapley_value.py b/captum/attr/_core/shapley_value.py
index 32f8c95c2a..7b6a9ded9c 100644
--- a/captum/attr/_core/shapley_value.py
+++ b/captum/attr/_core/shapley_value.py
@@ -37,6 +37,7 @@
     _is_mask_valid,
     _is_tuple,
     _run_forward,
+    _validate_input,
 )
 from captum._utils.exceptions import ShapleyValueFutureError
 from captum._utils.progress import progress
@@ -320,6 +321,7 @@ def attribute(
         # converting it into a tuple.
         is_inputs_tuple = _is_tuple(inputs)
         inputs_tuple, baselines = _format_input_baseline(inputs, baselines)
+        _validate_input(inputs_tuple, baselines)
         additional_forward_args = _format_additional_forward_args(
             additional_forward_args
         )
@@ -487,6 +489,7 @@ def attribute_future(
     ) -> Future[TensorOrTupleOfTensorsGeneric]:
         is_inputs_tuple = _is_tuple(inputs)
         inputs_tuple, baselines = _format_input_baseline(inputs, baselines)
+        _validate_input(inputs_tuple, baselines)
         additional_forward_args = _format_additional_forward_args(
             additional_forward_args
         )
diff --git a/captum/attr/_utils/common.py b/captum/attr/_utils/common.py
index cab1a31050..5ffd3b51ea 100644
--- a/captum/attr/_utils/common.py
+++ b/captum/attr/_utils/common.py
@@ -292,6 +292,10 @@ def _tensorize_single_baseline(baseline, input):
             type(baselines), type(inputs)
         )
     )
+    assert len(inputs) == len(baselines), (
+        "Input and baseline must have the same dimensions, baseline has "
+        f"{len(baselines)} features whereas input has {len(inputs)}."
+    )
     return tuple(
         _tensorize_single_baseline(baseline, input)
         for baseline, input in zip(baselines, inputs)
diff --git a/tests/attr/test_common.py b/tests/attr/test_common.py
index a3e95a3a01..2052383c14 100644
--- a/tests/attr/test_common.py
+++ b/tests/attr/test_common.py
@@ -56,6 +56,12 @@ def test_validate_input(self) -> None:
             (torch.tensor([-1.0]),), (torch.tensor([-2.0]),), method="gausslegendre"
         )
 
+        with self.assertRaisesRegex(AssertionError, "Baseline can be provided"):
+            _validate_input(
+                (torch.zeros((2, 3)),),
+                (torch.zeros((1, 4)),),
+            )
+
     def test_validate_nt_type(self) -> None:
         with self.assertRaises(
             AssertionError,
diff --git a/tests/attr/test_feature_ablation.py b/tests/attr/test_feature_ablation.py
index db293217c9..2917d9bd74 100644
--- a/tests/attr/test_feature_ablation.py
+++ b/tests/attr/test_feature_ablation.py
@@ -131,6 +131,42 @@ def test_simple_ablation_with_baselines(self) -> None:
             perturbations_per_eval=(1, 2, 3),
         )
 
+    def test_perturbation_methods_reject_baseline_tuple_arity_mismatch(self) -> None:
+        inputs = (torch.tensor([[1.0]]), torch.tensor([[2.0]]))
+        attribution = FeatureAblation(lambda first, second: first[:, 0] + second[:, 0])
+
+        for baselines in (
+            (torch.tensor([[0.0]]),),
+            (torch.tensor([[0.0]]),) * 3,
+        ):
+            with self.subTest(baseline_count=len(baselines)):
+                with self.assertRaisesRegex(
+                    AssertionError, "Input and baseline must have the same"
+                ):
+                    attribution.attribute(inputs, baselines=baselines)
+
+    def test_perturbation_methods_reject_invalid_baseline_batch_size(self) -> None:
+        inputs = torch.tensor([[1.0], [2.0], [3.0]])
+
+        attribution = FeatureAblation(lambda values: values[:, 0])
+        for baseline in (
+            torch.tensor([[10.0], [20.0]]),
+            torch.tensor([[10.0, 20.0]]),
+        ):
+            with self.subTest(baseline_shape=baseline.shape):
+                with self.assertRaisesRegex(AssertionError, "Baseline can be provided"):
+                    attribution.attribute(inputs, baselines=baseline)
+
+    def test_accepts_broadcastable_tensor_baselines(self) -> None:
+        inputs = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        attribution = FeatureAblation(lambda values: values.sum(dim=1))
+
+        for baseline in (torch.tensor([0.5, 1.5]), torch.tensor(0.5)):
+            with self.subTest(baseline_shape=baseline.shape):
+                result = attribution.attribute(inputs, baselines=baseline)
+                expected = inputs - baseline
+                torch.testing.assert_close(result, expected)
+
     def test_simple_ablation_boolean(self) -> None:
         ablation_algo = FeatureAblation(BasicModelBoolInput())
         inp = torch.tensor([[True, False, True]])
diff --git a/tests/attr/test_shapley.py b/tests/attr/test_shapley.py
index d7d9db56d1..1784143787 100644
--- a/tests/attr/test_shapley.py
+++ b/tests/attr/test_shapley.py
@@ -29,6 +29,41 @@
 
 
 class Test(BaseTest):
+    def test_rejects_baseline_tuple_arity_mismatch(self) -> None:
+        inputs = (torch.tensor([[1.0]]), torch.tensor([[2.0]]))
+        attribution = ShapleyValueSampling(
+            lambda first, second: first[:, 0] + second[:, 0]
+        )
+
+        for baselines in (
+            (torch.tensor([[0.0]]),),
+            (torch.tensor([[0.0]]),) * 3,
+        ):
+            with self.subTest(baseline_count=len(baselines)):
+                with self.assertRaisesRegex(
+                    AssertionError, "Input and baseline must have the same"
+                ):
+                    attribution.attribute(inputs, baselines=baselines)
+                with self.assertRaisesRegex(
+                    AssertionError, "Input and baseline must have the same"
+                ):
+                    attribution.attribute_future(inputs, baselines=baselines)
+
+    def test_rejects_invalid_baseline_shape(self) -> None:
+        inputs = torch.tensor([[1.0], [2.0], [3.0]])
+        attribution = ShapleyValueSampling(lambda values: values[:, 0])
+
+        for baseline in (
+            torch.tensor([[10.0], [20.0]]),
+            torch.tensor([[10.0, 20.0]]),
+            torch.tensor(10.0),
+        ):
+            with self.subTest(baseline_shape=baseline.shape):
+                with self.assertRaisesRegex(AssertionError, "Baseline can be provided"):
+                    attribution.attribute(inputs, baselines=baseline)
+                with self.assertRaisesRegex(AssertionError, "Baseline can be provided"):
+                    attribution.attribute_future(inputs, baselines=baseline)
+
     @parameterized.expand([True, False])
     def test_simple_shapley_sampling(self, use_future: bool) -> None:
         inp = torch.tensor([[20.0, 50.0, 30.0]], requires_grad=True)

From 98d2ffdab7021e57544fe6ac5793a04b93c52b0e Mon Sep 17 00:00:00 2001
From: Zach Carmichael <craymichael@meta.com>
Date: Fri, 28 Aug 2026 10:47:40 -0700
Subject: [PATCH 2/6] Select perturbed values without arithmetic masking
 (#1921)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary:

Summary

Perturbation paths combined original and replacement values as `old * (1 - mask) + new * mask`. IEEE arithmetic makes `NaN * 0` and `Inf * 0` non-finite, so values outside the selected feature could corrupt model inputs and attributions.

Counterexample

For input `[1, 2]`, baseline `[0, NaN]`, and mask `[0, 1]`, ablating feature 0 should evaluate `[0, 2]`. Arithmetic masking instead produced `[0, NaN]`, so feature 0 incorrectly received a `NaN` attribution.

Fix

* Use `torch.where` for Feature Ablation, Feature Permutation, Shapley feature updates, and attribution-mask accumulation.
* Apply the same selection semantics to within-group baseline construction, add-back, and permutation helpers.
* Mirror the fix in the legacy UFO implementations so alternate call paths cannot retain the corruption.
* Keep masks and donor indices on the destination tensor device.

Test Plan

Before: six focused regressions failed across core and within-group paths (`Pass 159, Fail 6`).

After:
* `buck test fbcode//pytorch/captum/tests/attr:test_feature_ablation fbcode//pytorch/captum/tests/attr:test_feature_permutation fbcode//pytorch/captum/tests/attr:test_shapley fbcode//pytorch/captum/tests/attr/fb:test_within_groups_utils fbcode//pytorch/captum/tests/attr/fb:test_shapley_value_permutation` — Pass 178, Fail 0.
* Regressions cover `NaN`, `+Inf`, and `-Inf` in selected and inactive positions.
* `arc lint -a` on all changed Python files — no new lint issues; existing legacy UFO line-length advice remains unchanged.
* `arc lint -a --engine extra --take CITRINEAGENT` on implementation files — no issues.
* `arc pyre check-owning-targets` on changed files — no type errors.

Differential Revision: D117601315
---
 captum/attr/_core/feature_ablation.py    | 12 ++++--
 captum/attr/_core/feature_permutation.py |  4 +-
 captum/attr/_core/shapley_value.py       | 15 ++++---
 tests/attr/test_feature_ablation.py      | 14 +++++++
 tests/attr/test_feature_permutation.py   | 14 +++++++
 tests/attr/test_shapley.py               | 51 ++++++++++++++++++++++++
 6 files changed, 98 insertions(+), 12 deletions(-)

diff --git a/captum/attr/_core/feature_ablation.py b/captum/attr/_core/feature_ablation.py
index 3111a2ab72..b1983270a3 100644
--- a/captum/attr/_core/feature_ablation.py
+++ b/captum/attr/_core/feature_ablation.py
@@ -784,9 +784,9 @@ def _construct_ablated_input_across_tensors(
                 tensor_mask.append(mask)
 
                 assert baseline is not None, "baseline must be provided"
-                ablated_feature = input_tensor[start_idx:end_idx] * (1 - mask).to(
-                    input_tensor.dtype
-                ) + (baseline * mask.to(input_tensor.dtype))
+                ablated_feature = torch.where(
+                    mask.bool(), baseline, input_tensor[start_idx:end_idx]
+                )
                 ablated_input = ablated_input.to(ablated_feature.dtype)
                 ablated_input[start_idx:end_idx] = ablated_feature
             current_masks.append(torch.stack(tensor_mask, dim=0))
@@ -1211,6 +1211,10 @@ def _process_ablated_out_full(
                 eval_diff_shape + (inputs[i].dim() - 1) * (1,)
             )
             eval_diff = eval_diff.to(total_attrib[i].device)
-            total_attrib[i] += (eval_diff * mask.to(attrib_type)).sum(dim=0)
+            total_attrib[i] += torch.where(
+                mask.to(device=eval_diff.device, dtype=torch.bool),
+                eval_diff,
+                torch.zeros_like(eval_diff),
+            ).sum(dim=0)
 
         return total_attrib, weights
diff --git a/captum/attr/_core/feature_permutation.py b/captum/attr/_core/feature_permutation.py
index fdb383639f..d3f3f4ef9a 100644
--- a/captum/attr/_core/feature_permutation.py
+++ b/captum/attr/_core/feature_permutation.py
@@ -26,9 +26,7 @@ def _permute_feature(x: Tensor, feature_mask: Tensor) -> Tensor:
     while (perm == no_perm).all():
         perm = torch.randperm(n)
 
-    return (x[perm] * feature_mask.to(dtype=x.dtype)) + (
-        x * feature_mask.bitwise_not().to(dtype=x.dtype)
-    )
+    return torch.where(feature_mask.to(device=x.device, dtype=torch.bool), x[perm], x)
 
 
 class FeaturePermutation(FeatureAblation):
diff --git a/captum/attr/_core/shapley_value.py b/captum/attr/_core/shapley_value.py
index 7b6a9ded9c..e87eee7529 100644
--- a/captum/attr/_core/shapley_value.py
+++ b/captum/attr/_core/shapley_value.py
@@ -757,7 +757,11 @@ def _eval_fut_to_prev_results_tuple(
             )
 
             # aggregate n_perturb
-            cur_attr = (formatted_eval_diff * cur_mask.float()).sum(dim=0)
+            cur_attr = torch.where(
+                cur_mask.to(device=formatted_eval_diff.device, dtype=torch.bool),
+                formatted_eval_diff,
+                torch.zeros_like(formatted_eval_diff),
+            ).sum(dim=0)
             # (*output_shape, *input_feature_shape)
             total_attrib[j] += cur_attr
 
@@ -808,10 +812,11 @@ def _update_current_tensors(
         for i in range(len(current_tensors)):
             if i in feat_list:
                 output_tensors.append(
-                    current_tensors[i]
-                    * (~(mask[i] == feature_index)).to(current_tensors[i].dtype)
-                    + input_tensors[i]
-                    * (mask[i] == feature_index).to(input_tensors[i].dtype)
+                    torch.where(
+                        (mask[i] == feature_index).to(current_tensors[i].device),
+                        input_tensors[i],
+                        current_tensors[i],
+                    )
                 )
 
             else:
diff --git a/tests/attr/test_feature_ablation.py b/tests/attr/test_feature_ablation.py
index 2917d9bd74..b0ba30a9e8 100644
--- a/tests/attr/test_feature_ablation.py
+++ b/tests/attr/test_feature_ablation.py
@@ -167,6 +167,20 @@ def test_accepts_broadcastable_tensor_baselines(self) -> None:
                 expected = inputs - baseline
                 torch.testing.assert_close(result, expected)
 
+    def test_inactive_nonfinite_baseline_does_not_contaminate_ablation(self) -> None:
+        inputs = torch.tensor([[1.0, 2.0]])
+        attribution = FeatureAblation(lambda values: values.sum(dim=1))
+
+        for inactive_value in (float("nan"), float("inf"), float("-inf")):
+            with self.subTest(inactive_value=inactive_value):
+                result = attribution.attribute(
+                    inputs,
+                    baselines=torch.tensor([[0.0, inactive_value]]),
+                    feature_mask=torch.tensor([[0, 1]]),
+                )
+
+                self.assertEqual(result[0, 0].item(), 1.0)
+
     def test_simple_ablation_boolean(self) -> None:
         ablation_algo = FeatureAblation(BasicModelBoolInput())
         inp = torch.tensor([[True, False, True]])
diff --git a/tests/attr/test_feature_permutation.py b/tests/attr/test_feature_permutation.py
index 5a3a5036a9..abb798a4e1 100644
--- a/tests/attr/test_feature_permutation.py
+++ b/tests/attr/test_feature_permutation.py
@@ -7,6 +7,7 @@
 
 # pyre-strict
 
+import unittest.mock
 from typing import Any, Callable, List, Tuple
 
 import torch
@@ -95,6 +96,19 @@ def test_perm_fn_broadcastable_masks(self) -> None:
 
             self._check_perm_fn_with_mask(inp, mask)
 
+    def test_inactive_nonfinite_donor_does_not_contaminate_permutation(self) -> None:
+        mask = torch.tensor([False, True])
+
+        for inactive_value in (float("nan"), float("inf"), float("-inf")):
+            with self.subTest(inactive_value=inactive_value):
+                inp = torch.tensor([[inactive_value, 2.0], [1.0, 3.0]])
+                with unittest.mock.patch(
+                    "torch.randperm", return_value=torch.tensor([1, 0])
+                ):
+                    result = _permute_feature(inp, mask)
+
+                self.assertEqual(result[1, 0].item(), 1.0)
+
     def test_single_input(self) -> None:
         batch_size = 2
         input_size = (6,)
diff --git a/tests/attr/test_shapley.py b/tests/attr/test_shapley.py
index 1784143787..054be5765f 100644
--- a/tests/attr/test_shapley.py
+++ b/tests/attr/test_shapley.py
@@ -64,6 +64,57 @@ def test_rejects_invalid_baseline_shape(self) -> None:
                 with self.assertRaisesRegex(AssertionError, "Baseline can be provided"):
                     attribution.attribute_future(inputs, baselines=baseline)
 
+    def test_selected_nonfinite_baseline_is_replaced_cleanly(self) -> None:
+        attribution = ShapleyValueSampling(lambda values: values.sum(dim=1))
+
+        for selected_value in (float("nan"), float("inf"), float("-inf")):
+            with self.subTest(selected_value=selected_value):
+                result = attribution._update_current_tensors(
+                    (torch.tensor([[0.0, selected_value]]),),
+                    (torch.tensor([[1.0, 2.0]]),),
+                    1,
+                    (torch.tensor([[0, 1]]),),
+                    {1: [0]},
+                )
+
+                torch.testing.assert_close(result[0], torch.tensor([[0.0, 2.0]]))
+
+    def test_nonfinite_marginal_does_not_contaminate_other_features(self) -> None:
+        attribution = ShapleyValueSampling(lambda values: values.sum(dim=1))
+        previous_result: Future[
+            Tuple[Tensor, Tensor, torch.Size, List[Tensor], bool]
+        ] = Future()
+        previous_result.set_result(
+            (
+                torch.tensor([0.0]),
+                torch.tensor([0.0]),
+                torch.Size([1]),
+                [torch.zeros((1, 2))],
+                False,
+            )
+        )
+        modified_result: Future[Tensor] = Future()
+        modified_result.set_result(torch.tensor([1.0, float("nan")]))
+        evaluations: Future[
+            List[
+                Union[
+                    Future[Tuple[Tensor, Tensor, torch.Size, List[Tensor], bool]],
+                    Future[Tensor],
+                ]
+            ]
+        ] = Future()
+        evaluations.set_result([previous_result, modified_result])
+
+        result = attribution._eval_fut_to_prev_results_tuple(
+            evaluations,
+            1,
+            (torch.zeros((1, 2)),),
+            (torch.tensor([[[1, 0]], [[0, 1]]]),),
+        )
+
+        self.assertEqual(result[3][0][0, 0].item(), 1.0)
+        self.assertTrue(torch.isnan(result[3][0][0, 1]))
+
     @parameterized.expand([True, False])
     def test_simple_shapley_sampling(self, use_future: bool) -> None:
         inp = torch.tensor([[20.0, 50.0, 30.0]], requires_grad=True)

From 1da1fd260945cc13b2aaf8f32c9093c34aea33a0 Mon Sep 17 00:00:00 2001
From: Zach Carmichael <craymichael@meta.com>
Date: Fri, 28 Aug 2026 10:47:40 -0700
Subject: [PATCH 3/6] Validate perturbation feature masks before attribution
 (#1922)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary:

Summary

Perturbation methods accepted feature-mask tuples with the wrong arity. Feature Ablation could then return zero attribution for an omitted input, while Shapley failed later or enumerated an invalid feature space. Shapley also accepted fractional, negative, non-finite, and complex group IDs even though its permutation loop requires non-negative integer IDs.

Counterexample

With two model inputs and a one-element feature-mask tuple, Feature Ablation perturbed only the first input and silently left the second input’s attribution at zero. With Shapley mask `[0, 0.5]`, the integer feature loop omitted group `0.5`.

Fix

* Require explicit feature-mask tuples to have one tensor per input in the shared formatter.
* Validate Shapley group IDs before feature enumeration.
* Accept integral-valued float masks and bool masks for compatibility.
* Keep value-domain validation out of Feature Ablation/Permutation so Greedy Feature Selection’s internal `-1` sentinel and existing float-typed masks continue to work.

Test Plan

Before: four focused regression methods failed (`Pass 132, Fail 4`), including silent acceptance of short and long mask tuples and invalid Shapley IDs.

After:
* Broad perturbation and wrapper suite — Pass 422, Fail 0, covering Feature Ablation, Feature Permutation, Shapley, DataLoaderAttribution, internal wrappers, WithinGroupSVS, SVS-P, AddOneBack, MarginalWithinGroups, Greedy Feature Selection, and shared mask utilities.
* `arc lint -a` on changed files — no issues.
* `arc lint -a --engine extra --take CITRINEAGENT` on implementation files — no issues.
* `arc pyre check-owning-targets` on changed files — no type errors.

Differential Revision: D117601317
---
 captum/_utils/common.py             |  4 +++
 captum/attr/_core/shapley_value.py  | 17 ++++++++++
 tests/attr/test_feature_ablation.py | 14 +++++++++
 tests/attr/test_shapley.py          | 49 +++++++++++++++++++++++++++++
 4 files changed, 84 insertions(+)

diff --git a/captum/_utils/common.py b/captum/_utils/common.py
index 703e7a5c28..230046f266 100644
--- a/captum/_utils/common.py
+++ b/captum/_utils/common.py
@@ -241,6 +241,10 @@ def _format_feature_mask(
 
     else:
         formatted_mask = _format_tensor_into_tuples(feature_mask)
+        assert len(formatted_mask) == len(inputs), (
+            "Input and feature mask must have the same number of tensors, "
+            f"but input has {len(inputs)} and feature mask has {len(formatted_mask)}."
+        )
 
     return formatted_mask
 
diff --git a/captum/attr/_core/shapley_value.py b/captum/attr/_core/shapley_value.py
index e87eee7529..d2065af107 100644
--- a/captum/attr/_core/shapley_value.py
+++ b/captum/attr/_core/shapley_value.py
@@ -84,6 +84,21 @@ def _shape_feature_mask(
     return tuple(mask_list)
 
 
+def _validate_shapley_feature_mask(feature_mask: Tuple[Tensor, ...]) -> None:
+    for mask in feature_mask:
+        if mask.is_complex():
+            values_are_integral = False
+        elif mask.is_floating_point():
+            values_are_integral = bool(
+                torch.isfinite(mask).all() and (mask == mask.round()).all()
+            )
+        else:
+            values_are_integral = True
+        assert values_are_integral and (
+            mask.numel() == 0 or bool((mask >= 0).all())
+        ), "Feature mask values must be non-negative integers."
+
+
 class ShapleyValueSampling(PerturbationAttribution):
     """
     A perturbation based approach to compute attribution, based on the concept
@@ -326,6 +341,7 @@ def attribute(
             additional_forward_args
         )
         formatted_feature_mask = _format_feature_mask(feature_mask, inputs_tuple)
+        _validate_shapley_feature_mask(formatted_feature_mask)
         reshaped_feature_mask = _shape_feature_mask(
             formatted_feature_mask, inputs_tuple
         )
@@ -494,6 +510,7 @@ def attribute_future(
             additional_forward_args
         )
         formatted_feature_mask = _format_feature_mask(feature_mask, inputs_tuple)
+        _validate_shapley_feature_mask(formatted_feature_mask)
         reshaped_feature_mask = _shape_feature_mask(
             formatted_feature_mask, inputs_tuple
         )
diff --git a/tests/attr/test_feature_ablation.py b/tests/attr/test_feature_ablation.py
index b0ba30a9e8..0df43a0a03 100644
--- a/tests/attr/test_feature_ablation.py
+++ b/tests/attr/test_feature_ablation.py
@@ -181,6 +181,20 @@ def test_inactive_nonfinite_baseline_does_not_contaminate_ablation(self) -> None
 
                 self.assertEqual(result[0, 0].item(), 1.0)
 
+    def test_rejects_feature_mask_tuple_arity_mismatch(self) -> None:
+        inputs = (torch.tensor([[1.0]]), torch.tensor([[2.0]]))
+        attribution = FeatureAblation(lambda first, second: first[:, 0] + second[:, 0])
+
+        for feature_mask in (
+            (torch.tensor([[0]]),),
+            (torch.tensor([[0]]),) * 3,
+        ):
+            with self.subTest(mask_count=len(feature_mask)):
+                with self.assertRaisesRegex(
+                    AssertionError, "Input and feature mask must have the same"
+                ):
+                    attribution.attribute(inputs, feature_mask=feature_mask)
+
     def test_simple_ablation_boolean(self) -> None:
         ablation_algo = FeatureAblation(BasicModelBoolInput())
         inp = torch.tensor([[True, False, True]])
diff --git a/tests/attr/test_shapley.py b/tests/attr/test_shapley.py
index 054be5765f..81c2cc552a 100644
--- a/tests/attr/test_shapley.py
+++ b/tests/attr/test_shapley.py
@@ -115,6 +115,55 @@ def test_nonfinite_marginal_does_not_contaminate_other_features(self) -> None:
         self.assertEqual(result[3][0][0, 0].item(), 1.0)
         self.assertTrue(torch.isnan(result[3][0][0, 1]))
 
+    def test_rejects_feature_mask_tuple_arity_mismatch(self) -> None:
+        inputs = (torch.tensor([[1.0]]), torch.tensor([[2.0]]))
+        attribution = ShapleyValueSampling(
+            lambda first, second: first[:, 0] + second[:, 0]
+        )
+
+        for feature_mask in (
+            (torch.tensor([[0]]),),
+            (torch.tensor([[0]]),) * 3,
+        ):
+            with self.subTest(mask_count=len(feature_mask)):
+                for attribute in (
+                    attribution.attribute,
+                    attribution.attribute_future,
+                ):
+                    with self.assertRaisesRegex(
+                        AssertionError, "Input and feature mask must have the same"
+                    ):
+                        attribute(inputs, feature_mask=feature_mask)
+
+    def test_rejects_nonintegral_or_negative_feature_masks(self) -> None:
+        inputs = torch.tensor([[1.0, 2.0]])
+        attribution = ShapleyValueSampling(lambda values: values.sum(dim=1))
+
+        for feature_mask in (
+            torch.tensor([[0.0, 0.5]]),
+            torch.tensor([[0, -1]]),
+            torch.tensor([[0.0, float("nan")]]),
+            torch.tensor([[0.0, float("inf")]]),
+            torch.tensor([[0.0 + 0.0j, 1.0 + 0.0j]]),
+        ):
+            with self.subTest(feature_mask=feature_mask):
+                for attribute in (
+                    attribution.attribute,
+                    attribution.attribute_future,
+                ):
+                    with self.assertRaisesRegex(
+                        AssertionError, "non-negative integers"
+                    ):
+                        attribute(inputs, feature_mask=feature_mask)
+
+        for feature_mask in (
+            torch.tensor([[0.0, 1.0]]),
+            torch.tensor([[False, True]]),
+        ):
+            with self.subTest(valid_feature_mask=feature_mask):
+                result = attribution.attribute(inputs, feature_mask=feature_mask)
+                torch.testing.assert_close(result, inputs)
+
     @parameterized.expand([True, False])
     def test_simple_shapley_sampling(self, use_future: bool) -> None:
         inp = torch.tensor([[20.0, 50.0, 30.0]], requires_grad=True)

From 3053973e6a77acc797d963f58cd8fc50571c1190 Mon Sep 17 00:00:00 2001
From: Zach Carmichael <craymichael@meta.com>
Date: Fri, 28 Aug 2026 10:47:40 -0700
Subject: [PATCH 4/6] Preserve fractional scalar baselines for integer inputs
 (#1924)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary:

Summary

`_tensorize_baseline` used `torch.full_like(input, baseline)`, which silently cast a floating-point scalar baseline to the input dtype. Integer inputs therefore lost fractional baseline values before Shapley attribution.

Counterexample

For integer input `[1, 2]` and scalar baseline `0.5`, additive Shapley attribution should be `[0.5, 1.5]`. The old tensorization converted the baseline to `[0, 0]` and returned `[1, 2]`.

Fix

Use PyTorch result-type promotion when materializing scalar baselines, while retaining `full_like` so device, layout, and memory format are preserved. Boolean inputs keep their historical boolean baseline dtype.

Test Plan

Before: the new sync regression failed with both elements off by `0.5`.

After:
* `buck test` across Shapley, WithinGroupSVS, DeepLift, LayerDeepLift, LayerIntegratedGradients, IntegratedGradients, and LayerConductance consumers — Pass 196, Fail 0.
* Focused common/Shapley/WithinGroupSVS rerun — Pass 97, Fail 0.
* Unit coverage verifies fractional integer promotion, boolean dtype preservation, and channels-last layout preservation.
* `arc lint -a` — no issues.
* `arc lint -a --engine extra --take CITRINEAGENT` — no issues.
* `arc pyre check-owning-targets` — no type errors.

Differential Revision: D117601321
---
 captum/attr/_utils/common.py | 11 ++++++++++-
 tests/attr/test_common.py    | 23 ++++++++++++++++++++++-
 tests/attr/test_shapley.py   | 19 +++++++++++++++++++
 3 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/captum/attr/_utils/common.py b/captum/attr/_utils/common.py
index 5ffd3b51ea..8ae09d043a 100644
--- a/captum/attr/_utils/common.py
+++ b/captum/attr/_utils/common.py
@@ -281,7 +281,16 @@ def _tensorize_baseline(
     # pyre-fixme[2]: Parameter must be annotated.
     def _tensorize_single_baseline(baseline, input):
         if isinstance(baseline, (int, float)):
-            return torch.full_like(input, baseline)
+            baseline_dtype = (
+                input.dtype
+                if input.dtype == torch.bool
+                else torch.result_type(input, baseline)
+            )
+            return torch.full_like(
+                input,
+                baseline,
+                dtype=baseline_dtype,
+            )
         if input.shape[0] > baseline.shape[0] and baseline.shape[0] == 1:
             return torch.cat([baseline] * input.shape[0])
         return baseline
diff --git a/tests/attr/test_common.py b/tests/attr/test_common.py
index 2052383c14..cf9afed3b7 100644
--- a/tests/attr/test_common.py
+++ b/tests/attr/test_common.py
@@ -9,12 +9,33 @@
 
 import torch
 from captum.attr._core.noise_tunnel import SUPPORTED_NOISE_TUNNEL_TYPES
-from captum.attr._utils.common import _validate_input, _validate_noise_tunnel_type
+from captum.attr._utils.common import (
+    _tensorize_baseline,
+    _validate_input,
+    _validate_noise_tunnel_type,
+)
 from captum.testing.helpers import BaseTest
 
 
 # pyrefly: ignore [invalid-inheritance]
 class Test(BaseTest):
+    def test_tensorize_scalar_baseline_preserves_value_dtype_and_layout(self) -> None:
+        integer_input = torch.tensor([[1, 2]])
+        integer_baseline = _tensorize_baseline((integer_input,), (0.5,))[0]
+        torch.testing.assert_close(integer_baseline, torch.tensor([[0.5, 0.5]]))
+
+        boolean_input = torch.tensor([[True, False]])
+        boolean_baseline = _tensorize_baseline((boolean_input,), (0,))[0]
+        self.assertEqual(boolean_baseline.dtype, torch.bool)
+
+        channels_last_input = torch.empty((2, 3, 4, 5)).to(
+            memory_format=torch.channels_last
+        )
+        channels_last_baseline = _tensorize_baseline((channels_last_input,), (0.5,))[0]
+        self.assertTrue(
+            channels_last_baseline.is_contiguous(memory_format=torch.channels_last)
+        )
+
     def test_validate_input(self) -> None:
         with self.assertRaises(AssertionError) as err:
             _validate_input(
diff --git a/tests/attr/test_shapley.py b/tests/attr/test_shapley.py
index 81c2cc552a..bd91635dba 100644
--- a/tests/attr/test_shapley.py
+++ b/tests/attr/test_shapley.py
@@ -64,6 +64,25 @@ def test_rejects_invalid_baseline_shape(self) -> None:
                 with self.assertRaisesRegex(AssertionError, "Baseline can be provided"):
                     attribution.attribute_future(inputs, baselines=baseline)
 
+    def test_float_scalar_baseline_preserves_dtype_for_integer_input(self) -> None:
+        inputs = torch.tensor([[1, 2]])
+        expected = torch.tensor([[0.5, 1.5]])
+
+        attribution = ShapleyValueSampling(lambda values: values.sum(dim=1))
+        result = attribution.attribute(inputs, baselines=0.5, n_samples=1)
+        torch.testing.assert_close(result, expected)
+
+        def future_sum(values: Tensor) -> Future[Tensor]:
+            result: Future[Tensor] = Future()
+            result.set_result(values.sum(dim=1))
+            return result
+
+        future_attribution = ShapleyValueSampling(future_sum)
+        future_result = future_attribution.attribute_future(
+            inputs, baselines=0.5, n_samples=1
+        ).wait()
+        torch.testing.assert_close(future_result, expected)
+
     def test_selected_nonfinite_baseline_is_replaced_cleanly(self) -> None:
         attribution = ShapleyValueSampling(lambda values: values.sum(dim=1))
 

From 20aa92c7ec23f5602d6e985e3584a8453ce8354b Mon Sep 17 00:00:00 2001
From: Zach Carmichael <craymichael@meta.com>
Date: Fri, 28 Aug 2026 10:47:40 -0700
Subject: [PATCH 5/6] Add an opt-in no-self donor policy (#1926)

Summary:

Problem
- Standard permutation may select a row as its own donor, violating explicitly requested leave-one-out semantics.
- The initial implementation changed scripted defaults, duplicated retry logic, and leaked arbitrary config fields.
- Rank Suggest then hard-coded --n-samples to SVS-P, regressing plain SVS.

Fix
- Preserve standard defaults and expose public exclude_self_donors options.
- Generate one device-side random cycle and share donor ordering across grouped/nested tensors.
- Route only supported kwargs with typed validation.
- Determine --n-samples support from the selected args dataclass fields, preserving plain SVS and Kernel SHAP while rejecting shuffling.

Differential Revision: D117601323
---
 captum/_utils/common.py                  | 18 +++++++
 captum/attr/_core/feature_permutation.py | 36 ++++++++++++--
 tests/attr/test_feature_permutation.py   | 60 +++++++++++++++++++++++-
 3 files changed, 110 insertions(+), 4 deletions(-)

diff --git a/captum/_utils/common.py b/captum/_utils/common.py
index 230046f266..eeb3168120 100644
--- a/captum/_utils/common.py
+++ b/captum/_utils/common.py
@@ -36,6 +36,24 @@
 from torch.nn import Module
 
 
+def _generate_derangement(n: int, device: torch.device | None = None) -> Tensor:
+    """Generate a random cyclic donor ordering without fixed points.
+
+    A random cycle gives every row a uniformly distributed donor among the
+    other rows while requiring only one device-side ``randperm`` and no
+    device-to-host synchronization. The result is a valid derangement for
+    every ``n > 1``.
+    """
+    row_indices = torch.arange(n, device=device)
+    if n <= 1:
+        return row_indices
+
+    cycle = torch.randperm(n, device=device)
+    permutation = torch.empty_like(cycle)
+    permutation[cycle] = torch.roll(cycle, shifts=-1)
+    return permutation
+
+
 def parse_version(v: str) -> Tuple[int, ...]:
     """
     Parse version strings into tuples for comparison.
diff --git a/captum/attr/_core/feature_permutation.py b/captum/attr/_core/feature_permutation.py
index d3f3f4ef9a..37498bf08a 100644
--- a/captum/attr/_core/feature_permutation.py
+++ b/captum/attr/_core/feature_permutation.py
@@ -9,7 +9,11 @@
 from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Union
 
 import torch
-from captum._utils.common import _format_output, _format_tensor_into_tuples
+from captum._utils.common import (
+    _format_output,
+    _format_tensor_into_tuples,
+    _generate_derangement,
+)
 from captum._utils.typing import BaselineType, TargetType, TensorOrTupleOfTensorsGeneric
 from captum.attr._core.feature_ablation import FeatureAblation
 from captum.log import log_usage
@@ -29,6 +33,15 @@ def _permute_feature(x: Tensor, feature_mask: Tensor) -> Tensor:
     return torch.where(feature_mask.to(device=x.device, dtype=torch.bool), x[perm], x)
 
 
+def _permute_feature_without_self_donors(x: Tensor, feature_mask: Tensor) -> Tensor:
+    n = x.size(0)
+    assert n > 1, "cannot permute features with batch_size = 1"
+
+    perm = _generate_derangement(n, x.device)
+
+    return torch.where(feature_mask.to(device=x.device, dtype=torch.bool), x[perm], x)
+
+
 class FeaturePermutation(FeatureAblation):
     r"""
     A perturbation based approach to compute attribution, which
@@ -58,6 +71,11 @@ class FeaturePermutation(FeatureAblation):
     This method, unlike other attribution methods, requires a batch
     of examples to compute attributions and cannot be performed on a single example.
 
+    The default permutation may leave individual rows fixed. Set
+    ``exclude_self_donors=True`` to require every selected value to come from a
+    different row. Distinct donor rows may still contain equal feature values;
+    those are valid samples from the empirical feature distribution.
+
     By default, each scalar value within
     each input tensor is taken as a feature and shuffled independently. Passing
     a feature mask allows grouping features to be shuffled together (including
@@ -80,7 +98,8 @@ class FeaturePermutation(FeatureAblation):
     def __init__(
         self,
         forward_func: Callable[..., Union[int, float, Tensor, Future[Tensor]]],
-        perm_func: Callable[[Tensor, Tensor], Tensor] = _permute_feature,
+        perm_func: Callable[[Tensor, Tensor], Tensor] | None = None,
+        exclude_self_donors: bool = False,
     ) -> None:
         r"""
         Args:
@@ -93,9 +112,20 @@ def __init__(
                 which applies a random permutation, this argument only needs
                 to be provided if a custom permutation behavior is desired.
                 Default: `_permute_feature`
+            exclude_self_donors (bool, optional): If True, every selected row
+                receives its value from a different donor row. This cannot be
+                combined with a custom ``perm_func``. Default: False
         """
         FeatureAblation.__init__(self, forward_func=forward_func)
-        self.perm_func = perm_func
+        if exclude_self_donors and perm_func is not None:
+            raise ValueError(
+                "exclude_self_donors cannot be combined with a custom perm_func"
+            )
+        self.perm_func = (
+            _permute_feature_without_self_donors
+            if exclude_self_donors
+            else perm_func if perm_func is not None else _permute_feature
+        )
         # Considering the case when we permute multiple input tensors at once
         # through `feature_mask`, we disregard the feature group if the 0th
         # dim of *any* input tensor in the group is less than
diff --git a/tests/attr/test_feature_permutation.py b/tests/attr/test_feature_permutation.py
index abb798a4e1..c8b27df05e 100644
--- a/tests/attr/test_feature_permutation.py
+++ b/tests/attr/test_feature_permutation.py
@@ -11,7 +11,12 @@
 from typing import Any, Callable, List, Tuple
 
 import torch
-from captum.attr._core.feature_permutation import _permute_feature, FeaturePermutation
+from captum.attr._core.feature_permutation import (
+    _generate_derangement,
+    _permute_feature,
+    _permute_feature_without_self_donors,
+    FeaturePermutation,
+)
 from captum.testing.helpers import BaseTest
 from captum.testing.helpers.basic import assertTensorAlmostEqual, set_all_random_seeds
 from captum.testing.helpers.basic_models import BasicModelWithSparseInputs
@@ -109,6 +114,59 @@ def test_inactive_nonfinite_donor_does_not_contaminate_permutation(self) -> None
 
                 self.assertEqual(result[1, 0].item(), 1.0)
 
+    def test_deranged_permutation_rejects_self_donors(self) -> None:
+        inp = torch.tensor([[1.0], [2.0], [3.0]])
+
+        with unittest.mock.patch(
+            "torch.randperm", return_value=torch.tensor([0, 1, 2])
+        ) as randperm:
+            result = _permute_feature_without_self_donors(inp, torch.tensor([True]))
+
+        self.assertEqual(randperm.call_count, 1)
+        torch.testing.assert_close(result, torch.tensor([[2.0], [3.0], [1.0]]))
+
+    def test_feature_permutation_exposes_no_self_donor_policy(self) -> None:
+        attribution = FeaturePermutation(
+            lambda inputs: inputs.sum(dim=1),
+            exclude_self_donors=True,
+        )
+
+        with unittest.mock.patch(
+            "torch.randperm", return_value=torch.tensor([0, 1, 2])
+        ):
+            result = attribution.perm_func(
+                torch.tensor([[1.0], [2.0], [3.0]]),
+                torch.tensor([True]),
+            )
+
+        torch.testing.assert_close(result, torch.tensor([[2.0], [3.0], [1.0]]))
+
+    def test_feature_permutation_rejects_conflicting_no_self_configuration(
+        self,
+    ) -> None:
+        with self.assertRaisesRegex(ValueError, "custom perm_func"):
+            FeaturePermutation(
+                lambda inputs: inputs.sum(dim=1),
+                perm_func=self._deterministic_perm,
+                exclude_self_donors=True,
+            )
+
+    def test_distinct_donor_can_have_equal_feature_value(self) -> None:
+        inp = torch.tensor([[1.0], [1.0]])
+
+        result = _permute_feature_without_self_donors(inp, torch.tensor([True]))
+
+        torch.testing.assert_close(result, inp)
+
+    @unittest.mock.patch("torch.randperm", return_value=torch.tensor([2, 1, 0]))
+    def test_derangement_uses_one_random_cycle(
+        self, mock_randperm: unittest.mock.MagicMock
+    ) -> None:
+        permutation = _generate_derangement(3)
+
+        torch.testing.assert_close(permutation, torch.tensor([2, 0, 1]))
+        self.assertEqual(mock_randperm.call_count, 1)
+
     def test_single_input(self) -> None:
         batch_size = 2
         input_size = (6,)

From 4a350e9c9e84b6270d92731900082762ed4acb18 Mon Sep 17 00:00:00 2001
From: Zach Carmichael <craymichael@meta.com>
Date: Fri, 28 Aug 2026 10:47:40 -0700
Subject: [PATCH 6/6] Weight distributed attribution aggregation by samples
 (#1925)

Summary:

- Weight rank/batch attribution means and coverage by their actual example counts.
- Infer the local count from supervision tensors, including uneven final or filtered batches, with the configured size only as a fallback.
- Gather counts alongside results and skip zero-sample ranks without skipping collectives.
- Scope every barrier, gather, and reduction to the supplied process group and translate subgroup rank 0 to its global destination rank.
- Extend Captum online mean, variance, standard deviation, and sum statistics with backward-compatible frequency weights.

Counterexample:
Rank 0 reports mean attribution 1 from one example. Rank 1 reports mean attribution 3 from three examples. Equal rank weighting returns 2.0; the correct example-weighted global mean is `(1 * 1 + 3 * 3) / 4 = 2.5`.

Implementation:
The evaluator derives each local batch count from label/weight tensors. The publisher gathers those counts in the same process group as attribution tensors, and the singleton summarizer applies them as frequency weights. Existing unweighted callers retain weight 1 and custom Stat implementations still receive their legacy one-argument `update` call.

Differential Revision: D117601324
---
 captum/attr/_utils/stat.py       | 30 ++++++++++++++++--------------
 captum/attr/_utils/summarizer.py | 23 +++++++++++++++++++----
 tests/attr/test_stat.py          | 26 ++++++++++++++++++++++++++
 3 files changed, 61 insertions(+), 18 deletions(-)

diff --git a/captum/attr/_utils/stat.py b/captum/attr/_utils/stat.py
index 81648de747..f704d711a9 100644
--- a/captum/attr/_utils/stat.py
+++ b/captum/attr/_utils/stat.py
@@ -62,7 +62,7 @@ def _get_stat(self, stat: _S) -> Optional[_S]:
         assert self._other_stats is not None
         return cast(Optional[_S], self._other_stats.get(stat))
 
-    def update(self, x: Tensor) -> None:
+    def update(self, x: Tensor, weight: float = 1) -> None:
         raise NotImplementedError()
 
     def get(self) -> Optional[StatValue]:
@@ -100,21 +100,20 @@ def name(self) -> str:
 
 class Count(Stat):
     """
-    Counts the number of elements, i.e. the
-    number of `update`'s called
+    Counts observations, including frequency weights supplied to `update`.
     """
 
     def __init__(self, name: Optional[str] = None) -> None:
         super().__init__(name=name)
-        self.n: Optional[int] = None
+        self.n: Optional[Union[int, float]] = None
 
-    def get(self) -> Optional[int]:
+    def get(self) -> Optional[Union[int, float]]:
         return self.n
 
-    def update(self, x: Tensor) -> None:
+    def update(self, x: Tensor, weight: float = 1) -> None:
         if self.n is None:
             self.n = 0
-        self.n += 1
+        self.n += weight
 
 
 class Mean(Stat):
@@ -133,7 +132,7 @@ def get(self) -> Optional[Tensor]:
     def init(self) -> None:
         self.n = self._get_stat(Count())
 
-    def update(self, x: Tensor) -> None:
+    def update(self, x: Tensor, weight: float = 1) -> None:
         assert self.n is not None
         n = self.n.get()
         assert n is not None
@@ -144,7 +143,7 @@ def update(self, x: Tensor) -> None:
             self.rolling_mean = x.clone() if x.is_floating_point() else x.double()
         else:
             delta = x - rolling_mean
-            self.rolling_mean = rolling_mean + delta / n
+            self.rolling_mean = rolling_mean + delta * weight / n
 
 
 class MSE(Stat):
@@ -166,12 +165,12 @@ def get(self) -> Optional[Tensor]:
             return torch.zeros_like(self.prev_mean)
         return self.mse
 
-    def update(self, x: Tensor) -> None:
+    def update(self, x: Tensor, weight: float = 1) -> None:
         assert self.mean is not None
         mean = self.mean.get()
 
         if mean is not None and self.prev_mean is not None:
-            rhs = (x - self.prev_mean) * (x - mean)
+            rhs = weight * (x - self.prev_mean) * (x - mean)
             if self.mse is None:
                 self.mse = rhs
             else:
@@ -207,7 +206,7 @@ def init(self) -> None:
         self.mse_stat = self._get_stat(MSE())
         self.n_stat = self._get_stat(Count())
 
-    def update(self, x: Tensor) -> None:
+    def update(self, x: Tensor, weight: float = 1) -> None:
         pass
 
     def get(self) -> Optional[Tensor]:
@@ -251,7 +250,7 @@ def __init__(self, name: Optional[str] = None, order: int = 0) -> None:
     def init(self) -> None:
         self.var_stat = self._get_stat(Var(order=self.order))
 
-    def update(self, x: Tensor) -> None:
+    def update(self, x: Tensor, weight: float = 1) -> None:
         pass
 
     def get(self) -> Optional[Tensor]:
@@ -276,7 +275,7 @@ def __init__(
     def get(self) -> Optional[Tensor]:
         return self.result
 
-    def update(self, x: Tensor) -> None:
+    def update(self, x: Tensor, weight: float = 1) -> None:
         if self.result is None:
             self.result = x
         else:
@@ -309,6 +308,9 @@ def __init__(
     ) -> None:
         super().__init__(name=name, fn=add_fn)
 
+    def update(self, x: Tensor, weight: float = 1) -> None:
+        super().update(x * weight)
+
 
 def CommonStats() -> List[Stat]:
     r"""
diff --git a/captum/attr/_utils/summarizer.py b/captum/attr/_utils/summarizer.py
index e4002e31cf..99e8fdaf7e 100644
--- a/captum/attr/_utils/summarizer.py
+++ b/captum/attr/_utils/summarizer.py
@@ -7,6 +7,7 @@
 
 # pyre-strict
 
+import math
 from typing import Dict, List, Optional, Tuple, Type, Union
 
 import torch
@@ -63,14 +64,25 @@ def _copy_stats(self) -> List[Stat]:
 
         return copy.deepcopy(self._stats)
 
-    def update(self, x: Union[float, Tensor, Tuple[Union[float, Tensor], ...]]) -> None:
+    def update(
+        self,
+        x: Union[float, Tensor, Tuple[Union[float, Tensor], ...]],
+        weight: float = 1,
+    ) -> None:
         r"""
         Calls `update` on each `Stat` object within the summarizer
 
         Args:
             x (Tensor or Tuple[Tensor, ...]):
                 The input(s) you wish to summarize
+            weight (float):
+                Frequency weight for this update. This is useful when ``x`` is
+                already a mean over multiple observations.
         """
+        if not math.isfinite(weight) or weight < 0:
+            raise ValueError(f"weight must be finite and nonnegative, got {weight}")
+        if weight == 0:
+            return
         if self._is_inputs_tuple is None:
             self._is_inputs_tuple = isinstance(x, tuple)
         else:
@@ -99,7 +111,7 @@ def update(self, x: Union[float, Tensor, Tuple[Union[float, Tensor], ...]]) -> N
                 )
             if not isinstance(inp, torch.Tensor):
                 inp = torch.tensor(inp, dtype=torch.float)
-            self._summarizers[i].update(inp)
+            self._summarizers[i].update(inp, weight)
 
     @property
     def summary(
@@ -213,7 +225,7 @@ def __init__(self, stats: List[Stat], summary_stats_indices: List[int]) -> None:
             stat._other_stats = self
             stat.init()
 
-    def update(self, x: Tensor) -> None:
+    def update(self, x: Tensor, weight: float = 1) -> None:
         r"""
         Updates the summary of a given tensor `x`
 
@@ -222,7 +234,10 @@ def update(self, x: Tensor) -> None:
                 The tensor to summarize
         """
         for stat in self._stats:
-            stat.update(x)
+            if weight == 1:
+                stat.update(x)
+            else:
+                stat.update(x, weight)
 
     def get(self, stat: Stat) -> Optional[Stat]:
         r"""
diff --git a/tests/attr/test_stat.py b/tests/attr/test_stat.py
index 32b5ccf0c0..16d2507280 100644
--- a/tests/attr/test_stat.py
+++ b/tests/attr/test_stat.py
@@ -6,6 +6,7 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
+import math
 import random
 from typing import Callable, Generator, List, Union
 
@@ -30,6 +31,31 @@ def get_values(
 
 # pyrefly: ignore [invalid-inheritance]
 class Test(BaseTest):
+    def test_weighted_updates_match_uneven_batches(self) -> None:
+        summarizer = Summarizer([Mean(), StdDev(order=0)])
+
+        summarizer.update(torch.tensor(1.0), weight=1)
+        summarizer.update(torch.tensor(3.0), weight=3)
+
+        summary = summarizer.summary
+        assert isinstance(summary, dict)
+        assertTensorAlmostEqual(self, summary["mean"], 2.5)
+        assertTensorAlmostEqual(self, summary["std_dev"], math.sqrt(0.75))
+
+    def test_weighted_updates_reject_invalid_weights(self) -> None:
+        for weight in (-1, float("inf"), float("nan")):
+            with self.subTest(weight=weight):
+                summarizer = Summarizer([Mean()])
+                with self.assertRaisesRegex(ValueError, "finite and nonnegative"):
+                    summarizer.update(torch.tensor(1.0), weight=weight)
+
+    def test_zero_weight_update_is_ignored(self) -> None:
+        summarizer = Summarizer([Mean()])
+
+        summarizer.update(torch.tensor(1.0), weight=0)
+
+        self.assertIsNone(summarizer.summary)
+
     def test_div0(self) -> None:
         summarizer = Summarizer([Var(), Mean()])
         summ = summarizer.summary