Using `argnum` in parameter shift rule

PennyLane Help
2

Hi @Pratibha_Hegde , welcome to the Forum!

argnum counts the arguments of your circuit so argnum=0 indicates ‘w’ and argnum=1 indicates ‘x’. Note that argnum starts counting at zero.

I think your issue here was in how you were calculating the gradients. I’ve adapted your code below to show the output of calculating the gradient with respect to both arguments and each of them individually. I hope this helps!

import pennylane as qml
from pennylane import numpy as np
from pennylane.templates import StronglyEntanglingLayers
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')
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=2, 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
g = qml.grad(circuit) # Gradient wrt w and x
g0 = qml.grad(circuit,argnum=0) # Gradient wrt w
g1 = qml.grad(circuit,argnum=1) # Gradient wrt x


d = []
d0 = []
d1 = []

for i in X:
    d.append(g(weights, i))
    d0.append(g0(weights, i))
    d1.append(g1(weights, i))
  
print('d: ',d)
print('d0: ',d0)
print('d1: ',d1)