Simulation of quantum physics with Tensor Processing Units: brute-force computation of ground states and time evolution

Tensor Processing Units (TPUs) were developed by Google exclusively to support large-scale machine learning tasks. TPUs can, however, also be used to accelerate and scale up other computationally demanding tasks. In this paper we repurpose TPUs for the challenging problem of simulating quantum spin systems. Consider a lattice model made of N spin-frac{1}{2} quantum spins, or qubits, with a Hamiltonian H = sum_i h_i that is a sum of local terms h_i and a wavefunction |Ψrangle consisting of 2^N complex amplitudes. We demonstrate the usage of TPUs for both (i) computing the ground state |Ψ_gsrangle of the Hamiltonian H, and (ii) simulating the time evolution |Ψ(t)rangle=e^-itH|Ψ(0)rangle generated by this Hamiltonian starting from some initial state |Ψ(0)rangle. The bottleneck of the above tasks is computing the product H |Ψrangle, which can be implemented with remarkable efficiency utilising the native capabilities of TPUs. With a TPU v3 pod, with 2048 cores, we simulate wavefunctions |Ψrangle of up to N=38 qubits. The dedicated matrix multiplication units (MXUs), the high bandwidth memory (HBM) on each core, and the fast inter-core interconnects (ICIs) together provide performance far beyond the capabilities of general purpose processors.