L'électron rare
5eea6bcc5b
Extracted from micro-kiki for open-source release. Pure-torch VQC matching
PennyLane default.qubit at 1e-5 numerical precision, with autograd training
and batched inference.
Contents:
src/torch_vqc/
circuit.py — torch_vqc_forward (6-qubit StronglyEntanglingLayers)
router.py — TorchVQCRouter nn.Module + optional learned projection
__init__.py — public API exports
tests/ (10 tests, all passing)
test_circuit.py — forward match vs PennyLane, single + batched
test_training.py — loss decrease, accuracy, gradient flow, 20x speedup
test_projection.py — learned projection rescues hard tasks, shape checks
conftest.py — sys.path shim for tests without pip install -e
pyproject.toml — Apache-2.0, Python 3.10+, deps: torch+numpy, test: pennylane
README.md — quick start, benchmarks, conventions, citation
LICENSE — Apache 2.0 full text
docs/findings.md — scientific background (from micro-kiki Plan 6)
Measured speedup on 6-qubit 6-layer circuit with 120 training samples × 3 epochs:
PennyLane parameter-shift: ~99 s
torch-vqc autograd: ~30 ms
= ~3000x speedup.
89 lines
3.2 KiB
Python
89 lines
3.2 KiB
Python
"""Verify torch VQC forward matches PennyLane reference numerically.
|
|
|
|
TDD: this test must fail before torch_vqc.py exists, pass after.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pennylane as qml
|
|
import pytest
|
|
import torch
|
|
|
|
|
|
N_QUBITS = 6
|
|
N_LAYERS = 6
|
|
|
|
|
|
def _pennylane_reference(features: np.ndarray, weights: np.ndarray) -> np.ndarray:
|
|
"""Run the identical circuit through PennyLane — ground truth."""
|
|
dev = qml.device("default.qubit", wires=N_QUBITS)
|
|
|
|
@qml.qnode(dev)
|
|
def circuit(w, f):
|
|
qml.AngleEmbedding(f[:N_QUBITS], wires=range(N_QUBITS))
|
|
qml.StronglyEntanglingLayers(w, wires=range(N_QUBITS))
|
|
return [qml.expval(qml.PauliZ(i)) for i in range(N_QUBITS)]
|
|
|
|
return np.array(circuit(weights, features))
|
|
|
|
|
|
def test_forward_matches_pennylane_single_sample():
|
|
from torch_vqc.circuit import torch_vqc_forward
|
|
|
|
rng = np.random.default_rng(0)
|
|
features = rng.uniform(0.0, 2 * np.pi, size=10).astype(np.float64)
|
|
weights = rng.uniform(0.0, 2 * np.pi, size=(N_LAYERS, N_QUBITS, 3)).astype(np.float64)
|
|
|
|
ref = _pennylane_reference(features, weights) # shape (N_QUBITS,)
|
|
|
|
f_t = torch.from_numpy(features).double()
|
|
w_t = torch.from_numpy(weights).double()
|
|
out = torch_vqc_forward(f_t, w_t, n_qubits=N_QUBITS, n_layers=N_LAYERS)
|
|
|
|
assert out.shape == (N_QUBITS,), f"expected ({N_QUBITS},), got {tuple(out.shape)}"
|
|
np.testing.assert_allclose(
|
|
out.detach().cpu().numpy(),
|
|
ref,
|
|
atol=1e-5,
|
|
err_msg=f"torch={out}, pennylane={ref}",
|
|
)
|
|
|
|
|
|
def test_forward_matches_pennylane_batched():
|
|
from torch_vqc.circuit import torch_vqc_forward
|
|
|
|
rng = np.random.default_rng(1)
|
|
B = 3
|
|
features = rng.uniform(0.0, 2 * np.pi, size=(B, 10)).astype(np.float64)
|
|
weights = rng.uniform(0.0, 2 * np.pi, size=(N_LAYERS, N_QUBITS, 3)).astype(np.float64)
|
|
|
|
refs = np.stack([_pennylane_reference(features[b], weights) for b in range(B)], axis=0)
|
|
|
|
f_t = torch.from_numpy(features).double()
|
|
w_t = torch.from_numpy(weights).double()
|
|
out = torch_vqc_forward(f_t, w_t, n_qubits=N_QUBITS, n_layers=N_LAYERS)
|
|
|
|
assert out.shape == (B, N_QUBITS), f"expected ({B}, {N_QUBITS}), got {tuple(out.shape)}"
|
|
np.testing.assert_allclose(
|
|
out.detach().cpu().numpy(),
|
|
refs,
|
|
atol=1e-5,
|
|
err_msg=f"batched mismatch: torch={out}, pennylane={refs}",
|
|
)
|
|
|
|
|
|
def test_forward_different_seeds_different_outputs():
|
|
"""Sanity check: distinct weights → distinct outputs (not constant function)."""
|
|
from torch_vqc.circuit import torch_vqc_forward
|
|
|
|
rng = np.random.default_rng(42)
|
|
features = rng.uniform(0.0, 2 * np.pi, size=10).astype(np.float64)
|
|
w1 = rng.uniform(0.0, 2 * np.pi, size=(N_LAYERS, N_QUBITS, 3)).astype(np.float64)
|
|
w2 = rng.uniform(0.0, 2 * np.pi, size=(N_LAYERS, N_QUBITS, 3)).astype(np.float64)
|
|
|
|
f_t = torch.from_numpy(features).double()
|
|
o1 = torch_vqc_forward(f_t, torch.from_numpy(w1).double(), n_qubits=N_QUBITS, n_layers=N_LAYERS)
|
|
o2 = torch_vqc_forward(f_t, torch.from_numpy(w2).double(), n_qubits=N_QUBITS, n_layers=N_LAYERS)
|
|
|
|
assert not torch.allclose(o1, o2, atol=1e-3), "different weights should produce different outputs"
|