from typing import Optional
from typing import Sequence
import torch


def uniform_superposition(
        num_qubits: int,
        batch_dims: Optional[Sequence] = None,
        device: torch.device = torch.device("cpu"),
        dtype: torch.dtype = torch.complex64,
        requires_grad: bool = False
):
    """Create the uniform superposition state |+...+>.

    This function is defined for the complex case only. Optional batch
    dimensions can be provided in which case the same uniform superposition
    is copied to form a batch with specified shape.
    """
    scaling = 2 ** (-num_qubits/2.)
    state = scaling * torch.ones(
        2 ** num_qubits,
        device=device,
        dtype=dtype
    )

    if batch_dims is not None:
        batch_dims = torch.Size(batch_dims)
        state = state.repeat(batch_dims + (1,))

    state.requires_grad_(requires_grad)
    return state