import torch


def act_last_qubit_tensor(
        single_qubit_operator: torch.Tensor,
        state_tensor: torch.Tensor,
) -> torch.Tensor:
    """Apply an operator to the last qubit of a state in tensor layout.

        Args:
            single_qubit_operator: Tensor with size (2, 2, 2) or (2, 2) in
                the complex and real cases respectively. For the complex case,
                the first dimension is for real and imaginary parts.

            state_tensor: State or batch of states in tensor layout.
        """
    def act(matrix, tensor):
        """Contract matrix with the last index of tensor."""
        return torch.einsum('ab, ...b -> ...a', matrix, tensor)

    return act(single_qubit_operator, state_tensor)