mlx-quantum 0.1.0__tar.gz

mlx_quantum-0.1.0/.gitignore

@@ -0,0 +1,18 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+# macOS
+.DS_Store
+
+# Generated example output (transient)
+examples/*.png
+benchmarks/run.log
+.pytest_cache/

mlx_quantum-0.1.0/CHANGELOG.md

@@ -0,0 +1,41 @@
+# Changelog
+
+All notable changes to this project are documented here. The format follows
+[Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project
+adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [Unreleased]
+
+## [0.1.0] - 2026-07-05
+
+Initial release: a native, GPU-accelerated, differentiable quantum-machine-learning
+library in pure Apple MLX.
+
+### Added
+- `mlx_quantum.statevector` — batched complex statevector simulation with gates
+  applied as `einsum` contractions: `zero_state`, `apply_1q`, `apply_2q`,
+  `expval_z`, `expval_all_z`, and gates `H, X, Y, Z, rx, ry, rz, CX, CZ`. Qubit
+  ordering is little-endian, matching Qiskit's amplitude layout exactly.
+- `QuantumLayer` — a trainable `mlx.nn.Module` running a hardware-efficient
+  ansatz; its weights train by ordinary MLX autodiff, no custom vjp.
+- Examples: `simple_mlp.py` (hybrid classical-quantum MLP) and
+  `benchmark_vs_qiskit.py` (accuracy + speed vs Qiskit's `EstimatorQNN`).
+- Validation suite (`benchmarks/validate.py`) producing five scope-labeled graphs:
+  accuracy vs qubits, two-baseline speedup (vs `EstimatorQNN` and vs Aer),
+  wall-time to the memory cliff, error distribution over a random-circuit sweep,
+  and a training-curve overlay.
+- Tests covering statevector/gate correctness, gate unitarity, state-norm
+  preservation, the little-endian convention, layer training, a finite-difference
+  gradient check, and forward + weight-gradient + input-gradient parity with
+  Qiskit across a random gate sweep (shared engine in `tests/_circuits.py`;
+  Qiskit-dependent tests auto-skip when Qiskit is absent).
+
+### Notes
+- Runtime dependencies are MLX and NumPy only; Qiskit, Aer, and matplotlib are
+  optional (`[examples]` extra), used for validation and benchmarking.
+- Forward values and gradients match Qiskit to ~1e-6 (noiseless statevector,
+  float32); ~100–400× faster end-to-end than driving `EstimatorQNN` from Python,
+  and ~1.7–3× faster than Aer's compiled statevector on the forward pass.
+
+[Unreleased]: https://github.com/UmarGit/mlx_quantum/compare/v0.1.0...HEAD
+[0.1.0]: https://github.com/UmarGit/mlx_quantum/releases/tag/v0.1.0

mlx_quantum-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Umar Ahmed
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mlx_quantum-0.1.0/PKG-INFO

@@ -0,0 +1,173 @@
+Metadata-Version: 2.4
+Name: mlx-quantum
+Version: 0.1.0
+Summary: Fast, differentiable quantum-machine-learning in pure Apple MLX — GPU statevector simulation with native autodiff.
+Project-URL: Repository, https://github.com/UmarGit/mlx_quantum
+License-Expression: MIT
+License-File: LICENSE
+Keywords: apple-silicon,differentiable,mlx,quantum,quantum-machine-learning
+Requires-Python: >=3.13
+Requires-Dist: mlx>=0.31.2
+Requires-Dist: numpy>=2.5.1
+Provides-Extra: examples
+Requires-Dist: matplotlib>=3.9; extra == 'examples'
+Requires-Dist: qiskit-aer>=0.17.2; extra == 'examples'
+Requires-Dist: qiskit-algorithms>=0.4.0; extra == 'examples'
+Requires-Dist: qiskit-machine-learning>=0.9.0; extra == 'examples'
+Requires-Dist: qiskit>=2.5.0; extra == 'examples'
+Description-Content-Type: text/markdown
+
+# mlx-quantum
+
+[![CI](https://github.com/UmarGit/mlx_quantum/actions/workflows/ci.yml/badge.svg)](https://github.com/UmarGit/mlx_quantum/actions/workflows/ci.yml)
+
+Fast, differentiable quantum machine learning in pure [Apple MLX](https://github.com/ml-explore/mlx).
+
+Statevector simulation runs on the Metal GPU and is differentiable end-to-end
+through MLX autodiff — so a quantum layer trains like any other `mlx.nn` module,
+with no custom gradient code. Forward values **and** gradients match Qiskit to
+~1e-6 (float32), and it is **~100–400× faster** end-to-end than driving Qiskit's
+`EstimatorQNN` from Python. See [Validation](#validation) for the evidence.
+
+## Install
+
+```bash
+uv add mlx-quantum
+# or
+pip install mlx-quantum
+```
+
+Requires Python ≥ 3.13 and Apple Silicon. The library depends only on MLX and
+NumPy; Qiskit is optional and used solely to cross-validate/benchmark.
+
+## Quickstart
+
+`QuantumLayer` is a trainable `mlx.nn.Module`. Drop it into a model and train:
+
+```python
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_quantum import QuantumLayer
+
+class HybridMLP(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.pre = nn.Linear(8, 4)
+        self.qnn = QuantumLayer(num_qubits=4, reps=2)   # trainable quantum layer
+        self.post = nn.Linear(4, 3)
+
+    def __call__(self, x):
+        x = mx.tanh(self.pre(x)) * mx.pi                # encode into rotation angles
+        return self.post(self.qnn(x))
+```
+
+The quantum weights are ordinary MLX parameters — `nn.value_and_grad` and any
+optimizer update them automatically:
+
+```python
+loss_and_grad = nn.value_and_grad(model, loss_fn)
+loss, grads = loss_and_grad(model, x, y)
+optimizer.update(model, grads)
+```
+
+## Building custom circuits
+
+`QuantumLayer` runs a hardware-efficient ansatz, but the simulator primitives are
+public — build any circuit as a plain differentiable function:
+
+```python
+import mlx.core as mx
+from mlx_quantum import zero_state, apply_1q, apply_2q, expval_all_z, H, ry, CX
+
+def circuit(x, weights):            # x: (batch, n) angles, weights: (n,)
+    n = x.shape[1]
+    state = zero_state(x.shape[0], n)
+    for q in range(n):
+        state = apply_1q(state, H, q)
+    for q in range(n):
+        state = apply_1q(state, ry(x[:, q]), q)          # per-sample encoding
+    for q in range(n):
+        state = apply_1q(state, ry(weights[q]), q)        # trainable
+    for q in range(n - 1):
+        state = apply_2q(state, CX, q, q + 1)
+    return expval_all_z(state)       # <Z> per qubit, shape (batch, n)
+
+grads = mx.grad(lambda w: mx.sum(circuit(x, w)))(weights)  # just works
+```
+
+Gates provided: `H, X, Y, Z, rx, ry, rz, CX, CZ`. Add your own — a single-qubit
+gate is any `(2, 2)` complex `mx.array`; a two-qubit gate is a `(2, 2, 2, 2)`
+tensor `[out0, out1, in0, in1]`.
+
+## How it works
+
+A statevector is a complex `mx.array` of shape `(batch,) + (2,) * num_qubits`;
+qubit ordering is little-endian to match Qiskit exactly (qubit 0 is the fastest
+index, so flattening reproduces Qiskit's amplitude order). Gates are
+applied as `einsum` contractions, so the entire simulation is differentiable and
+GPU-resident. Because there is no custom `vjp` and no NumPy round-trip, `mx.grad`
+differentiates the circuit directly — including through complex amplitudes.
+
+Two MLX specifics the implementation works around: the initial state is built as a
+constant (not an in-place assignment, which compiles to an unsupported complex
+GPU scatter), and gates are contractions rather than `take`/gather (whose backward
+is also a scatter).
+
+## Examples
+
+```bash
+uv run python examples/simple_mlp.py                             # hybrid MLP training
+uv run --extra examples python examples/benchmark_vs_qiskit.py   # quick speed + accuracy vs Qiskit
+```
+
+## Validation
+
+Two tracks — is it correct, and is the speed claim fair? All measurements are
+**noiseless statevector, float32**. Regenerate with
+`uv run --extra examples python benchmarks/validate.py` (details in
+[`benchmarks/`](benchmarks/)).
+
+**Correctness.** Forward values and gradients are compared against Qiskit
+(`Statevector` and `ReverseEstimatorGradient`) over 142 random circuits covering
+every gate (`H, X, Y, Z, rx, ry, rz, CX, CZ`); per-circuit error stays at ~1e-6.
+(The *batch-summed* gradient error on the accuracy plot climbs to ~1e-5 by 8
+qubits — that is float32 accumulation from summing 128 terms into one number,
+still ≥5 significant figures, not a modelling error.) Gates are checked for
+unitarity, the state norm is checked after every layer, and an asymmetric circuit
+pins the little-endian qubit order to Qiskit's.
+
+| | ![accuracy](benchmarks/accuracy_vs_qubits.png) | ![error distribution](benchmarks/error_distribution.png) |
+|---|---|---|
+
+**Performance.** Two honest baselines. End-to-end vs `EstimatorQNN` driven from
+Python (~100–400×), and kernel-level vs Aer's compiled statevector estimator
+(~1.7–3×, forward only) — so the win is not just deleted orchestration. Wall-time
+is shown until MLX hits the memory cliff (~22–26 qubits, single statevector).
+
+| | ![speedup](benchmarks/speedup_vs_qubits.png) | ![wall time](benchmarks/walltime_vs_qubits.png) |
+|---|---|---|
+
+**Trains identically.** Same circuit, init, data, and optimizer (SGD): the MLX
+layer (`mx.grad`) and the Qiskit QNN (`qnn.backward`) produce the same loss curve
+to ~1e-7 — same training, just faster.
+
+![training overlay](benchmarks/training_overlay.png)
+
+## Tests
+
+```bash
+uv run pytest
+```
+
+Covers gate/statevector correctness, gate unitarity and norm preservation, the
+little-endian convention, layer training, a finite-difference gradient check, and
+forward + weight-gradient + input-gradient parity with Qiskit across a random gate
+sweep (Qiskit-dependent tests skip automatically if Qiskit is absent).
+
+## Changelog
+
+See [CHANGELOG.md](CHANGELOG.md).
+
+## License
+
+MIT — see [LICENSE](LICENSE).

mlx_quantum-0.1.0/README.md

@@ -0,0 +1,154 @@
+# mlx-quantum
+
+[![CI](https://github.com/UmarGit/mlx_quantum/actions/workflows/ci.yml/badge.svg)](https://github.com/UmarGit/mlx_quantum/actions/workflows/ci.yml)
+
+Fast, differentiable quantum machine learning in pure [Apple MLX](https://github.com/ml-explore/mlx).
+
+Statevector simulation runs on the Metal GPU and is differentiable end-to-end
+through MLX autodiff — so a quantum layer trains like any other `mlx.nn` module,
+with no custom gradient code. Forward values **and** gradients match Qiskit to
+~1e-6 (float32), and it is **~100–400× faster** end-to-end than driving Qiskit's
+`EstimatorQNN` from Python. See [Validation](#validation) for the evidence.
+
+## Install
+
+```bash
+uv add mlx-quantum
+# or
+pip install mlx-quantum
+```
+
+Requires Python ≥ 3.13 and Apple Silicon. The library depends only on MLX and
+NumPy; Qiskit is optional and used solely to cross-validate/benchmark.
+
+## Quickstart
+
+`QuantumLayer` is a trainable `mlx.nn.Module`. Drop it into a model and train:
+
+```python
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_quantum import QuantumLayer
+
+class HybridMLP(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.pre = nn.Linear(8, 4)
+        self.qnn = QuantumLayer(num_qubits=4, reps=2)   # trainable quantum layer
+        self.post = nn.Linear(4, 3)
+
+    def __call__(self, x):
+        x = mx.tanh(self.pre(x)) * mx.pi                # encode into rotation angles
+        return self.post(self.qnn(x))
+```
+
+The quantum weights are ordinary MLX parameters — `nn.value_and_grad` and any
+optimizer update them automatically:
+
+```python
+loss_and_grad = nn.value_and_grad(model, loss_fn)
+loss, grads = loss_and_grad(model, x, y)
+optimizer.update(model, grads)
+```
+
+## Building custom circuits
+
+`QuantumLayer` runs a hardware-efficient ansatz, but the simulator primitives are
+public — build any circuit as a plain differentiable function:
+
+```python
+import mlx.core as mx
+from mlx_quantum import zero_state, apply_1q, apply_2q, expval_all_z, H, ry, CX
+
+def circuit(x, weights):            # x: (batch, n) angles, weights: (n,)
+    n = x.shape[1]
+    state = zero_state(x.shape[0], n)
+    for q in range(n):
+        state = apply_1q(state, H, q)
+    for q in range(n):
+        state = apply_1q(state, ry(x[:, q]), q)          # per-sample encoding
+    for q in range(n):
+        state = apply_1q(state, ry(weights[q]), q)        # trainable
+    for q in range(n - 1):
+        state = apply_2q(state, CX, q, q + 1)
+    return expval_all_z(state)       # <Z> per qubit, shape (batch, n)
+
+grads = mx.grad(lambda w: mx.sum(circuit(x, w)))(weights)  # just works
+```
+
+Gates provided: `H, X, Y, Z, rx, ry, rz, CX, CZ`. Add your own — a single-qubit
+gate is any `(2, 2)` complex `mx.array`; a two-qubit gate is a `(2, 2, 2, 2)`
+tensor `[out0, out1, in0, in1]`.
+
+## How it works
+
+A statevector is a complex `mx.array` of shape `(batch,) + (2,) * num_qubits`;
+qubit ordering is little-endian to match Qiskit exactly (qubit 0 is the fastest
+index, so flattening reproduces Qiskit's amplitude order). Gates are
+applied as `einsum` contractions, so the entire simulation is differentiable and
+GPU-resident. Because there is no custom `vjp` and no NumPy round-trip, `mx.grad`
+differentiates the circuit directly — including through complex amplitudes.
+
+Two MLX specifics the implementation works around: the initial state is built as a
+constant (not an in-place assignment, which compiles to an unsupported complex
+GPU scatter), and gates are contractions rather than `take`/gather (whose backward
+is also a scatter).
+
+## Examples
+
+```bash
+uv run python examples/simple_mlp.py                             # hybrid MLP training
+uv run --extra examples python examples/benchmark_vs_qiskit.py   # quick speed + accuracy vs Qiskit
+```
+
+## Validation
+
+Two tracks — is it correct, and is the speed claim fair? All measurements are
+**noiseless statevector, float32**. Regenerate with
+`uv run --extra examples python benchmarks/validate.py` (details in
+[`benchmarks/`](benchmarks/)).
+
+**Correctness.** Forward values and gradients are compared against Qiskit
+(`Statevector` and `ReverseEstimatorGradient`) over 142 random circuits covering
+every gate (`H, X, Y, Z, rx, ry, rz, CX, CZ`); per-circuit error stays at ~1e-6.
+(The *batch-summed* gradient error on the accuracy plot climbs to ~1e-5 by 8
+qubits — that is float32 accumulation from summing 128 terms into one number,
+still ≥5 significant figures, not a modelling error.) Gates are checked for
+unitarity, the state norm is checked after every layer, and an asymmetric circuit
+pins the little-endian qubit order to Qiskit's.
+
+| | ![accuracy](benchmarks/accuracy_vs_qubits.png) | ![error distribution](benchmarks/error_distribution.png) |
+|---|---|---|
+
+**Performance.** Two honest baselines. End-to-end vs `EstimatorQNN` driven from
+Python (~100–400×), and kernel-level vs Aer's compiled statevector estimator
+(~1.7–3×, forward only) — so the win is not just deleted orchestration. Wall-time
+is shown until MLX hits the memory cliff (~22–26 qubits, single statevector).
+
+| | ![speedup](benchmarks/speedup_vs_qubits.png) | ![wall time](benchmarks/walltime_vs_qubits.png) |
+|---|---|---|
+
+**Trains identically.** Same circuit, init, data, and optimizer (SGD): the MLX
+layer (`mx.grad`) and the Qiskit QNN (`qnn.backward`) produce the same loss curve
+to ~1e-7 — same training, just faster.
+
+![training overlay](benchmarks/training_overlay.png)
+
+## Tests
+
+```bash
+uv run pytest
+```
+
+Covers gate/statevector correctness, gate unitarity and norm preservation, the
+little-endian convention, layer training, a finite-difference gradient check, and
+forward + weight-gradient + input-gradient parity with Qiskit across a random gate
+sweep (Qiskit-dependent tests skip automatically if Qiskit is absent).
+
+## Changelog
+
+See [CHANGELOG.md](CHANGELOG.md).
+
+## License
+
+MIT — see [LICENSE](LICENSE).

mlx_quantum-0.1.0/benchmarks/README.md

@@ -0,0 +1,37 @@
+# Validation & benchmarks
+
+Correctness and performance evidence for `mlx-quantum`. Scope for every
+measurement: **noiseless statevector, float32**.
+
+```bash
+uv run --extra examples python benchmarks/validate.py
+```
+
+Regenerates the five graphs in this directory. Requires the dev/`examples`
+dependencies (Qiskit, qiskit-machine-learning, qiskit-algorithms, qiskit-aer,
+matplotlib).
+
+## What each graph shows
+
+| File | Question it answers |
+|------|---------------------|
+| `accuracy_vs_qubits.png` | Forward error vs Qiskit stays at ~1e-6. The batch-summed gradient error rises to ~1e-5 by 8 qubits — pure float32 accumulation (16×8 terms summed into one number), still ≥5 significant figures; the per-circuit gradient error stays ~1e-6 (see the distribution). |
+| `speedup_vs_qubits.png` | Two honest baselines: end-to-end vs `EstimatorQNN` (forward+gradient), and kernel-level vs Aer's compiled statevector (forward only). |
+| `walltime_vs_qubits.png` | Absolute wall-time, extended until MLX hits the memory cliff. EstimatorQNN is forward+grad; Aer and MLX are forward only. |
+| `error_distribution.png` | Error histogram over a random-circuit sweep covering every gate (up to 150 draws; parameter-free draws are skipped, leaving n≈142) — the 1e-6 is not cherry-picked. |
+| `training_overlay.png` | Same circuit/init/data/optimizer trains identically via `mx.grad` and via `qnn.backward`. |
+
+## References used
+
+- **Forward:** `qiskit.quantum_info.Statevector`.
+- **Gradients:** `qiskit_algorithms.gradients.ReverseEstimatorGradient` (adjoint /
+  reverse-mode — the efficient, best-case Qiskit gradient, so the speedup claim
+  is conservative).
+- **End-to-end baseline:** `qiskit_machine_learning.neural_networks.EstimatorQNN`
+  (`forward` + `backward` driven from Python).
+- **Kernel baseline:** `qiskit_aer.primitives.EstimatorV2` (compiled statevector),
+  forward only — Aer does not expose a matching gradient, so the kernel line is a
+  forward-vs-forward comparison.
+
+The gate-level parity engine (`tests/_circuits.py`) is shared with the test suite,
+so the graphs and `pytest` check the same thing.