In [5]:
!pip install pennylane
!pip install jax==0.5.3
Requirement already satisfied: pennylane in /usr/local/lib/python3.12/dist-packages (0.42.3) Requirement already satisfied: scipy in /usr/local/lib/python3.12/dist-packages (from pennylane) (1.16.2) Requirement already satisfied: networkx in /usr/local/lib/python3.12/dist-packages (from pennylane) (3.5) Requirement already satisfied: rustworkx>=0.14.0 in /usr/local/lib/python3.12/dist-packages (from pennylane) (0.17.1) Requirement already satisfied: autograd in /usr/local/lib/python3.12/dist-packages (from pennylane) (1.8.0) Requirement already satisfied: appdirs in /usr/local/lib/python3.12/dist-packages (from pennylane) (1.4.4) Requirement already satisfied: autoray<0.8,>=0.6.11 in /usr/local/lib/python3.12/dist-packages (from pennylane) (0.7.2) Requirement already satisfied: cachetools in /usr/local/lib/python3.12/dist-packages (from pennylane) (5.5.2) Requirement already satisfied: pennylane-lightning>=0.42 in /usr/local/lib/python3.12/dist-packages (from pennylane) (0.43.0) Requirement already satisfied: requests in /usr/local/lib/python3.12/dist-packages (from pennylane) (2.32.4) Requirement already satisfied: tomlkit in /usr/local/lib/python3.12/dist-packages (from pennylane) (0.13.3) Requirement already satisfied: typing_extensions in /usr/local/lib/python3.12/dist-packages (from pennylane) (4.15.0) Requirement already satisfied: packaging in /usr/local/lib/python3.12/dist-packages (from pennylane) (25.0) Requirement already satisfied: diastatic-malt in /usr/local/lib/python3.12/dist-packages (from pennylane) (2.15.2) Requirement already satisfied: numpy in /usr/local/lib/python3.12/dist-packages (from pennylane) (2.0.2) Requirement already satisfied: scipy-openblas32>=0.3.26 in /usr/local/lib/python3.12/dist-packages (from pennylane-lightning>=0.42->pennylane) (0.3.30.0.2) Requirement already satisfied: astunparse in /usr/local/lib/python3.12/dist-packages (from diastatic-malt->pennylane) (1.6.3) Requirement already satisfied: gast in /usr/local/lib/python3.12/dist-packages (from diastatic-malt->pennylane) (0.6.0) Requirement already satisfied: termcolor in /usr/local/lib/python3.12/dist-packages (from diastatic-malt->pennylane) (3.1.0) Requirement already satisfied: charset_normalizer<4,>=2 in /usr/local/lib/python3.12/dist-packages (from requests->pennylane) (3.4.3) Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.12/dist-packages (from requests->pennylane) (3.10) Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.12/dist-packages (from requests->pennylane) (2.5.0) Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.12/dist-packages (from requests->pennylane) (2025.10.5) Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.12/dist-packages (from astunparse->diastatic-malt->pennylane) (0.45.1) Requirement already satisfied: six<2.0,>=1.6.1 in /usr/local/lib/python3.12/dist-packages (from astunparse->diastatic-malt->pennylane) (1.17.0) Requirement already satisfied: jax==0.5.3 in /usr/local/lib/python3.12/dist-packages (0.5.3) Requirement already satisfied: jaxlib<=0.5.3,>=0.5.3 in /usr/local/lib/python3.12/dist-packages (from jax==0.5.3) (0.5.3) Requirement already satisfied: ml_dtypes>=0.4.0 in /usr/local/lib/python3.12/dist-packages (from jax==0.5.3) (0.5.3) Requirement already satisfied: numpy>=1.25 in /usr/local/lib/python3.12/dist-packages (from jax==0.5.3) (2.0.2) Requirement already satisfied: opt_einsum in /usr/local/lib/python3.12/dist-packages (from jax==0.5.3) (3.4.0) Requirement already satisfied: scipy>=1.11.1 in /usr/local/lib/python3.12/dist-packages (from jax==0.5.3) (1.16.2)
In [21]:
import pennylane.numpy as np
import pennylane as qml
import jax
dev = qml.device("default.mixed", wires=2)
@qml.qnode(dev)
def circuit(param):
qml.RX(param, wires=0)
qml.DepolarizingChannel(p=0.1, wires=1)
qml.CNOT(wires=[0, 1])
return qml.expval(qml.PauliZ(1))
In [22]:
print("\nGradient Descent")
print("---------------")
grad_circuit = jax.grad(circuit)
param = 0.42
print(f"Initial param: {param:0.3f}")
print(f"Initial cost: {circuit(param):0.3f}")
for _ in range(100):
param -= grad_circuit(param)
print(f"Tuned param: {param:0.3f}")
print(f"Tuned cost: {circuit(param):0.3f}")
Gradient Descent --------------- Initial param: 0.420 Initial cost: 0.913 Tuned param: 3.142 Tuned cost: -1.000
In [23]:
import jax
import jax.numpy as jnp
from jax import random, grad
# data
key = random.PRNGKey(0)
x = random.normal(key, (10, 3))
y = jnp.array([[1.0], [0.0], [1.0], [0.0], [1.0],
[0.0], [1.0], [0.0], [1.0], [0.0]])
# init weights
def init_params(key):
k1, k2 = random.split(key)
W1 = random.normal(k1, (3, 4))
b1 = jnp.zeros(4)
W2 = random.normal(k2, (4, 1))
b2 = jnp.zeros(1)
return (W1, b1, W2, b2)
# model
def forward(params, x):
W1, b1, W2, b2 = params
h = jnp.tanh(x @ W1 + b1)
h = jax.nn.sigmoid(h @ W2 + b2) * jnp.pi
return circuit(jnp.mean(h))
# loss
def loss_fn(params, x, y):
preds = forward(params, x)
return jnp.mean((preds - y) ** 2)
# training step
@jax.jit
def update(params, x, y, lr=0.1):
grads = grad(loss_fn)(params, x, y)
return jax.tree_util.tree_map(lambda p, g: p - lr * g, params, grads)
# train
params = init_params(key)
for i in range(200):
params = update(params, x, y)
if i % 50 == 0:
print(i, loss_fn(params, x, y))
0 0.4422666 50 0.25000006 100 0.25 150 0.25
In [ ]: