electron-rare/torch-vqc
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: initial import — torch-vqc v0.1.0

Browse Source 
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.
This commit is contained in:
L'électron rare
2026-04-19 16:34:24 +02:00
commit 5eea6bcc5b
12 changed files with 1121 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+13
View File
@@ -0,0 +1,13 @@
__pycache__/
*.py[cod]
*.so
.venv/
venv/
.env
*.egg-info/
dist/
build/
.pytest_cache/
.ruff_cache/
.mypy_cache/
.DS_Store
LICENSE
+202
View File
@@ -0,0 +1,202 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
README.md
+108
View File
@@ -0,0 +1,108 @@
# torch-vqc
Pure-torch variational quantum circuit (VQC) with autograd training — **~3000× faster** than PennyLane parameter-shift on 6-qubit StronglyEntanglingLayers.
## Why
PennyLane's `default.qubit` + parameter-shift gradient is the standard for VQC research, but it's painfully slow:
- **Per gradient step**: 2 extra forward passes per parameter (~216 for a 6×6×3 StronglyEntanglingLayers)
- **No batching**: samples processed one at a time in a Python loop
- **Result**: training a small VQC classifier on 400 samples × 10 epochs takes minutes
`torch-vqc` re-implements the same circuit as explicit state-vector math in torch:
- **Autograd backprop** replaces parameter-shift: one backward pass for the whole gradient
- **Batched**: all samples in parallel via `einsum`/`index_select` on the state tensor
- **Validated**: forward output matches PennyLane `default.qubit` to `1e-5`
Benchmark (6 qubits, 6 layers, 120 train samples × 3 epochs):
| Backend | Time | Speedup |
|---|---|---|
| PennyLane parameter-shift | ~99 s | 1× |
| torch-vqc autograd | ~30 ms | **~3000×** |
## Install
```bash
pip install torch-vqc
# or for dev:
git clone https://github.com/electron-rare/torch-vqc
cd torch-vqc && pip install -e ".[dev]"
```
## Quick start
```python
import torch
from torch_vqc import torch_vqc_forward, TorchVQCRouter
# Low-level: run a circuit forward pass directly
features = torch.randn(32, 10, dtype=torch.float64) # batch of 32, 10 input features
weights = torch.randn(6, 6, 3, dtype=torch.float64) # (n_layers, n_qubits, 3)
z_expvals = torch_vqc_forward(features, weights, n_qubits=6, n_layers=6)
# → shape (32, 6), values in [-1, 1]
# High-level: VQC + classical head, ready to train
model = TorchVQCRouter(
n_qubits=4,
n_layers=6,
n_classes=10,
input_dim=384, # enables learned projection (384 → n_qubits) before circuit
weight_decay=1e-4, # L2 regularization
lr=0.05,
seed=0,
)
X = torch.randn(400, 384, dtype=torch.float64)
y = torch.randint(0, 10, (400,))
losses = model.train_batched(X, y, epochs=300)
preds = model.predict(X)
```
## What's included
| Module | Purpose |
|---|---|
| `torch_vqc.circuit.torch_vqc_forward` | Pure forward pass — AngleEmbedding + StronglyEntanglingLayers + PauliZ expectations, validated vs PennyLane at 1e-5 |
| `torch_vqc.router.TorchVQCRouter` | `nn.Module` wrapping circuit + classical head, optional learned projection, full-batch SGD with `train_batched()` |
## Conventions (matched to PennyLane default.qubit)
- Wire 0 = most significant bit in the ket notation `|q₀ q₁ ... q_{N-1}⟩`
- `Rot(φ, θ, ω) = RZ(ω) · RY(θ) · RZ(φ)` (PennyLane convention)
- `AngleEmbedding` uses RX (PennyLane default), not RY — easy to miss
- `StronglyEntanglingLayers` CNOT ring range `r_l = l % (N-1) + 1`
## Tests
```bash
pip install -e ".[test]"
pytest tests/ -v
```
Ten tests covering:
- Forward matches PennyLane (single + batched)
- Training reduces loss, reaches non-trivial accuracy
- Gradients flow through VQC weights
- Speedup ≥ 20× vs PennyLane parameter-shift reference
- Learned projection rescues architecture on tasks where truncation fails
## Background
This work was extracted from the [micro-kiki](https://github.com/genial-lab/micro-kiki) project, where it was used to execute a VQC reproducibility investigation in 2 hours of compute instead of the 3 days that PennyLane's parameter-shift would have required. See `docs/findings.md` for the full scientific context.
## License
Apache 2.0 — see `LICENSE`.
## Citation
If you use this in academic work, cite as:
```bibtex
@software{torch_vqc_2026,
author = {L'électron rare},
title = {torch-vqc: Pure-torch VQC with autograd training at 3000× PennyLane speedup},
year = {2026},
url = {https://github.com/electron-rare/torch-vqc}
}
```
docs/findings.md
+143
View File
@@ -0,0 +1,143 @@
# Plan 6 — VQC Reproducibility: final findings
**Date:** 2026-04-19
**Branch:** `poc/torch-vqc-mps`
**Method:** Torch VQC (3000× faster than PennyLane parameter-shift) used to execute
Plan 6 tasks 5-7 in minutes instead of days.
## TL;DR
**The Task 14 accuracy of 0.925 is not reproducible**, on identical code, identical
data, with a rigorous modern VQC training loop. The "97% retention at 3× compression"
claim in Paper A §4.4 is built on a non-reproducible reference number.
## Investigation trail
### Task A: git archaeology
- `quantum_router.py` at commit `8bdd376` (Task 14 timestamp, 18 apr 13:51) is
**byte-identical** to HEAD. Same `QuantumRouterConfig` defaults.
- `eval_text_jepa_vqc.py` at `a804daa` (first commit of script) is **byte-identical**
to the version that produced Task 15.5. Same `n_qubits=4, n_layers=6`, same split.
- Dataset `data/final/<domain>/train.jsonl` unchanged since before Task 14.
**Conclusion**: nothing in git explains the gap. Task 14 must have been run with
uncommitted local modifications.
### Task B: n_qubits sweep (20 runs × 200 epochs in 22 min on Studio)
```
n_qubits=4: test_acc = 0.082 ± 0.025
n_qubits=6: test_acc = 0.080 ± 0.033
n_qubits=8: test_acc = 0.076 ± 0.016
n_qubits=10: test_acc = 0.058 ± 0.013
```
All configurations stuck near chance (0.10 for 10 classes). More qubits does NOT
help. The architectural limit is **VQC only sees `features[:n_qubits]`** — the
first few dims of MiniLM mean-pooled embeddings don't carry domain-discriminating
information.
### Task C: Task 14 config reconstruction
Decoded Task 14's reported `n_params=20` + `n_test=80`:
| n_classes × n/domain | n_test | nq=4 n_layers=0 params |
|---|---|---|
| 4 × 100 | 80 | **20** ✓✓ exact match |
So Task 14 was run with:
- **4 classes** (not 10): dsp, electronics, emc, embedded (likely)
- **n_layers=0** (no variational layers — just angle embedding + measurement + linear head)
- **100 samples/domain** (400 total)
Reproducing this config with torch VQC:
```
nq=4 nl=6 n_params= 92 → test_acc=0.212
nq=4 nl=0 n_params= 20 → test_acc=0.200 ← Task 14 config reconstructed
nq=2 nl=1 n_params= 18 → test_acc=0.200
nq=3 nl=1 n_params= 25 → test_acc=0.200
```
**Task 14's 0.925 is NOT reproducible.** Current rigorous run gives ~0.20 (chance
for 4 classes).
## Implications
1. **Paper A §4.4** claim of "97% retention at 3× compression" compares Text-JEPA
at 0.900 vs baseline at 0.925, where 0.925 is an irreproducible number. The
ratio argument is invalid.
2. **Plan 6's kill criterion is inverted**: it was designed to detect non-determinism
in the pipeline. Actual result: the pipeline IS deterministic (matched grid v3 +
torch sweep confirm baseline stability at 0.19 ± 0.01). The gap is in the
historical number, not the pipeline.
3. **The real finding**: 4-qubit VQC on top of MiniLM embeddings cannot do
10-class routing. The information isn't in `features[:4]`. No amount of qubits
or layers rescues this — you'd need a different encoding (PCA to n_qubits dims,
or a learned projection).
## What this unlocks (mid-term)
With torch VQC at 3000× speedup, we can now actually do the experiments Paper A
needs to be main-track:
- Learned projection layer (384 → n_qubits) instead of truncation: tractable to sweep
- n_qubits × n_layers grid across 5 seeds: 15 min instead of 6 weeks
- DataEmbedding ablation (MiniLM vs BERT vs domain-tuned): hours instead of quarter
- Proper Bayesian bootstrap CIs on every reported number: trivial
## Recommendations
1. **Retract §4.4 numerical claims** in Paper A draft — keep the conceptual
argument but remove the "97% retention" line until we have a reproducible
baseline.
2. **Pivot Paper A** from "compression ratio" frame to "torch-native batched
VQC training" frame — this is the real contribution: making VQC research
cheap enough to iterate on.
3. **Do real-projection experiments** before submitting anywhere.
## Update (2026-04-19, later): learned projection rescues the architecture
Added optional `input_dim` to `TorchVQCRouter` — inserts a learned linear layer
`(input_dim → n_qubits)` followed by `π·tanh`, trained jointly. This replaces
the naive `features[:n_qubits]` truncation with a class-discriminative projection.
Results on 10-class real data (500 samples, 400 train, 100 test, 300 epochs):
| Config | Test acc (trunc) | Test acc (+ proj) | Δ |
|---|---|---|---|
| nq=4, nl=6 | 0.090 | **0.300** | **+21 pt** (3× chance) |
| nq=6, nl=6 | 0.070 | 0.190 | +12 pt |
| nq=8, nl=6 | 0.110 | 0.200 | +9 pt |
**The architecture is rescued.** Best is counter-intuitively the smallest qubit
count (nq=4 + proj = 0.300) — more qubits optimizer-harder on 400 samples, and
the extra capacity overfits (train 0.40 / test 0.30 for nq=4+proj already).
This is the **honest main-track-publishable contribution**: torch-native batched
VQC + learned projection makes quantum classifier research on real embeddings
tractable, at 3000× speedup vs PennyLane. The "97% retention" claim in Paper A
is still invalid — but the direction (VQC + compression) is alive again under
a proper framing.
## Update 2 (2026-04-19, evening): 3-axis rescue ablation
36 runs on Studio (~5 min) over {data=50 vs 500/dom} × {wd=0, 1e-4, 1e-3} × {linear vs MLP proj} × 3 seeds, 10 classes, nq=4, nl=6, 300 epochs.
**Headlines**:
- **Best single run: 0.390 test_acc** (50/dom + wd=1e-4 + linear + seed 0). Train=0.435, gap=0.045 — not memorizing.
- **Median: 0.250**, range 0.170–0.390. Huge seed-to-seed variance → optimization landscape rugged.
- **Linear always beats MLP** by 7.2 pt mean (likely `π·tanh` saturates MLP hidden units).
- **10× more data adds only 3.8 pt** mean — 4-qubit info channel is the real bottleneck.
- **wd=1e-4 linear is the sweet spot** (0.304 mean across data sizes).
**Upper bound estimate**: 4× chance (~0.40) appears to be the ceiling for this architecture on 10-class routing. A classical linear probe on MiniLM would hit 0.7-0.8. The VQC is NOT competitive classically — but is a rigorous quantum-ML benchmark now.
**Final recommendation for Paper A**:
- Contribution #1 (methodological): torch-native batched VQC @ 3000× over PennyLane.
- Contribution #2 (architectural): learned projection rescues VQC from uselessness on arbitrary pretrained embeddings.
- Contribution #3 (empirical): thorough ablation showing linear proj + small wd is optimal, MLP and more data give diminishing returns due to information-capacity ceiling at 4 qubits.
- **Drop any "competitive with classical" claim**. Frame as "enabling rigorous quantum-classifier research" rather than "winning".
pyproject.toml
+51
View File
@@ -0,0 +1,51 @@
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[project]
name = "torch-vqc"
version = "0.1.0"
description = "Pure-torch variational quantum circuit (VQC) with autograd training — ~3000× faster than PennyLane parameter-shift on 6-qubit StronglyEntanglingLayers."
readme = "README.md"
requires-python = ">=3.10"
license = { text = "Apache-2.0" }
authors = [{ name = "L'électron rare" }]
keywords = ["quantum-machine-learning", "vqc", "pytorch", "autograd", "pennylane", "variational-quantum-circuit"]
classifiers = [
"Development Status :: 3 - Alpha",
"Intended Audience :: Science/Research",
"License :: OSI Approved :: Apache Software License",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"Topic :: Scientific/Engineering :: Physics",
]
dependencies = [
"torch>=2.0",
"numpy>=1.24",
]
[project.optional-dependencies]
test = [
"pytest>=7.0",
"pennylane>=0.40",
]
dev = [
"pytest>=7.0",
"pennylane>=0.40",
"ruff>=0.4",
]
[project.urls]
Repository = "https://github.com/electron-rare/torch-vqc"
Issues = "https://github.com/electron-rare/torch-vqc/issues"
[tool.hatch.build.targets.wheel]
packages = ["src/torch_vqc"]
[tool.pytest.ini_options]
testpaths = ["tests"]
addopts = "-v"
src/torch_vqc/__init__.py
+9
View File
@@ -0,0 +1,9 @@
"""torch-vqc — pure-torch variational quantum circuit with autograd training.
3000× speedup over PennyLane parameter-shift on 6-qubit StronglyEntanglingLayers.
"""
from torch_vqc.circuit import torch_vqc_forward
from torch_vqc.router import TorchVQCRouter
__version__ = "0.1.0"
__all__ = ["torch_vqc_forward", "TorchVQCRouter"]
src/torch_vqc/circuit.py
+154
View File
@@ -0,0 +1,154 @@
"""Pure-torch VQC forward pass matching PennyLane's default.qubit.
Replaces PennyLane for our specific 6-qubit, 6-layer StronglyEntanglingLayers
circuit. Enables MPS/CUDA execution + autograd-based gradients (vs parameter
shift) + batched inference.
Conventions (matched to PennyLane default.qubit):
- Wire 0 = most significant bit in |q_0 q_1 ... q_{N-1}⟩.
- Basis index: idx = sum_i q_i · 2^{N-1-i}.
- Rot(φ, θ, ω) = RZ(ω) · RY(θ) · RZ(φ).
- StronglyEntanglingLayers ranges: r_l = l % (N-1) + 1.
"""
from __future__ import annotations
import torch
def _ry(theta: torch.Tensor, cdtype: torch.dtype) -> torch.Tensor:
"""RY(θ) = [[cos(θ/2), -sin(θ/2)], [sin(θ/2), cos(θ/2)]]. θ scalar → (2, 2)."""
c = torch.cos(theta / 2)
s = torch.sin(theta / 2)
zero = torch.zeros_like(c)
row0 = torch.stack([torch.complex(c, zero), torch.complex(-s, zero)])
row1 = torch.stack([torch.complex(s, zero), torch.complex(c, zero)])
return torch.stack([row0, row1]).to(cdtype)
def _rz(theta: torch.Tensor, cdtype: torch.dtype) -> torch.Tensor:
"""RZ(θ) = diag(exp(-iθ/2), exp(iθ/2)). θ scalar → (2, 2)."""
half = theta / 2
cos_h = torch.cos(half)
sin_h = torch.sin(half)
zero_c = torch.complex(torch.zeros_like(cos_h), torch.zeros_like(cos_h))
e_minus = torch.complex(cos_h, -sin_h) # exp(-iθ/2)
e_plus = torch.complex(cos_h, sin_h) # exp(+iθ/2)
row0 = torch.stack([e_minus, zero_c])
row1 = torch.stack([zero_c, e_plus])
return torch.stack([row0, row1]).to(cdtype)
def _rot(phi: torch.Tensor, theta: torch.Tensor, omega: torch.Tensor, cdtype: torch.dtype) -> torch.Tensor:
"""PennyLane convention: Rot(φ, θ, ω) = RZ(ω) @ RY(θ) @ RZ(φ)."""
return _rz(omega, cdtype) @ _ry(theta, cdtype) @ _rz(phi, cdtype)
def _ry_batched(theta: torch.Tensor, cdtype: torch.dtype) -> torch.Tensor:
"""RY for a batch of angles. theta shape (B,) → (B, 2, 2)."""
c = torch.cos(theta / 2)
s = torch.sin(theta / 2)
zero = torch.zeros_like(c)
row0 = torch.stack([torch.complex(c, zero), torch.complex(-s, zero)], dim=-1)
row1 = torch.stack([torch.complex(s, zero), torch.complex(c, zero)], dim=-1)
return torch.stack([row0, row1], dim=-2).to(cdtype)
def _rx_batched(theta: torch.Tensor, cdtype: torch.dtype) -> torch.Tensor:
"""RX for a batch of angles. theta shape (B,) → (B, 2, 2).
RX(θ) = [[cos(θ/2), -i·sin(θ/2)], [-i·sin(θ/2), cos(θ/2)]].
PennyLane AngleEmbedding default is RX.
"""
c = torch.cos(theta / 2)
s = torch.sin(theta / 2)
zero = torch.zeros_like(c)
row0 = torch.stack([torch.complex(c, zero), torch.complex(zero, -s)], dim=-1)
row1 = torch.stack([torch.complex(zero, -s), torch.complex(c, zero)], dim=-1)
return torch.stack([row0, row1], dim=-2).to(cdtype)
def _apply_1q(state: torch.Tensor, U: torch.Tensor, q: int, n: int) -> torch.Tensor:
"""Apply a single-qubit gate U (shape (2,2)) on wire q. state: (B, 2^n)."""
B = state.shape[0]
s = state.reshape(B, 2**q, 2, 2**(n - q - 1))
s = torch.einsum("ij,xajk->xaik", U, s)
return s.reshape(B, 2**n)
def _apply_1q_batched(state: torch.Tensor, U_batch: torch.Tensor, q: int, n: int) -> torch.Tensor:
"""Apply a batched single-qubit gate. U_batch: (B, 2, 2), state: (B, 2^n)."""
B = state.shape[0]
s = state.reshape(B, 2**q, 2, 2**(n - q - 1))
s = torch.einsum("xij,xajk->xaik", U_batch, s)
return s.reshape(B, 2**n)
def _apply_cnot(state: torch.Tensor, ctrl: int, targ: int, n: int) -> torch.Tensor:
"""Apply CNOT(ctrl, targ) via index permutation. CNOT is an involution."""
idx = torch.arange(2**n, device=state.device)
q_c = (idx >> (n - 1 - ctrl)) & 1
new_idx = idx ^ (q_c * (1 << (n - 1 - targ)))
return state.index_select(-1, new_idx)
def torch_vqc_forward(
features: torch.Tensor,
weights: torch.Tensor,
n_qubits: int = 6,
n_layers: int = 6,
) -> torch.Tensor:
"""Forward pass of the VQC matching PennyLane default.qubit output.
Args:
features: Shape (>=n_qubits,) single sample, or (B, >=n_qubits) batched.
Only the first n_qubits entries are used as RY embedding angles.
weights: Shape (n_layers, n_qubits, 3). Same weights applied across batch.
n_qubits: Number of wires.
n_layers: Number of StronglyEntanglingLayers layers.
Returns:
Tensor of expectation values <Z_i>, shape (n_qubits,) if unbatched input
or (B, n_qubits) if batched input. Values in [-1, 1].
"""
batched = features.dim() == 2
if not batched:
features = features.unsqueeze(0)
B = features.shape[0]
device = features.device
cdtype = torch.complex128 if features.dtype == torch.float64 else torch.complex64
state = torch.zeros(B, 2**n_qubits, dtype=cdtype, device=device)
state[:, 0] = 1.0
# AngleEmbedding = RX(features[i]) on wire i (PennyLane default is 'X')
for i in range(n_qubits):
theta = features[..., i] # (B,)
U_batch = _rx_batched(theta, cdtype)
state = _apply_1q_batched(state, U_batch, i, n_qubits)
# StronglyEntanglingLayers
for l in range(n_layers):
for q in range(n_qubits):
phi = weights[l, q, 0]
theta = weights[l, q, 1]
omega = weights[l, q, 2]
U = _rot(phi, theta, omega, cdtype)
state = _apply_1q(state, U, q, n_qubits)
r = l % (n_qubits - 1) + 1
for q in range(n_qubits):
state = _apply_cnot(state, q, (q + r) % n_qubits, n_qubits)
# Measure <Z_i> on each wire
probs = (state.conj() * state).real # (B, 2^n)
idx = torch.arange(2**n_qubits, device=device)
real_dtype = torch.float64 if cdtype == torch.complex128 else torch.float32
z_expvals = torch.zeros(B, n_qubits, dtype=real_dtype, device=device)
for i in range(n_qubits):
bit = (idx >> (n_qubits - 1 - i)) & 1
sign = 1.0 - 2.0 * bit.to(real_dtype)
z_expvals[:, i] = (probs * sign).sum(dim=-1)
if not batched:
return z_expvals.squeeze(0)
return z_expvals
src/torch_vqc/router.py
+112
View File
@@ -0,0 +1,112 @@
"""Torch-native VQC router: autograd training + batched inference.
Replaces PennyLane's parameter-shift gradient computation with torch autograd
backprop through the explicit state-vector simulation in `torch_vqc.py`.
On 6-qubit, 6-layer StronglyEntanglingLayers, this yields 2-3 orders of
magnitude speedup over `quantum_router.QuantumRouter.train` (per-sample SGD
with 108 parameter-shift evaluations per gradient step).
"""
from __future__ import annotations
import math
import numpy as np
import torch
from torch import nn
from torch_vqc.circuit import torch_vqc_forward
class TorchVQCRouter(nn.Module):
"""VQC + classical head with autograd-based training.
Parameters:
vqc_weights: (n_layers, n_qubits, 3) rotation angles
linear_w: (n_qubits, n_classes) classical head weights
linear_b: (n_classes,) classical head biases
"""
def __init__(
self,
n_qubits: int = 6,
n_layers: int = 6,
n_classes: int = 35,
lr: float = 0.01,
seed: int = 42,
input_dim: int | None = None,
hidden_dim: int | None = None,
weight_decay: float = 0.0,
) -> None:
super().__init__()
self.n_qubits = n_qubits
self.n_layers = n_layers
self.n_classes = n_classes
self.lr = lr
self.weight_decay = weight_decay
self.input_dim = input_dim
self.hidden_dim = hidden_dim
# Match QuantumRouter init: vqc weights uniform on [0, 2π)
rng = np.random.default_rng(seed)
w_init = rng.uniform(0.0, 2 * math.pi, size=(n_layers, n_qubits, 3))
self.vqc_weights = nn.Parameter(torch.from_numpy(w_init).double())
# Optional projection: linear (input_dim → n_qubits) OR MLP (input_dim → hidden_dim → n_qubits)
# Without: circuit sees features[:n_qubits] only — severe bottleneck.
# Linear: layer learns which n_qubits-D subspace carries class-discriminative signal.
# MLP (hidden_dim set): non-linear compression for richer feature extraction.
self.projection: nn.Module | None
if input_dim is not None:
if hidden_dim is not None:
self.projection = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, n_qubits),
).double()
else:
self.projection = nn.Linear(input_dim, n_qubits).double()
else:
self.projection = None
# Classical head: normal(0, 0.1) init, zero bias (match QuantumRouter)
rng2 = np.random.default_rng(seed + 1)
lw_init = rng2.standard_normal((n_qubits, n_classes)) * 0.1
self.linear_w = nn.Parameter(torch.from_numpy(lw_init).double())
self.linear_b = nn.Parameter(torch.zeros(n_classes, dtype=torch.float64))
def forward(self, features: torch.Tensor) -> torch.Tensor:
"""Compute logits for a (B, feat_dim) batch of embeddings."""
if features.dim() == 1:
features = features.unsqueeze(0)
if self.projection is not None:
# Project input → n_qubits, bound to [-π, π] for rotations
features = math.pi * torch.tanh(self.projection(features))
qubits = torch_vqc_forward(
features, self.vqc_weights, n_qubits=self.n_qubits, n_layers=self.n_layers
)
return qubits @ self.linear_w + self.linear_b
def train_batched(
self,
embeddings: torch.Tensor,
labels: torch.Tensor,
epochs: int = 10,
) -> list[float]:
"""Full-batch SGD over `epochs` passes. Returns per-epoch loss."""
opt = torch.optim.SGD(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
losses: list[float] = []
for _ in range(epochs):
opt.zero_grad()
logits = self.forward(embeddings)
loss = nn.functional.cross_entropy(logits, labels)
loss.backward()
opt.step()
losses.append(loss.item())
return losses
@torch.no_grad()
def predict(self, embeddings: torch.Tensor) -> torch.Tensor:
"""Argmax class prediction for a (B, feat_dim) batch."""
logits = self.forward(embeddings)
return logits.argmax(dim=-1)
tests/conftest.py
+7
View File
@@ -0,0 +1,7 @@
"""pytest config: add src/ to sys.path so tests work without pip install -e."""
import sys
from pathlib import Path
_SRC = Path(__file__).resolve().parent.parent / "src"
if str(_SRC) not in sys.path:
sys.path.insert(0, str(_SRC))
tests/test_circuit.py
+88
View File
@@ -0,0 +1,88 @@
"""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"
tests/test_projection.py
+77
View File
@@ -0,0 +1,77 @@
"""Projection tests: learned 384→n_qubits layer rescues accuracy on hard tasks."""
from __future__ import annotations
import numpy as np
import torch
def _make_hard_task(n_classes: int = 5, samples_per_class: int = 40, input_dim: int = 64, seed: int = 0):
"""Class info lives ONLY in dims [input_dim-6 : input_dim], so truncation to dims [0:n_qubits] sees pure noise."""
rng = np.random.default_rng(seed)
class_centers = rng.uniform(0.3, np.pi - 0.3, size=(n_classes, 6)) # mild, in rotation range
X, y = [], []
for c in range(n_classes):
for _ in range(samples_per_class):
x = rng.normal(0, 0.5, size=input_dim) # noise everywhere
x[-6:] = class_centers[c] + rng.normal(0, 0.08, size=6) # signal in last 6 dims
X.append(x)
y.append(c)
X = np.array(X, dtype=np.float64)
y = np.array(y, dtype=np.int64)
perm = rng.permutation(len(X))
return X[perm], y[perm]
def test_projection_beats_no_projection_on_hard_task():
"""Signal in last 6 dims only: truncation sees pure noise, projection can find it."""
from torch_vqc.router import TorchVQCRouter
X, y = _make_hard_task(n_classes=5, samples_per_class=40, input_dim=64, seed=3)
X_t = torch.from_numpy(X).double()
y_t = torch.from_numpy(y)
# Without projection: VQC sees X[:, :4] (first 4 dims = pure noise here)
model_no_proj = TorchVQCRouter(n_qubits=4, n_layers=4, n_classes=5, lr=0.05, seed=0)
model_no_proj.train_batched(X_t, y_t, epochs=200)
with torch.no_grad():
acc_no_proj = (model_no_proj.predict(X_t).numpy() == y).mean()
# With projection: a 4×64 layer can learn to pick from the informative last dims
model_proj = TorchVQCRouter(
n_qubits=4, n_layers=4, n_classes=5, lr=0.05, seed=0, input_dim=64
)
model_proj.train_batched(X_t, y_t, epochs=200)
with torch.no_grad():
acc_proj = (model_proj.predict(X_t).numpy() == y).mean()
print(f"\nNo projection acc = {acc_no_proj:.3f} (chance = {1/5:.3f})")
print(f"With projection acc = {acc_proj:.3f}")
assert acc_proj > acc_no_proj + 0.15, (
f"projection should give ≥15pt boost, got {acc_no_proj:.3f} → {acc_proj:.3f}"
)
def test_projection_forward_shape():
"""Projection doesn't break shape invariants."""
from torch_vqc.router import TorchVQCRouter
model = TorchVQCRouter(n_qubits=6, n_layers=6, n_classes=10, input_dim=384, seed=0)
x = torch.randn(32, 384, dtype=torch.float64)
logits = model.forward(x)
assert logits.shape == (32, 10), f"expected (32, 10), got {tuple(logits.shape)}"
def test_projection_params_counted():
"""Projection adds n_qubits × input_dim + n_qubits parameters."""
from torch_vqc.router import TorchVQCRouter
no_proj = TorchVQCRouter(n_qubits=4, n_layers=6, n_classes=10, seed=0)
with_proj = TorchVQCRouter(n_qubits=4, n_layers=6, n_classes=10, input_dim=384, seed=0)
n_no = sum(p.numel() for p in no_proj.parameters())
n_yes = sum(p.numel() for p in with_proj.parameters())
expected_extra = 4 * 384 + 4 # projection_w + projection_b
assert n_yes - n_no == expected_extra, (
f"expected {expected_extra} extra params, got {n_yes - n_no}"
)
tests/test_training.py
+157
View File
@@ -0,0 +1,157 @@
"""Training tests: autograd loop reduces loss and reaches non-trivial accuracy.
TDD: verify that TorchVQCRouter learns on a small synthetic task, replacing
PennyLane's parameter-shift training with torch autograd.
"""
from __future__ import annotations
import numpy as np
import pytest
import torch
def _make_synthetic_task(n_classes: int = 4, samples_per_class: int = 20, seed: int = 0):
"""Each class has embeddings clustered around a distinct mean in first 6 dims.
Classifiable by the first 6 features alone — what the VQC sees.
"""
rng = np.random.default_rng(seed)
X, y = [], []
for c in range(n_classes):
center = rng.uniform(0.5, 2 * np.pi - 0.5, size=6) # stay away from 0/2π
for _ in range(samples_per_class):
noise = rng.normal(0, 0.15, size=10)
x = np.zeros(10)
x[:6] = center + noise[:6]
x[6:] = noise[6:]
X.append(x)
y.append(c)
X = np.array(X, dtype=np.float64)
y = np.array(y, dtype=np.int64)
# Shuffle
perm = rng.permutation(len(X))
return X[perm], y[perm]
def test_training_reduces_loss():
from torch_vqc.router import TorchVQCRouter
X, y = _make_synthetic_task(n_classes=4, samples_per_class=10, seed=0)
model = TorchVQCRouter(n_qubits=6, n_layers=6, n_classes=4, lr=0.1, seed=42)
X_t = torch.from_numpy(X).double()
y_t = torch.from_numpy(y)
losses = model.train_batched(X_t, y_t, epochs=30)
assert len(losses) == 30
assert losses[-1] < losses[0] - 0.05, (
f"expected loss decrease >0.05, got {losses[0]:.4f} → {losses[-1]:.4f}"
)
def test_training_reaches_nontrivial_accuracy():
"""After training, train-set accuracy should clearly beat chance (25% for 4 classes)."""
from torch_vqc.router import TorchVQCRouter
X, y = _make_synthetic_task(n_classes=4, samples_per_class=15, seed=1)
model = TorchVQCRouter(n_qubits=6, n_layers=6, n_classes=4, lr=0.1, seed=7)
X_t = torch.from_numpy(X).double()
y_t = torch.from_numpy(y)
model.train_batched(X_t, y_t, epochs=50)
with torch.no_grad():
logits = model.forward(X_t)
preds = logits.argmax(dim=-1)
acc = (preds == y_t).float().mean().item()
assert acc > 0.50, f"expected >50% (well above 25% chance), got {acc:.3f}"
def test_gradients_flow_through_vqc_weights():
"""Sanity: VQC weights must receive non-zero gradient."""
from torch_vqc.router import TorchVQCRouter
X, y = _make_synthetic_task(n_classes=3, samples_per_class=5, seed=2)
model = TorchVQCRouter(n_qubits=6, n_layers=6, n_classes=3, lr=0.05, seed=3)
X_t = torch.from_numpy(X).double()
y_t = torch.from_numpy(y)
# One forward + backward
logits = model.forward(X_t)
loss = torch.nn.functional.cross_entropy(logits, y_t)
loss.backward()
assert model.vqc_weights.grad is not None
assert model.vqc_weights.grad.abs().max().item() > 1e-6, (
"VQC weight gradient is zero — autograd not flowing through circuit"
)
assert model.linear_w.grad is not None
assert model.linear_w.grad.abs().max().item() > 1e-6
def test_training_faster_than_pennylane_parameter_shift():
"""Benchmark: torch autograd training must be >=20× faster than PennyLane parameter-shift.
Builds a minimal per-sample parameter-shift loop with PennyLane to compare
head-to-head. This is the whole point of the package — if this regresses,
the main contribution is broken.
"""
import time
import pennylane as qml
from torch_vqc.router import TorchVQCRouter
X, y = _make_synthetic_task(n_classes=4, samples_per_class=10, seed=4)
epochs = 3
# Torch autograd (full batch)
model = TorchVQCRouter(n_qubits=6, n_layers=6, n_classes=4, lr=0.05, seed=0)
X_t = torch.from_numpy(X).double()
y_t = torch.from_numpy(y)
_ = model.train_batched(X_t[:5], y_t[:5], epochs=1) # warmup
t0 = time.perf_counter()
model.train_batched(X_t, y_t, epochs=epochs)
torch_time = time.perf_counter() - t0
# PennyLane parameter-shift reference — minimal per-sample SGD loop
n_q, n_l, n_c = 6, 6, 4
dev = qml.device("default.qubit", wires=n_q)
@qml.qnode(dev)
def circuit(w, f):
qml.AngleEmbedding(f[:n_q], wires=range(n_q))
qml.StronglyEntanglingLayers(w, wires=range(n_q))
return [qml.expval(qml.PauliZ(i)) for i in range(n_q)]
rng = np.random.default_rng(0)
w = rng.uniform(0, 2 * np.pi, size=(n_l, n_q, 3))
lw = rng.standard_normal((n_q, n_c)) * 0.1
lb = np.zeros(n_c)
lr = 0.05
t0 = time.perf_counter()
for _ in range(epochs):
for x, label in zip(X, y):
q = np.array(circuit(w, x))
logits = q @ lw + lb
probs = np.exp(logits - logits.max())
probs /= probs.sum()
d_logits = probs.copy()
d_logits[int(label)] -= 1
lw -= lr * np.outer(q, d_logits)
lb -= lr * d_logits
dl_dq = lw @ d_logits
# parameter-shift for each VQC weight
shift = np.pi / 2
for idx in np.ndindex(*w.shape):
w_p, w_m = w.copy(), w.copy()
w_p[idx] += shift
w_m[idx] -= shift
dq = (np.array(circuit(w_p, x)) - np.array(circuit(w_m, x))) / 2.0
w[idx] -= lr * float(dl_dq @ dq)
pl_time = time.perf_counter() - t0
speedup = pl_time / torch_time
print(f"\nSpeedup: {speedup:.1f}× (torch {torch_time*1000:.0f}ms vs PL parameter-shift {pl_time*1000:.0f}ms)")
assert speedup >= 20.0, f"expected ≥20× speedup, got {speedup:.1f}×"