Yao.jl - Differentiable Quantum Programming In Julia

We introduce Yao (check our latest paper), an open-source Julia package for solving practical problems in quantum computation research. The name Yao comes from the first Chinese character for unitary (幺正).

Why we created Yao? To be short, we are as greedy as Julia itself. We want something that is:

Differentiable

Like many other Julia blog posts (as well as the Zygote paper) have mentioned: gradients can be a better programmer than humans sometimes. In quantum computing, they can be used for variational algorithms, quantum control, gate learning, etc. Thus, we want to have differentiable programming on quantum circuits as well!

However, automatic differentiation (AD) for quantum circuits is quite different from regular programs: the memory allocation cost in circuit simulation can be extremely high due to caching the intermediate states in the general context. And in forward mode AD, we need extra semantics to preserve the quantum circuit – so that it can be implemented on the real device.

In Yao, to accomplish this goal, we implemented a builtin domain-specific automatic differentiation (AD) engine to make use of the reversible nature of quantum circuits. The following is an example of a variational quantum eigensolver algorithm. With our efficient AD engine, you should be able to try it on your laptop

using Yao, YaoExtensions
# number of qubits and circuit depth
n, d = 16, 100
circuit = dispatch!(variational_circuit(n, d),:random)

h = heisenberg(n)

for i in 1:100
 _, grad = expect'(h, zero_state(n) => circuit)
 dispatch!(-, circuit, 1e-1 * grad)
 println("Step $i, energy = $(real.(expect(h, zero_state(n)=>circuit)))")
end

This example trains a 100 layer parametrized circuit (4816 parameters) to find the ground state of a 16 site Heisenberg model. The engine can also be integrated with general AD framework such as Zygote seamlessly, e.g., in our gate learning example. We use differentiable programming to find the decomposition of a given unitary. You can learn more in our tutorial and Quantum Algorithm Zoo.

Extensible

New research ideas keep emerging every day and every hour. The field quantum software itself grows rapidly. We want a framework that is flexible enough for researchers and developers to extend it at any level for any possible type of research.

First, we designed and developed a hardware-free intermediate representation (the Quantum Block Intermediate Representation, QBIR) to represent and manipulate quantum circuits and a set of quantum register interface. This design enables one to extend Yao to customized algorithms, hardware and more by overloading the relevant interfaces. For example, while achieving the state-of-the-art performance as shown in next section, we extend our CUDA backend in CuYao with only a few hundred lines of code written in native Julia with CUDAnative. With some patches and syntax sugar, Yao just works with the symbolic engine SymEngine that allows you to differentiate, and calculate a quantum circuit with symbolic computation.

Second, like other projects in Julia, we make Yao extensible at the architecture level. The package Yao (or its CUDA backend CuYao) is only a meta-package that re-exports other component packages. Developers can customize their own software with light-weight dependencies and develop new features rapidly.

Efficient

Efficiency also matters especially in research that evolves numerical experiments, such as variational quantum algorithms. Besides all other exciting features, we still want this framework to achieve the-state-of-art performance in the simulation.

By making use of native GPU programming in Julia and specialization based on multiple dispatch, Yao achieves state-of-the-art performance on intermediate-sized quantum circuits.

You can view a more detailed benchmark here.

What's more?

So far, we are happy to announce its birth, but the journey just starts.

We still want actual hardware compilation (e.g. to OpenQASM), circuit simplification and compilation (YaoIR), visualization, tensor network (check the online playground), and more!

Although some beta users helped us shape this software during real research work, we still need more use cases to develop it further and more people to join us. If you are interested in this idea, join us, and let's build a powerful tool for quantum computing research!