Using `argnum` in parameter shift rule

PennyLane Help
5

Hi @CatalinaAlbornoz, thank you once again and I apologize for getting back quite late on the topic. I am not using the second derivate for the optimization loop for now. I used qml.jacobian(qml.jacobian(circuit, argnum=1), argnum=1) to compute the second gradient and it works fine. I also tried to use qml.gradients.param_shift_hessian() function to compute the second derivatives directly w.r.t. x. Again, I am stuck in using argnum in this case. I resorted to count the arguments like before(starting from 0 this time) where I was counting the rotation angles parameters of the first layer of gates followed the by the second layer which encodes the variable x. Since x appears r times in the second layer, it seems I have to give all the argnum= between 3r to 4r. With this method, the hessian that I obtained is the same as qml.jacobian(qml.jacobian(circuit, argnum=1), argnum=1) method, and it is considerably faster and that is why I would like to use the qml.gradients.param_shift_hessian(). However, I wanted to make sure that counting argument number this way is correct (which was clearly wrong when using qml.grad()). I would be very grateful if you could share your thoughts and suggestions on this. I have attached the code that I am using below:

import pennylane as qml
from pennylane import numpy as np
from pennylane.templates import StronglyEntanglingLayers
import time

dev = qml.device('default.qubit', wires=3)

r=3 # num_qubits

np.random.seed(42)
def S(x):
    """Data encoding circuit block."""
    for w in range(r):
        qml.RX(x, wires=w)


def W(theta):
    """Trainable circuit block."""
    StronglyEntanglingLayers(theta, wires=range(r))

# You can add the differentiation method when you define your QNode
@qml.qnode(dev,diff_method='parameter-shift', max_diff=2)
def circuit(w, x):
    W(w[0])
    S(x)
    W(w[1])
    return qml.expval(qml.PauliZ(wires=0))

X = np.linspace(0., 1., num=10, requires_grad=True)
weights = 2 * np.pi * np.random.random(size=(2, 1, r, 3), requires_grad=True)

# Draw your circuit
qml.draw_mpl(circuit,decimals=1,expansion_strategy='device')(weights, X[0])

# Calculate your gradient function with respect to all or just one argument
g2 = qml.jacobian(qml.jacobian(circuit, argnum=1), argnum=1)

#Defining Hessian
def hessian(func, w, x):
    argnum = np.arange(3*r, 4*r, 1) #argument index corresponding to x
    hess = [qml.gradients.param_shift_hessian(func, argnum=argnum)(w, x_)[1] for x_ in x]
    return hess

d2 = []
h = []

t1 = time.time()
for i in X:
    d2.append(g2(weights, i))
t2 = time.time()
print('time to compute jacobian(jacobian):', t2-t1)

t3 = time.time()
h.append(hessian(circuit, weights, X))
t4 = time.time()
print('time to compute hessian:', t4-t3)

print('d2 using jacobian(jacobian):', d2)
print('h using hessian:', h)