import torch


def count_qubits_tensor(
        state_tensor: torch.Tensor,
        num_batch_dims: int,
):
    """Get the number of qubits of a state in tensor layout."""
    return state_tensor.dim() - num_batch_dims