How to perform optimisation in parallel?

PennyLane Help
1

Hello, I am trying to optimise quantum circuits. Below is my minimum working example.

import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev)
def circuit(initial_state, params):
    qml.StatePrep(initial_state, wires=0, normalize=True)
    qml.Rot(params[0], params[1], params[2], wires=0)
    qml.Rot(params[3], params[4], params[5], wires=1)
    qml.CNOT(wires=[0, 1])
            
    return qml.probs(wires=[0, 1])

def cost_function(params, initial_state_vec, pprep_vec):
    cost = 0
    for i in range(len(initial_state_vec)):
       results = circuit(initial_state_vec[i], params)
       cost += pprep_vec[i]*(1 - results[i])
    return cost

opt = qml.GradientDescentOptimizer()
n_iter = 100
initial_state_vec = [[1, 0], [0, 1], [1, 1], [1, -1]]
n_state = len(initial_state_vec)
pprep_vec = np.ones(n_state)/n_state

rng = np.random.default_rng()
params = rng.random(6)*np.pi
cost = np.empty(n_iter)
for i in range(n_iter):
    params, cost[i] = opt.step_and_cost(lambda params: cost_function(params, initial_state_vec, pprep_vec), params)

As I scale up the number of qubits and parameters, I would like to do this optimisation in parallel to save time. I know Pennylane has parameter broadcasting but I do not know how to get it working with my example. If I include parameter broadcasting by changing the params and cost function as below,

params = rng.random((6,2))*np.pi

def cost_function(params, initial_state_vec, pprep_vec):
    cost = [0]*2
    for j in range(len(cost)):    
        for i in range(len(initial_state_vec)):
           results = circuit(initial_state_vec[i], params)
           cost[j] += pprep_vec[i]*(1 - results[j][i])
    return cost

I get the following error.

C:\Users\Jackson\venv-lightning-no-spyder-notebook\Lib\site-packages\autograd\tracer.py:14: UserWarning: Output seems independent of input.
  warnings.warn("Output seems independent of input.")
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File ~\venv-lightning-no-spyder-notebook\Lib\site-packages\spyder_kernels\customize\utils.py:209, in exec_encapsulate_locals(code_ast, globals, locals, exec_fun, filename)
    207     if filename is None:
    208         filename = "<stdin>"
--> 209     exec_fun(compile(code_ast, filename, "exec"), globals, None)
    210 finally:
    211     if use_locals_hack:
    212         # Cleanup code

File c:\users\jackson\documents\phd applications 2023\singapore\astar\python codes\untitled0.py:40
     38 cost = np.empty(n_iter)
     39 for i in range(n_iter):
---> 40     params, cost[i] = opt.step_and_cost(lambda params: cost_function(params, initial_state_vec, pprep_vec), params)

File ~\venv-lightning-no-spyder-notebook\Lib\site-packages\pennylane\optimize\gradient_descent.py:64, in GradientDescentOptimizer.step_and_cost(self, objective_fn, grad_fn, *args, **kwargs)
     44 def step_and_cost(self, objective_fn, *args, grad_fn=None, **kwargs):
     45     """Update trainable arguments with one step of the optimizer and return the corresponding
     46     objective function value prior to the step.
     47 
   (...)
     61         If single arg is provided, list [array] is replaced by array.
     62     """
---> 64     g, forward = self.compute_grad(objective_fn, args, kwargs, grad_fn=grad_fn)
     65     new_args = self.apply_grad(g, args)
     67     if forward is None:

File ~\venv-lightning-no-spyder-notebook\Lib\site-packages\pennylane\optimize\gradient_descent.py:122, in GradientDescentOptimizer.compute_grad(objective_fn, args, kwargs, grad_fn)
    104 r"""Compute the gradient of the objective function at the given point and return it along with
    105 the objective function forward pass (if available).
    106 
   (...)
    119     will not be evaluated and instead ``None`` will be returned.
    120 """
    121 g = get_gradient(objective_fn) if grad_fn is None else grad_fn
--> 122 grad = g(*args, **kwargs)
    123 forward = getattr(g, "forward", None)
    125 num_trainable_args = sum(getattr(arg, "requires_grad", False) for arg in args)

File ~\venv-lightning-no-spyder-notebook\Lib\site-packages\pennylane\_grad.py:224, in grad.__call__(self, *args, **kwargs)
    221     self._forward = self._fun(*args, **kwargs)
    222     return ()
--> 224 grad_value, ans = grad_fn(*args, **kwargs)  # pylint: disable=not-callable
    225 self._forward = ans
    227 return grad_value

File ~\venv-lightning-no-spyder-notebook\Lib\site-packages\autograd\wrap_util.py:20, in unary_to_nary.<locals>.nary_operator.<locals>.nary_f(*args, **kwargs)
     18 else:
     19     x = tuple(args[i] for i in argnum)
---> 20 return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)

File ~\venv-lightning-no-spyder-notebook\Lib\site-packages\pennylane\_grad.py:245, in grad._grad_with_forward(fun, x)
    242 vjp, ans = _make_vjp(fun, x)  # pylint: disable=redefined-outer-name
    244 if vspace(ans).size != 1:
--> 245     raise TypeError(
    246         "Grad only applies to real scalar-output functions. "
    247         "Try jacobian, elementwise_grad or holomorphic_grad."
    248     )
    250 grad_value = vjp(vspace(ans).ones())
    251 return grad_value, ans

TypeError: Grad only applies to real scalar-output functions. Try jacobian, elementwise_grad or holomorphic_grad.

Here is the output of qml.about().

Name: PennyLane
Version: 0.39.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: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: C:\Users\Jackson\venv-lightning-no-spyder-notebook\Lib\site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           Windows-10-10.0.19045-SP0
Python version:          3.12.7
Numpy version:           2.0.2
Scipy version:           1.14.1
Installed devices:
- default.clifford (PennyLane-0.39.0)
- default.gaussian (PennyLane-0.39.0)
- default.mixed (PennyLane-0.39.0)
- default.qubit (PennyLane-0.39.0)
- default.qutrit (PennyLane-0.39.0)
- default.qutrit.mixed (PennyLane-0.39.0)
- default.tensor (PennyLane-0.39.0)
- null.qubit (PennyLane-0.39.0)
- reference.qubit (PennyLane-0.39.0)
- lightning.qubit (PennyLane_Lightning-0.39.0)

Any help with this is appreciated.

2

Hi @JacksonPhoong ,

Welcome to the Forum!

You can indeed do parameter broadcasting but you can’t find gradients when the cost function isn’t a scalar. In your code it seems that the cost was a vector of length two, and the GradientDescentOptimizer cannot find a gradient over anything that’s not a scalar. In your case it looks like you want a cost function that sums the cost over all batches. Below is a code that does this. I’ve added comments to show what I’ve changed with respect to your code.

import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev)
def circuit(initial_state, params):
    qml.StatePrep(initial_state, wires=0, normalize=True)
    qml.Rot(params[0], params[1], params[2], wires=0)
    qml.Rot(params[3], params[4], params[5], wires=1)
    qml.CNOT(wires=[0, 1])
            
    return qml.probs(wires=[0, 1])

# Keep the original cost function but slightly modified
def cost_function(params, initial_state_vec, pprep_vec):
    cost = [0]*n_batches
    for i in range(len(initial_state_vec)):
        results = circuit(initial_state_vec[i], params)
        # note: The leading dimension in results is the batch dimension
        
        # Numpy doesn't like in-place replacements.
        # Check this part of the PennyLane docs to learn more -> https://docs.pennylane.ai/en/stable/introduction/interfaces/numpy.html#advanced-autograd-usage
        cost = cost + pprep_vec[i]*np.array(1 - results[:,i])
    
    # The cost will be the sum over all batches. This way we get a single scalar that we can use for optimization.
    return np.sum(cost)

opt = qml.GradientDescentOptimizer()
n_iter = 10 # Tested with a smaller number of iterations
initial_state_vec = np.array([[1, 0], [0, 1], [1, 1], [1, -1]],requires_grad=False) # Made this non-trainable
n_state = len(initial_state_vec)
pprep_vec = np.ones(n_state,requires_grad=False)/n_state # Made this non-trainable

rng = np.random.default_rng()
n_batches = 2 # Specified this separately for ease
params = rng.random((6,n_batches))*np.pi # n_batches must be at least 1
params.requires_grad=True # Make the parameters trainable

cost_array = np.empty(n_iter) # modified for clarity
for i in range(n_iter):
    args, prev_cost = opt.step_and_cost(cost_function, params, initial_state_vec, pprep_vec) # modified
    cost_array[i] = prev_cost
    params = args[0] # args will contain all arguments of the cost function
    
print('cost_array: ',cost_array)

I hope this helps!

3

Hey @CatalinaAlbornoz, thanks for your clear and concise reply. I have learned some things from it. Sorry I wasn’t too clear in my previous comment. I don’t actually want a cost function that sums the cost over all batches. I want it to just sum the weighted probabilities from the measurements in each individual batch. Basically, I want to keep the batches separate and see how the different input parameters affect the convergence of the optimiser.

I know that GradientDescentOptimizer should only work for a scalar. The issue was that parameter broadcasting returns an array for the batches. I was hoping that GradientDescentOptimizer would work by accepting the array like how parameter broadcasting works. Am I correct in saying that parameter broadcasting wouldn’t really work in my case? If that is so, could you suggest another way to perform the optimisation in parallel?

Hope that clears it up.

4

Hi @JacksonPhoong ,

The issue here is that the GradientDescentOptimizer uses qml.grad under the hood, which isn’t designed for higher-order derivatives. You could potentially use qml.jacobian for this but in that case you wouldn’t be using the GradientDescentOptimizer.

If your goal is simply to speed things up you can instead try batching over the non-trainable inputs. Below is an example of how you could do this.

import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev)
def circuit(initial_state, params):
    qml.StatePrep(initial_state, wires=0, normalize=True)
    qml.Rot(params[0], params[1], params[2], wires=0)
    qml.Rot(params[3], params[4], params[5], wires=1)
    qml.CNOT(wires=[0, 1])
            
    return qml.probs(wires=[0, 1])

# new cost function batching over the non-trainable inputs
def cost_function(params,initial_state_vec,pprep_vec):
    r = circuit(initial_state_vec, params)
    c = pprep_vec*np.array(1 - r)
    cost = np.trace(c)
    return cost

opt = qml.GradientDescentOptimizer()
n_iter = 10 
initial_state_vec = np.array([[1, 0], [0, 1], [1, 1], [1, -1]],requires_grad=False) 
n_state = len(initial_state_vec)
pprep_vec = np.ones(n_state,requires_grad=False)/n_state 

rng = np.random.default_rng()
params = rng.random(6)*np.pi # single batch of parameters
params.requires_grad=True 

cost_array = np.empty(n_iter) 
for i in range(n_iter):
    args, prev_cost = opt.step_and_cost(cost_function, params, initial_state_vec, pprep_vec) 
    cost_array[i] = prev_cost
    params = args[0] 
    
print('cost_array: ',cost_array)

Let me know if this helps!

5

Hey @CatalinaAlbornoz, what do you mean by higher-order derivatives? I thought that the gradient and the jacobian are first-order derivatives. Could you suggest which other optimisers that uses qml.jacobian instead? The choice of optimiser is not fixed for me. Thanks for the example, that definitely helped to speed things up.

6

Hi @JacksonPhoong ,

You’re right, that was bad wording on my part. The Jacobian is still first order.

You can build your own gradient descent using the Jacobian. The Jacobian will provide your gradients which you can then multiply by a learning rate. Note that the Jacobian will have a lot of zeros everywhere so I added a for loop to extract the actual value of the updated parameters. See the code below.

Note that:

  1. This may not be the most optimized way of doing this (but it works).
  2. This may be slower than doing everything separately.

I do hope that this helps you though!

import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev)
def circuit(initial_state, params):
    qml.StatePrep(initial_state, wires=0, normalize=True)
    qml.Rot(params[0], params[1], params[2], wires=0)
    qml.Rot(params[3], params[4], params[5], wires=1)
    qml.CNOT(wires=[0, 1])
            
    return qml.probs(wires=[0, 1])

def cost_function(params_vec,initial_state_vec,pprep_vec):
    cost = [0]*(params_vec.shape[0]) # The cost length will be the number of batches
    # Find the cost for each batch
    for i,params in enumerate(params_vec):
        r = circuit(initial_state_vec, params)
        c = pprep_vec*np.array(1 - r)
        cost[i] = np.trace(c)
    return cost

initial_state_vec = np.array([[1, 0], [0, 1], [1, 1], [1, -1]],requires_grad=False)
n_state = len(initial_state_vec)
pprep_vec = np.ones(n_state,requires_grad=False)/n_state

rng = np.random.default_rng()
n_batches = 3
params_vec = rng.random((n_batches,6))*np.pi # The batched dimension must go first in this case
params_vec.requires_grad=True


# Optimization

learning_rate = 0.1
num_steps = 10
# The Jacobian will be your gradient function
grad_fn = qml.jacobian(lambda p: np.stack(cost_function(p, initial_state_vec, pprep_vec)))

cost_vec = []
for step in range(num_steps):
    cost = cost_function(params_vec, initial_state_vec, pprep_vec)
    # update the parameters based on a basic gradient descent
    # temp_params will have a lot of zeros
    temp_params = params_vec - learning_rate * grad_fn(params_vec)
    # We need to extract the values that will actually update the parameters
    for i in range(n_batches):
        params_vec[i] = temp_params[i,i]
    cost_vec.append(cost)  # Store cost value
    if step % 2 == 0: # optionally print the steps and cost
        print(f"{step:3d} Steps: {cost}")

cost_vec
7

Hey @CatalinaAlbornoz, thanks for the example. I have tested it and yes it seems to be slower than the previous example. For my use case right now, I am happy with the speedup provided by your previous example. Thanks again for all your help so far.

1 Like