import torch
import math


def count_qubits_gate_matrix(gate: torch.Tensor):
    """Get the number of qubits that a gate matrix acts on.

    By convention, a gate matrix has the shape

        (*optional_batch_dims, 2^k, 2^k)

    where k is the number of qubits that the gate acts on. Note that
    k might be smaller than the number of qubits in a state that we
    are going to apply the gate to.
    """
    length = gate.size()[-1]
    num_bits = round(math.log2(length))
    if 2 ** num_bits != length:
        raise RuntimeError(f'Given gate matrix has size {gate.size()} which '
                           f'is not consistent with any number of qubits.')
    return num_bits