Package: qujax
qujaxInstallationQuick startPure state simulationMixed state simulationConverting from TKETExamplesContributingCiting qujaxAPI Referencequjax is aJAX-based Python library for the classical simulation of quantum circuits. It is designed to besimple,fastandflexible.It follows a functional programming design by translating circuits into pure functions. This allows qujax toseamlessly interface with JAX, enabling direct access to its powerful automatic differentiation tools, just-in-time compiler, vectorization capabilities, GPU/TPU integration and growing ecosystem of packages.qujax can be used both for pure and for mixed quantum state simulation. It not only supports the standard gate set, but also allows user-defined custom operations, including general quantum channels, enabling the user to e.g. model device noise and errors.An overview of the core functionalities of qujax can be found in theQuick startsection. More advanced use-cases, including the training of parameterised quantum circuits, are listed inExamples.Installationqujax ishosted on PyPIand can be installed via the pip package managerpip install qujaxQuick startImportant note: qujax circuit parameters are expressed in units of $\pi$ (e.g. in the range $[0,2]$ as opposed to $[0, 2\pi]$).Pure state simulationWe start by defining the quantum gates making up the circuit, along with the qubits that they act on and the indices of the parameters for each gate.A list of all gates can be foundhere(custom operations can be included bypassing an array or functioninstead of a string).fromjaximportnumpyasjnpimportqujax# List of quantum gatescircuit_gates=['H','Ry','CZ']# Indices of qubits the gates will be applied tocircuit_qubit_inds=[[0],[0],[0,1]]# Indices of parameters each parameterised gate will usecircuit_params_inds=[[],[0],[]]qujax.print_circuit(circuit_gates,circuit_qubit_inds,circuit_params_inds);# q0: -----H-----Ry[0]-----◯---#                          |# q1: ---------------------CZ--We then translate the circuit to a pure functionparam_to_stthat takes a set of parameters and an (optional) initial quantum state as its input.param_to_st=qujax.get_params_to_statetensor_func(circuit_gates,circuit_qubit_inds,circuit_params_inds)param_to_st(jnp.array([0.1]))# Array([[0.58778524+0.j, 0.        +0.j],#        [0.80901706+0.j, 0.        +0.j]], dtype=complex64)The optional initial state can be passed toparam_to_stusing thestatetensor_inargument. When it is not provided, the initial state defaults to $\ket{0...0}$.Note that qujax represents quantum states asstatetensors. For example, for $N=4$ qubits, the corresponding vector space has $2^4$ dimensions, and a quantum state in this space is represented by an array with shape(2,2,2,2). The usual statevector representation with shape(16,)can be obtained by calling.flatten()or.reshape(-1)or.reshape(2**N)on this array.In the statetensor representation, the coefficient associated with e.g. basis state $\ket{0101}$ is given byarr[0,1,0,1]; each axis corresponds to one qubit.param_to_st(jnp.array([0.1])).flatten()# Array([0.58778524+0.j, 0.+0.j, 0.80901706+0.j, 0.+0.j], dtype=complex64)Finally, by defining an observable, we can map the statetensor to an expectation value. A general observable is specified using lists of Pauli matrices, the qubits they act on, and the associated coefficients.For example, $Z_1Z_2Z_3Z_4 - 2 X_3$ would be written as[['Z','Z','Z','Z'], ['X']], [[1,2,3,4], [3]], [1., -2.].st_to_expectation=qujax.get_statetensor_to_expectation_func([['Z']],[[0]],[1.])Combiningparam_to_standst_to_expectationgives us a parameter to expectation function that can be automatically differentiated using JAX.fromjaximportvalue_and_gradparam_to_expectation=lambdaparam:st_to_expectation(param_to_st(param))expectation_and_grad=value_and_grad(param_to_expectation)expectation_and_grad(jnp.array([0.1]))# (Array(-0.3090171, dtype=float32),#  Array([-2.987832], dtype=float32))Mixed state simulationMixed state simulations are analogous to the above, but with calls toget_params_to_densitytensor_funcandget_densitytensor_to_expectation_funcinstead.param_to_dt=qujax.get_params_to_densitytensor_func(circuit_gates,circuit_qubit_inds,circuit_params_inds)dt=param_to_dt(jnp.array([0.1]))dt.shape# (2, 2, 2, 2)dt_to_expectation=qujax.get_densitytensor_to_expectation_func([['Z']],[[0]],[1.])dt_to_expectation(dt)# Array(-0.3090171, dtype=float32)Similarly to a statetensor, which represents the reshaped $2^N$-dimensional statevector of a pure quantum state, adensitytensorrepresents the reshaped $2^N \times 2^N$ density matrix of a mixed quantum state. This densitytensor has shape(2,) * 2 * N.For example, for $N=2$, and a mixed state $\frac{1}{2} (\ket{00}\bra{11} + \ket{11}\bra{00} + \ket{11}\bra{11} + \ket{00}\bra{00})$, the corresponding densitytensordtis such thatdt[0,0,1,1] = dt[1,1,0,0] = dt[1,1,1,1] = dt[0,0,0,0] = 1/2, and all other entries are zero.The equivalent density matrix can be obtained by calling.reshape(2 ** N, 2 ** N).Converting from TKETOne can directly convert apytketcircuit using thetk_to_qujaxandtk_to_qujax_symbolicfunctions in thepytket-qujaxextension.An example of this can be found in thepytket-qujax_heisenberg_vqe.ipynbnotebook.ExamplesBelow are some use-case notebooks. These both illustrate the flexibility of qujax and the power of directly interfacing with JAX and its package ecosystem.heisenberg_vqe.ipynb- an implementation of the variational quantum eigensolver to find the ground state of a quantum Hamiltonian.maxcut_vqe.ipynb- an implementation of the variational quantum eigensolver to solve a MaxCut problem. Trains with Adam viaoptaxand uses more realistic stochastic parameter shift gradients.noise_channel.ipynb- uses the densitytensor simulator to fit the parameters of a depolarising noise channel.qaoa.ipynb- uses a problem-inspired QAOA ansatz to find the ground state of a quantum Hamiltonian. Demonstrates how to encode more sophisticated parameters that control multiple gates.barren_plateaus.ipynb- illustrates how to sample gradients of a cost function to identify the presence of barren plateaus. Uses batched/vectorized evaluation to speed up computation.reducing_jit_compilation_time.ipynb- explains how JAX compilation works and how that can lead to excessive compilation times when executing quantum circuits. Presents a solution for the case of circuits with a repeating structure.variational_inference.ipynb- uses a parameterised quantum circuit as a variational distribution to fit to a target probability mass function. Uses Adam viaoptaxto minimise the KL divergence between circuit and target distributions.classification.ipynb- train a quantum circuit for binary classification using data re-uploading.generative_modelling.ipynb- uses a parameterised quantum circuit as a generative model for a real life dataset. Trains via stochastic gradient Langevin dynamics on the maximum mean discrepancy between statetensor and dataset.Thepytketrepository also containstk_to_qujaximplementations for some of the above atpytket-qujax_classification.ipynb,pytket-qujax_heisenberg_vqe.ipynbandpytket-qujax_qaoa.ipynb.ContributingYou can open a bug report or a feature request by creating a newissue on GitHub.Pull requests are welcome! To open a new one, please go through the following steps:First fork the repo and create your branch fromdevelop.Commit your code and tests.Update the documentation, if required.Check the code lints (runblack . --checkandpylint */).Issue a pull request into thedevelopbranch.New commits ondevelopwill be merged intomainin the next release.Citing qujaxIf you have used qujax in your code or research, we kindly ask that you cite it. You can use the following BibTeX entry for this:@article{qujax2023,author={Duffield, Samuel and Matos, Gabriel and Johannsen, Melf},doi={10.21105/joss.05504},journal={Journal of Open Source Software},month=sep,number={89},pages={5504},title={{qujax: Simulating quantum circuits with JAX}},url={https://joss.theoj.org/papers/10.21105/joss.05504},volume={8},year={2023}}