def extract_bias_diagonal(module, S, sum_batch=True):
    """Extract diagonal of ``(Jᵀ S) (Jᵀ S)ᵀ`` where ``J`` is the bias Jacobian.

    Args:
        module (torch.nn.Conv1d or torch.nn.Conv2d or torch.nn.Conv3d): Convolution
            layer for which the diagonal is extracted w.r.t. the bias.
        S (torch.Tensor): Backpropagated (symmetric factorization) of the loss Hessian.
            Has shape ``(V, *module.output.shape)``.
        sum_batch (bool, optional): Sum out the batch dimension of the bias diagonals.
            Default value: ``True``.

    Returns:
        torch.Tensor: Per-sample bias diagonal if ``sum_batch=False`` (shape
            ``(N, module.bias.shape)`` with batch size ``N``) or summed bias
            diagonal if ``sum_batch=True`` (shape ``module.bias.shape``).
    """
    start_spatial = 3
    sum_before = list(range(start_spatial, S.dim()))
    sum_after = [0, 1] if sum_batch else [0]

    return S.sum(sum_before).pow_(2).sum(sum_after)