import torch


def abs_squared(x: torch.Tensor) -> torch.Tensor:
    """Compute absolute value squared of complex or real tensor.
    Avoids complex multiply and square root for efficiency."""
    if x.is_complex():
        return torch.view_as_real(x).square().sum(dim=-1)
    else:
        return x.square()