import torch


def apply_phase(angles: torch.Tensor, state: torch.Tensor):
    """Multiply the kth component of state by e^(-i angles_k).

    Batching for this function follows PyTorch multiplication batching.
    Here are some examples:

    If `angles` is a scalar (a Tensor with only one element) and `state`
    is a state in vector layout with arbitrary batch dimensions, then the
    angle is applied to every component of every state.

    If `angles` is a vector with size (2^n) (where n is the number of qubits)
    and `state` is a Tensor with arbitrary batch dimensions, then
    the specified phase angle for each component is applied to each state in
    the batch of states.

    When `angles` has the same batch dimensions as `state`, we apply phase
    angles to states in one-to-one correspondence.

    To apply a batch of phases to a batch of states, but to use the same
    angle for each component of a given state, we can give `angles` size
    (*batch_dims, 1) while `state` will have size (*batch_dims, 2^n).

    Args:
        angles: Size (*angle_batch_dims, Hilbert space dim)
        state: Size (*state_batch_dims, Hilbert space dim)
    """
    phase_factors = torch.exp(-1.j * angles)
    return phase_factors * state