mlx-quantum 0.1.0__py3-none-any.whl

mlx_quantum/__init__.py

@@ -0,0 +1,38 @@
+from . import statevector
+from .statevector import (
+    zero_state,
+    apply_1q,
+    apply_2q,
+    expval_z,
+    expval_all_z,
+    rx,
+    ry,
+    rz,
+    H,
+    X,
+    Y,
+    Z,
+    CX,
+    CZ,
+)
+from .layer import QuantumLayer
+
+__all__ = [
+    "QuantumLayer",
+    "statevector",
+    "zero_state",
+    "apply_1q",
+    "apply_2q",
+    "expval_z",
+    "expval_all_z",
+    "rx",
+    "ry",
+    "rz",
+    "H",
+    "X",
+    "Y",
+    "Z",
+    "CX",
+    "CZ",
+]
+__version__ = "0.1.0"

mlx_quantum/layer.py

@@ -0,0 +1,50 @@
+import numpy as np
+import mlx.core as mx
+import mlx.nn as nn
+
+from .statevector import zero_state, apply_1q, apply_2q, expval_all_z, H, ry, CX
+
+
+class QuantumLayer(nn.Module):
+    """A trainable quantum layer that behaves like any other ``mlx.nn`` module.
+
+    Runs a hardware-efficient ansatz natively in MLX, so its weights train by
+    ordinary gradient descent on the Metal GPU — no Qiskit, no custom autodiff.
+
+        H on all qubits
+        RY(input[q]) on each qubit q          (angle encoding)
+        reps x [ RY(weight) on each qubit; CX ring ]
+
+    Call with ``x`` of shape ``(num_qubits,)`` or ``(batch, num_qubits)``; returns
+    <Z> per qubit, shape ``(num_qubits,)`` or ``(batch, num_qubits)``.
+    """
+
+    def __init__(self, num_qubits: int, reps: int = 1, initial_weights: np.ndarray | None = None):
+        super().__init__()
+        self.num_qubits = num_qubits
+        self.reps = reps
+        if initial_weights is None:
+            initial_weights = np.random.uniform(-1, 1, reps * num_qubits)
+        self.weight = mx.array(np.asarray(initial_weights, np.float32))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        was_1d = x.ndim == 1
+        if was_1d:
+            x = x[None]
+
+        n = self.num_qubits
+        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)
+
+        w = self.weight.reshape(self.reps, n)
+        for r in range(self.reps):
+            for q in range(n):
+                state = apply_1q(state, ry(w[r, q]), q)
+            for q in range(n - 1):
+                state = apply_2q(state, CX, q, q + 1)
+
+        out = expval_all_z(state)
+        return out[0] if was_1d else out

mlx_quantum/statevector.py

@@ -0,0 +1,114 @@
+"""Batched, differentiable statevector simulation in pure MLX.
+
+A statevector is a complex ``mx.array`` of shape ``(batch,) + (2,) * num_qubits``.
+Qubit ordering is little-endian to match Qiskit: qubit 0 is the fastest-varying
+index, so flattening the state reproduces Qiskit's amplitude order exactly. Gates
+are applied as ``einsum`` contractions, so the whole simulation runs on the Metal
+GPU and is differentiable end to end via ``mx.grad`` — no custom vjp.
+"""
+
+import math
+import string
+
+import numpy as np
+import mlx.core as mx
+
+_AXES = string.ascii_lowercase
+
+
+def zero_state(batch: int, num_qubits: int) -> mx.array:
+    """The |0...0> state, broadcast over a batch."""
+    amp = np.zeros((1,) + (2,) * num_qubits, np.complex64)
+    amp.flat[0] = 1.0
+    return mx.array(np.broadcast_to(amp, (batch,) + (2,) * num_qubits).copy())
+
+
+def _axis(qubit: int, n: int) -> int:
+    """Little-endian: qubit 0 is the last axis (fastest index when flattened)."""
+    return n - 1 - qubit
+
+
+def apply_1q(state: mx.array, gate: mx.array, qubit: int) -> mx.array:
+    """Apply a single-qubit gate. ``gate`` is ``(2, 2)`` or, per-sample, ``(batch, 2, 2)``."""
+    n = state.ndim - 1
+    axes = _AXES[:n]
+    q = _axis(qubit, n)
+    src, tgt = "B" + axes, "B" + axes[:q] + "Z" + axes[q + 1:]
+    if gate.ndim == 2:
+        return mx.einsum(f"{src},Z{axes[q]}->{tgt}", state, gate)
+    return mx.einsum(f"{src},BZ{axes[q]}->{tgt}", state, gate)
+
+
+def apply_2q(state: mx.array, gate: mx.array, q0: int, q1: int) -> mx.array:
+    """Apply a two-qubit gate given as a ``(2, 2, 2, 2)`` tensor ``[out0, out1, in0, in1]``."""
+    n = state.ndim - 1
+    axes = _AXES[:n]
+    a0, a1 = _axis(q0, n), _axis(q1, n)
+    tgt = list("B" + axes)
+    tgt[1 + a0], tgt[1 + a1] = "Y", "Z"
+    return mx.einsum(f"B{axes},YZ{axes[a0]}{axes[a1]}->{''.join(tgt)}", state, gate)
+
+
+def expval_z(state: mx.array, qubit: int) -> mx.array:
+    """<Z> on one qubit, shape ``(batch,)``."""
+    probs = mx.abs(state) ** 2
+    keep = _axis(qubit, state.ndim - 1)
+    axes = [1 + j for j in range(state.ndim - 1) if j != keep]
+    marg = mx.sum(probs, axis=axes) if axes else probs
+    return marg[:, 0] - marg[:, 1]
+
+
+def expval_all_z(state: mx.array) -> mx.array:
+    """<Z> on every qubit, shape ``(batch, num_qubits)``."""
+    return mx.stack([expval_z(state, q) for q in range(state.ndim - 1)], axis=1)
+
+
+def _const(matrix) -> mx.array:
+    return mx.array(np.array(matrix, np.complex64))
+
+
+H = _const([[1, 1], [1, -1]]) / math.sqrt(2)
+X = _const([[0, 1], [1, 0]])
+Y = _const([[0, -1j], [1j, 0]])
+Z = _const([[1, 0], [0, -1]])
+
+
+def _two_qubit(fn) -> mx.array:
+    g = np.zeros((2, 2, 2, 2), np.complex64)
+    for c in (0, 1):
+        for t in (0, 1):
+            oc, ot, val = fn(c, t)
+            g[oc, ot, c, t] = val
+    return mx.array(g)
+
+
+CX = _two_qubit(lambda c, t: (c, t ^ c, 1.0))
+CZ = _two_qubit(lambda c, t: (c, t, -1.0 if c and t else 1.0))
+
+
+def rx(theta: mx.array) -> mx.array:
+    c, s = mx.cos(theta / 2), mx.sin(theta / 2)
+    return _rot(c, -1j * s, -1j * s, c)
+
+
+def ry(theta: mx.array) -> mx.array:
+    c, s = mx.cos(theta / 2), mx.sin(theta / 2)
+    return _rot(c, -s, s, c)
+
+
+def rz(theta: mx.array) -> mx.array:
+    p = _expi(theta / 2)
+    zero = mx.zeros_like(theta).astype(mx.complex64)
+    return _rot(mx.conj(p), zero, zero, p)
+
+
+def _expi(phase: mx.array) -> mx.array:
+    return (mx.cos(phase) + 1j * mx.sin(phase)).astype(mx.complex64)
+
+
+def _rot(a, b, c, d) -> mx.array:
+    """Stack four (scalar or batched) entries into a ``(*shape, 2, 2)`` gate."""
+    a, b, c, d = (mx.array(v).astype(mx.complex64) for v in (a, b, c, d))
+    row0 = mx.stack([a, b], axis=-1)
+    row1 = mx.stack([c, d], axis=-1)
+    return mx.stack([row0, row1], axis=-2)

mlx_quantum-0.1.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,8 @@
+mlx_quantum/__init__.py,sha256=7qjKB4AZ1N3xiUX43BN4iilxspCnPB46KJnJBEDTs6A,479
+mlx_quantum/layer.py,sha256=eSdMsA3vJNDAHTww4AP1yztWPJX7eTS67EA6Q7Fvu9I,1767
+mlx_quantum/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mlx_quantum/statevector.py,sha256=-4zG5KVPFXku6hO0iofgOP-0ULsjgJoVhpCKsS7PGuw,3743
+mlx_quantum-0.1.0.dist-info/METADATA,sha256=2bcRDqPEtIu9GIEz5Krt069FYxNO43HFeEn5HVIQtiw,6701
+mlx_quantum-0.1.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+mlx_quantum-0.1.0.dist-info/licenses/LICENSE,sha256=CthJDHR37kFu4ndMrsrTTwdJaBqDDyB-OyIn4AteLxM,1067
+mlx_quantum-0.1.0.dist-info/RECORD,,

mlx_quantum-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

mlx_quantum-0.1.0.dist-info/licenses/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.