import torch


def abs_squared(state: torch.Tensor):
    """Compute the vector of measurement probabilities for state.

    For an $N$-dimensional Hilbert space, in a given basis a state
    can be regarded as a vector (x_1, ..., x_N). This function
    simply computes (|x_1|^2, ..., |x_N|^2).

    When state is a batch, the vector of absolute value squared coefficients
    is computed for each batch entry.

    Note:
        This function trivially wraps torch operations. We recommend its usage
        when it improves readability.

    Args:
        state (Tensor): State or batch of states in vector layout. Tensor size
            should be (*batch_dims, N). State can be real or complex.

    Returns:
        Tensor with size (*batch_dims, N) giving all measurement
        probabilities for all states in the batch.
    """
    return (state.conj() * state).real