import torch


def complex_abs_sq(data: torch.Tensor) -> torch.Tensor:
    """
    Compute the squared absolute value of a complex tensor.

    Args:
        data: A complex valued tensor, where the size of the final dimension
            should be 2.

    Returns:
        Squared absolute value of data.
    """
    if data.shape[-1] != 2:
        raise ValueError("Tensor does not have separate complex dim.")

    return (data**2).sum(dim=-1)