Catalyst v0.3.0
New features
-
Catalyst now officially supports macOS ARM devices, such as Apple M1/M2 machines, with macOS binary wheels available on PyPI. For more details on the changes involved to support macOS, please see the improvements section. (#229) (#232) (#233) (#234)
-
Write Catalyst-compatible programs with native Python conditional statements. (#235)
AutoGraph is a new, experimental, feature that automatically converts Python conditional statements like
if
,else
, andelif
, into their equivalent functional forms provided by Catalyst (such ascatalyst.cond
).This feature is currently opt-in, and requires setting the
autograph=True
flag in theqjit
decorator:dev = qml.device("lightning.qubit", wires=1) @qjit(autograph=True) @qml.qnode(dev) def f(x): if x < 0.5: qml.RY(jnp.sin(x), wires=0) else: qml.RX(jnp.cos(x), wires=0) return qml.expval(qml.PauliZ(0))
The implementation is based on the AutoGraph module from TensorFlow, and requires a working TensorFlow installation be available. In addition, Python loops (
for
andwhile
) are not yet supported, and do not work in AutoGraph mode.Note that there are some caveats when using this feature especially around the ues of global variables or object mutation inside of methods. A functional style is always recommended when using
qjit
or AutoGraph. -
The quantum adjoint operation can now be used in conjunction with Catalyst control flow, such as loops and conditionals. For this purpose a new instruction,
catalyst.adjoint
, has been added. (#220)catalyst.adjoint
can wrap around quantum functions which contain the Catalystcond
,for_loop
, andwhile_loop
primitives. Previously, the usage ofqml.adjoint
on functions with these primitives would result in decomposition errors. Note that a future release of Catalyst will
merge the behaviour ofcatalyst.adjoint
intoqml.adjoint
for convenience.dev = qml.device("lightning.qubit", wires=3) @qjit @qml.qnode(dev) def circuit(x): @for_loop(0, 3, 1) def repeat_rx(i): qml.RX(x / 2, wires=i) adjoint(repeat_rx)() return qml.expval(qml.PauliZ(0))
>>> circuit(0.2) array(0.99500417)
Additionally, the ability to natively represent the adjoint construct in Catalyst's program representation (IR) was added.
-
QJIT-compiled programs now support (nested) container types as inputs and outputs of compiled functions. This includes lists and dictionaries, as well as any data structure implementing the PyTree protocol. (#215) (#221)
For example, a program that accepts and returns a mix of dictionaries, lists, and tuples:
@qjit def workflow(params1, params2): res1 = params1["a"][0][0] + params2[1] return {"y1": jnp.sin(res1), "y2": jnp.cos(res1)}
>>> params1 = {"a": [[0.1], 0.2]} >>> params2 = (0.6, 0.8) >>> workflow(params1, params2) array(0.78332691)
-
Compile-time backpropagation of arbitrary hybrid programs is now supported, via integration with Enzyme AD. (#158) (#193) (#224) (#225) (#239) (#244)
This allows
catalyst.grad
to differentiate hybrid functions that contain both classical pre-processing (inside & outside of QNodes), QNodes, as well as classical post-processing (outside of QNodes) via a combination of backpropagation and quantum gradient methods.The new default for the differentiation
method
attribute incatalyst.grad
has been changed to"auto"
, which performs Enzyme-based reverse mode AD on classical code, in conjunction with the quantumdiff_method
specified on each QNode:dev = qml.device("lightning.qubit", wires=1) @qml.qnode(dev, diff_method="parameter-shift") def circuit(theta): qml.RX(jnp.exp(theta ** 2) / jnp.cos(theta / 4), wires=0) return qml.expval(qml.PauliZ(wires=0))
>>> grad = qjit(catalyst.grad(circuit, method="auto")) >>> grad(jnp.pi) array(0.05938718)
The reworked differentiation pipeline means you can now compute exact derivatives of programs with both classical pre- and post-processing, as shown below:
@qml.qnode(qml.device("lightning.qubit", wires=1), diff_method="adjoint") def circuit(theta): qml.RX(jnp.exp(theta ** 2) / jnp.cos(theta / 4), wires=0) return qml.expval(qml.PauliZ(wires=0)) def loss(theta): return jnp.pi / jnp.tanh(circuit(theta)) @qjit def grad_loss(theta): return catalyst.grad(loss)(theta)
>>> grad_loss(1.0) array(-1.90958669)
You can also use multiple QNodes with different differentiation methods:
@qml.qnode(qml.device("lightning.qubit", wires=1), diff_method="parameter-shift") def circuit_A(params): qml.RX(jnp.exp(params[0] ** 2) / jnp.cos(params[1] / 4), wires=0) return qml.probs() @qml.qnode(qml.device("lightning.qubit", wires=1), diff_method="adjoint") def circuit_B(params): qml.RX(jnp.exp(params[1] ** 2) / jnp.cos(params[0] / 4), wires=0) return qml.expval(qml.PauliZ(wires=0)) def loss(params): return jnp.prod(circuit_A(params)) + circuit_B(params) @qjit def grad_loss(theta): return catalyst.grad(loss)(theta)
>>> grad_loss(jnp.array([1.0, 2.0])) array([ 0.57367285, 44.4911605 ])
And you can differentiate purely classical functions as well:
def square(x: float): return x ** 2 @qjit def dsquare(x: float): return catalyst.grad(square)(x)
>>> dsquare(2.3) array(4.6)
Note that the current implementation of reverse mode AD is restricted to 1st order derivatives, but you can still use
catalyst.grad(method="fd")
is still available to perform a finite differences approximation of any differentiable function. -
Add support for the new PennyLane arithmetic operators. (#250)
PennyLane is in the process of replacing
Hamiltonian
andTensor
observables with a set of general arithmetic operators. These consist of Prod, Sum and SProd.By default, using dunder methods (eg.
+
,-
,@
,*
) to combine operators with scalars or other operators will createHamiltonian
andTensor
objects. However, these two methods will be deprecated in coming releases of PennyLane.To enable the new arithmetic operators, one can use
Prod
,Sum
, andSprod
directly or activate them by calling enable_new_opmath at the beginning of your PennyLane program.dev = qml.device("lightning.qubit", wires=2) @qjit @qml.qnode(dev) def circuit(x: float, y: float): qml.RX(x, wires=0) qml.RX(y, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(0.2 * qml.PauliX(wires=0) - 0.4 * qml.PauliY(wires=1))
>>> qml.operation.enable_new_opmath() >>> qml.operation.active_new_opmath() True >>> circuit(np.pi / 4, np.pi / 2) array(0.28284271)
Improvements
-
Better support for Hamiltonian observables:
-
Allow Hamiltonian observables with integer coefficients. (#248)
For example, compiling the following circuit wasn't previously allowed, but is now supported in Catalyst:
dev = qml.device("lightning.qubit", wires=2) @qjit @qml.qnode(dev) def circuit(x: float, y: float): qml.RX(x, wires=0) qml.RY(y, wires=1) coeffs = [1, 2] obs = [qml.PauliZ(0), qml.PauliZ(1)] return qml.expval(qml.Hamiltonian(coeffs, obs))
-
Allow nested Hamiltonian observables. (#255)
@qjit @qml.qnode(qml.device("lightning.qubit", wires=3)) def circuit(x, y, coeffs1, coeffs2): qml.RX(x, wires=0) qml.RX(y, wires=1) qml.RY(x + y, wires=2) obs = [ qml.PauliX(0) @ qml.PauliZ(1), qml.Hamiltonian(coeffs1, [qml.PauliZ(0) @ qml.Hadamard(2)]), ] return qml.var(qml.Hamiltonian(coeffs2, obs))
-
-
Various performance improvements:
-
The execution and compile time of programs has been reduced, by generating more efficient code and avoiding unnecessary optimizations. Specifically, a scalarization procedure was added to the MLIR pass pipeline, and LLVM IR compilation is now invoked with optimization level 0. (#217)
-
The execution time of compiled functions has been improved in the frontend. (#213)
Specifically, the following changes have been made, which leads to a small but measurable improvement when using larger matrices as inputs, or functions with many inputs:
- only loading the user program library once per compilation,
- generating return value types only once per compilation,
- avoiding unnecessary type promotion, and
- avoiding unnecessary array copies.
-
Peak memory utilization of a JIT compiled program has been reduced, by allowing tensors to be scheduled for deallocation. Previously, the tensors were not deallocated until the end of the call to the JIT compiled function. (#201)
-
-
Various improvements have been made to enable Catalyst to compile on macOS:
-
Remove unnecessary
reinterpret_cast
fromObsManager
. Removal of thesereinterpret_cast
allows compilation of the runtime to succeed in macOS. macOS uses an ILP32 mode for Aarch64 where they use the full 64 bit mode but with 32 bit Integer, Long, and Pointers. This patch also changes a test file to prevent a mismatch in machines which compile using ILP32 mode. (#229) -
Allow runtime to be compiled on macOS. Substitute
nproc
with a call toos.cpu_count()
and use correct flags forld.64
. (#232) -
Improve portability on the frontend to be available on macOS. Use
.dylib
, remove unnecessary flags, and address behaviour difference in flags. (#233) -
Small compatibility changes in order for all integration tests to succeed on macOS. (#234)
-
-
Dialects can compile with older versions of clang by avoiding type mismatches. (#228)
-
The runtime is now built against
qir-stdlib
pre-build artifacts. (#236) -
Small improvements have been made to the CI/CD, including fixing the Enzyme cache, generalize caches to other operating systems, fix build wheel recipe, and remove references to QIR in runtime's Makefile. (#243) (#247)
Breaking changes
-
Support for Python 3.8 has been removed. (#231)
-
The default differentiation method on
grad
andjacobian
is reverse-mode automatic differentiation instead of finite differences. When a QNode does not have adiff_method
specified, it will default to using the parameter shift method instead of finite-differences. (#244) (#271) -
The JAX version used by Catalyst has been updated to
v0.4.14
, the minimum PennyLane version required is nowv0.32
. (#264) -
Due to the change allowing Python container objects as inputs to QJIT-compiled functions, Python lists are no longer automatically converted to JAX arrays. (#231)
This means that indexing on lists when the index is not static will cause a
TracerIntegerConversionError
, consistent with JAX's behaviour.That is, the following example is no longer support:
@qjit def f(x: list, index: int): return x[index]
However, if the parameter
x
above is a JAX or NumPy array, the compilation will continue to succeed. -
The
catalyst.grad
function has been renamed tocatalyst.jacobian
and supports differentiation of functions that return multiple or non-scalar outputs. A newcatalyst.grad
function has been added that enforces that it is differentiating a function with a single scalar return value. (#254)
Bug fixes
-
Fixed an issue preventing the differentiation of
qml.probs
with the parameter-shift method. (#211) -
Fixed the incorrect return value data-type with functions returning
qml.counts
. (#221) -
Fix segmentation fault when differentiating a function where a quantum measurement is used multiple times by the same operation. (#242)
Contributors
This release contains contributions from (in alphabetical order):
Ali Asadi, David Ittah, Erick Ochoa Lopez, Jacob Mai Peng, Romain Moyard, Sergei Mironov.