Qubit rotation

Using qubit rotation example to understand basic syntax of pennylane, gradient descent. See this link for more details.

import pennylane as qml
from jax import numpy as np
import jax
import jaxopt

dev1 = qml.device("lightning.qubit",wires=1)

@qml.qnode(dev1)
def circuit(params):
    qml.RX(params[0], wires=0)
    qml.RY(params[1], wires=0)
    return qml.expval(qml.PauliZ(0))


# how the gradiat is calculated?
dcircuit = jax.grad(circuit, argnums=0)

def cost(x):
    return circuit(x)

init_params = np.array([0.011, 0.012])

opt = jaxopt.GradientDescent(cost, stepsize=1, acceleration=False)

steps = 100
params = init_params
opt_state = opt.init_state(params)
for i in range(steps):
    params, opt_state = opt.update(params, opt_state)
    if (i + 1) % 5 == 0:
        print("Cost after step {:5d}: {: .7f}".format(i + 1, cost(params)))
print("Optimized rotation angles: {}".format(params))