Using QJIT with UCCSD ansatz

PennyLane Help
1

I have two questions regarding @qml.qjit.

  1. I am using @qml.qjit with the UCCSD ansatz, and it’s giving the following error.

    ValueError: Operation Sum is not written in terms of a single parameter

Here is my code:

dev = qml.device("lightning.qubit", wires=qubits)

s_wires, d_wires = qchem.excitations_to_wires(singles, doubles)

@qml.qnode(dev)
def circuit(params):
  qml.UCCSD(params,range(qubits),s_wires,d_wires,hf_state)
  return qml.expval(H)

opt = qml.QNGOptimizerQJIT(stepsize=0.1, approx='block-diag')

def update_step_qjit(i, args):
    params, state = args
    return opt.step(circuit, params, state)

@qml.qjit
def optimization_qjit(params, iters):
    state = opt.init(params)
    args = (params, state)
    params, state = qml.for_loop(iters)(update_step_qjit)(args)
    return params

params = jnp.zeros(num_params)

iters = 5
start = time.time()
optimization_qjit(params=params,iters=iters)
print(time.time() - start)

Any idea how to resolve this? I am running the code on an A100 GPU on Google Colab.

  1. When I run the code without @qml.qjit. The code is running ok. However, it’s taking more time than regular Pennylane numpy and QNGOptimizer. After reading this blog, PennyLane v0.43 and Catalyst v0.13 released | PennyLane Blog, I thought this would speed up (or atleast perform the same) the simulation. Could you suggest an alternative method to accelerate the VQE?

For reference: Pennylane Numpy + QNGOptimizer: 35 secs per iteration on CPU.

Here is qml.about()

!pip install pennylane==0.43
!pip install pennylane-catalyst==0.13
!pip install -U "jax[cuda12]==0.6.2"


import pennylane as qml
qml.about()

Name: pennylane
Version: 0.43.0
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: 
Author: 
Author-email: 
License: 
Location: /usr/local/lib/python3.12/dist-packages
Requires: appdirs, autograd, autoray, cachetools, diastatic-malt, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, tomlkit, typing_extensions
Required-by: pennylane_catalyst, pennylane_lightning

Platform info:           Linux-6.6.105+-x86_64-with-glibc2.35
Python version:          3.12.12
Numpy version:           2.0.2
Scipy version:           1.16.3
JAX version:             0.6.2
Installed devices:
- default.clifford (pennylane-0.43.0)
- default.gaussian (pennylane-0.43.0)
- default.mixed (pennylane-0.43.0)
- default.qubit (pennylane-0.43.0)
- default.qutrit (pennylane-0.43.0)
- default.qutrit.mixed (pennylane-0.43.0)
- default.tensor (pennylane-0.43.0)
- null.qubit (pennylane-0.43.0)
- reference.qubit (pennylane-0.43.0)
- nvidia.custatevec (pennylane_catalyst-0.13.0)
- nvidia.cutensornet (pennylane_catalyst-0.13.0)
- oqc.cloud (pennylane_catalyst-0.13.0)
- softwareq.qpp (pennylane_catalyst-0.13.0)
- lightning.qubit (pennylane_lightning-0.43.0)
2

Hi @mdaamir ,

I see a few things.

Regarding point 1 it looks like @qml.qjit is having issues with the Hamiltonian you’re trying to return. It could potentially be related to the following note in the documentation for QNGOptimizerQJIT:

QNGOptimizerQJIT does not work correctly if there is any classical processing in the QNode circuit (e.g., 2 * theta as a gate parameter).

A Hamiltonian is defined as a weighed sum of Pauli operators, e.g. 1.0*qml.Z(0) + 1.0*qml.X(1). The weights being applied to the Pauli operators might be causing the error, which means there wouldn’t be a workaround since this is inherent to the Hamiltonian itself.

For reference, I took the example code below from the v0.42 Improvements section of the release notes.

import jax.numpy as jnp

@qml.qjit(autograph=True)
def workflow():
    dev = qml.device("lightning.qubit", wires=2)

    @qml.qnode(dev)
    def circuit(params):
        qml.RX(params[0], wires=0)
        qml.RY(params[1], wires=1)
        return qml.expval(qml.Z(0) + qml.X(1))

    opt = qml.QNGOptimizerQJIT(stepsize=0.2)

    params = jnp.array([0.1, 0.2])
    state = opt.init(params)
    for _ in range(100):
        params, state = opt.step(circuit, params, state)

    return params

workflow()

The example code above works with no issues.
However, if you change the qnode return to qml.expval(1.0*qml.Z(0) + 1.0*qml.X(1)) then you get the error you identified before: ValueError: Operation Sum is not written in terms of a single parameter.

Regarding your reasoning for wanting to use qjit, note that it can make things slower for small programs (same as using GPUs). If you have less than 10 qubits then using qjit might make things slower, and if you have less than 20 qubits then using GPUs might make things slower too. Note that in your case you weren’t even fully leveraging the GPU since you were using "lightning.qubit" instead of "lightning.gpu".

You can learn more about when to use which simulator in our performance page.

Since you haven’t shared any information on your Hamiltonian it’s hard to tell what the best combination of tools is best. Assuming you have less than 15 qubits, my recommendation would be to use "lightning.qubit" without Catalyst, but feel free to explore different alternatives and choose the one that works best for your problem.

I hope this helps!

3

Hi @CatalinaAlbornoz,

Thanks a lot for your detailed answer. Your response to my Q1 makes a lot of sense now.

Regarding Q2, I am running simulations for LiH (12Q), H2O (14Q), and NH3 (16Q). I tried using lightning.gpu, but that’s not faster compared to lightning.qubit on CPU.

4

Hi @mdaamir ,

I’m glad my previous answer helped. Since you’re using less than 20 qubits it’s normal for lightning.qubit on CPU to be faster.

Let us know if you have any further questions!