def attribution_locality(
    attributions: tuple[Tensor, ...] | list[tuple[Tensor, ...]],
    feature_mask: tuple[Tensor, ...] | None,
    target_indices: list[int],
    bboxes: tuple[Tensor, ...],
    x_threshold: float = 0.025,
    y_threshold: float = 0.025,
    use_weighted_sum: bool = False,
    return_dict: bool = False,
) -> dict | Tensor | list[Tensor]:
    """Measures how spatially concentrated attributions are around a target bounding box. ↑ better.

    For each modality, computes the fraction of (normalised) attribution mass within an x/y
    threshold of the target feature's centre, plus a weighted-RMSE spread score. Useful for
    diagnosing whether multimodal explanations focus near the correct object.

    Args:
        attributions (tuple[Tensor, ...] or list[tuple[Tensor, ...]]): Attribution tensors, one
            per modality, each of shape ``(batch_size, *input_shape)``. Pass a list of such tuples
            for multi-target mode.
        feature_mask (tuple[Tensor, ...] or None): Feature group masks of the same shape as
            ``attributions``, used to pool raw features into semantic groups before scoring.
            If ``None``, each feature is its own group.
        target_indices (list[int]): Index of the target feature group per element in
            ``attributions`` (or per target when passing a list).
        bboxes (tuple[Tensor, ...]): Bounding boxes per modality, each of shape
            ``(batch_size, n_groups, 4)`` in ``[x1, y1, x2, y2]`` format (normalised 0–1).
        x_threshold (float): Maximum normalised x-distance from the target centre to count as
            local. Default: ``0.025``.
        y_threshold (float): Maximum normalised y-distance from the target centre to count as
            local. Default: ``0.025``.
        use_weighted_sum (bool): If ``True``, use weighted-sum pooling across feature groups
            instead of plain sum. Default: ``False``.
        return_dict (bool): If ``True``, return a dict with keys ``"x_locality"``,
            ``"y_locality"``, and ``"spread"`` instead of a stacked tensor.
            Default: ``False``.

    Returns:
        Tensor or list[Tensor] or dict: Shape ``(batch_size, n_modalities, 3)`` where the last
        dimension is ``[x_locality, y_locality, spread]``. Returns a list for multi-target input
        and a dict when ``return_dict=True``.

    Example:
        >>> import torch
        >>> attr = (torch.rand(2, 16),)          # one modality, batch of 2
        >>> boxes = (torch.rand(2, 16, 4),)      # bbox per feature group
        >>> score = attribution_locality(attr, feature_mask=None, target_indices=[3, 5], bboxes=boxes)
    """
    with torch.no_grad():
        is_list = isinstance(attributions, list)
        assert len(attributions) == len(target_indices), (
            "Length of attributions list must match length of target_indices"
        )

        scores = []
        for target_index, attribution in zip(target_indices, attributions, strict=True):
            if not isinstance(attribution, tuple):
                attribution = (attribution,)
            if not isinstance(feature_mask, tuple) and feature_mask is not None:
                feature_mask = (feature_mask,)
            if not isinstance(bboxes, tuple):
                bboxes = (bboxes,)

            bsz = attribution[0].size(0)
            batch_scores = []

            for i in range(bsz):
                score = _locality_single_sample(
                    attributions_single_sample=tuple(
                        attr[i].unsqueeze(0) for attr in attribution
                    ),
                    feature_mask_single_sample=(
                        tuple(mask[i].unsqueeze(0) for mask in feature_mask)
                        if feature_mask is not None
                        else None
                    ),
                    bboxes_single_sample=tuple(bbox[i].unsqueeze(0) for bbox in bboxes),
                    target_index=target_index,
                    x_threshold=x_threshold,
                    y_threshold=y_threshold,
                    use_weighted_sum=use_weighted_sum,
                )
                batch_scores.append(score)

            scores.append(torch.stack(batch_scores))  # [batch_size, n_modalities, 3]

        if not is_list:
            scores = scores[0]

        if return_dict:
            if isinstance(scores, list):
                return {
                    "x_locality": [s[..., 0] for s in scores],
                    "y_locality": [s[..., 1] for s in scores],
                    "spread": [s[..., 2] for s in scores],
                }
            return {
                "x_locality": scores[..., 0],  # [batch_size, n_modalities]
                "y_locality": scores[..., 1],
                "spread": scores[..., 2],
            }

        return scores