Machine Learning Integration

Use Superfermion circuits as differentiable layers in PyTorch, Flax/JAX, and TensorFlow.

Superfermion provides quantum layers for all three major ML frameworks. Each is a thin adapter (~20 lines) that wires sf.State.grad() into the framework's autograd mechanism.

Architecture

Framework autograd (PyTorch / JAX / TF)
         │
         ▼
  QuantumLayer wrapper (~20 lines)
         │
         ▼
  sf.simulate() + sf.State.grad()  (Rust)

The quantum layer:

Forward pass — calls sf.simulate() + state.expectation()
Backward pass — calls state.grad() (adjoint or parameter-shift)

Flax / JAX

from superfermion.nn.quantum_layer import QuantumLayer
import jax
import jax.numpy as jnp
import superfermion as sf

# Build a parameterized circuit
theta = sf.param("theta")
phi = sf.param("phi")
circuit = sf.Circuit(2).ry(theta, 0).ry(phi, 1).cnot(0, 1).rz(theta, 1)

# Define observable (pass a PauliString — Flax converts plain lists to tuples)
obs = sf.PauliString("ZZ", coeff=1.0)

# Create layer
layer = QuantumLayer(circuit, obs, device="cpu")

# Initialize
key = jax.random.PRNGKey(0)
params = layer.init(key)

# Forward pass (custom_vjp routes backward to sf.State.grad)
output = layer.apply(params)
print(f"Output: {output}")

# Gradient via JAX autograd
grad_fn = jax.grad(lambda p: layer.apply(p).sum())
grads = grad_fn(params)
print(f"Gradients: {grads}")

JAX integration uses jax.custom_vjp — the backward pass calls sf.State.grad() in Rust for exact parameter-shift or adjoint gradients.

PyTorch

from superfermion.nn.torch_layer import TorchQuantumLayer
import torch
import superfermion as sf

# Build circuit and observable
circuit = sf.Circuit(2).ry(sf.param("t0"), 0).ry(sf.param("t1"), 1).cnot(0, 1)
obs = [([3, 3], 1.0, 0.0)]  # ZZ

# Create layer — parameters live inside the module (layer.weights)
layer = TorchQuantumLayer(circuit, obs, device="cpu")

# Forward pass
output = layer()
print(f"Output: {output.item():.6f}")

# Backward pass (sf.State.grad via autograd.Function)
output.backward()
print(f"Gradients: {layer.weights.grad}")

PyTorch integration uses torch.autograd.Function. The backward pass calls sf.State.grad() for exact gradients.

TensorFlow

from superfermion.nn.tf_layer import TFQuantumLayer
import tensorflow as tf
import superfermion as sf

# Build circuit and observable
circuit = sf.Circuit(2).ry(sf.param("t0"), 0).ry(sf.param("t1"), 1).cnot(0, 1)
obs = [([3, 3], 1.0, 0.0)]  # ZZ

# Create layer — weights live inside the layer (layer.weights_var)
layer = TFQuantumLayer(circuit, obs, device="cpu")

# Forward pass
output = layer()
print(f"Output: {output.numpy():.6f}")

# Backward pass (tf.custom_gradient routes to sf.State.grad)
with tf.GradientTape() as tape:
    loss = layer()

grads = tape.gradient(loss, layer.weights_var)
print(f"Gradients: {grads.numpy()}")

TensorFlow integration uses tf.custom_gradient — the backward pass calls sf.State.grad() for exact gradients.

Data Embedding (`feature_names`)

All three layers accept a feature_names list naming which circuit parameters are fed from the input (x / inputs) instead of being trained. The remaining parameters become trainable weights, and gradients flow to both the input and the weights.

import torch
import superfermion as sf
from superfermion.nn.torch_layer import TorchQuantumLayer

x0, w0 = sf.param("x0"), sf.param("w0")
circuit = sf.Circuit(1).rx(x0, 0).ry(w0, 0)

layer = TorchQuantumLayer(circuit, [([3], 1.0, 0.0)], feature_names=["x0"])
out = layer(torch.tensor([0.5]))   # x0 <- 0.5; w0 stays the trainable weight

feature_names defaults to None — the legacy behavior where every parameter is a weight and the input is ignored.

Batch Evaluation

All three layers support batched parameter evaluation:

import jax
import torch
import superfermion as sf
from superfermion.nn.quantum_layer import QuantumLayer
from superfermion.nn.torch_layer import TorchQuantumLayer

circuit = sf.Circuit(2).ry(sf.param("t0"), 0).ry(sf.param("t1"), 1).cnot(0, 1)
obs = sf.PauliString("ZZ", coeff=1.0)

# JAX/Flax: evaluate a batch of weight vectors
flax_layer = QuantumLayer(circuit, obs, device="cpu")
params = flax_layer.init(jax.random.PRNGKey(0))
weights_batch = [params["params"]["weights"], params["params"]["weights"] * 2]
outputs = [flax_layer.apply({"params": {"weights": w}}) for w in weights_batch]
print(f"JAX batch outputs: {outputs}")

# PyTorch: stack repeated evaluations (weights live inside the module)
torch_layer = TorchQuantumLayer(circuit, obs, device="cpu")
outputs = torch.stack([torch_layer() for _ in range(4)])
print(f"PyTorch batch outputs: {outputs}")

Gradient Methods

The gradient engine is selected via the layer or directly:

Method	Mechanism	Speed	Best For
Adjoint	`state.grad()` with adjoint flag	O(1) vs params	Deep circuits, many params
Parameter-shift	`state.grad()` with shift rules	O(N)	Shallow circuits
SPSA	Stochastic perturbation	O(1) stochastic	Noisy, many params
QNG	Quantum Fisher metric	O(N²)	Fast convergence

Custom Training Loop

For full control, bypass the layer wrappers:

import superfermion as sf
import numpy as np

ansatz = sf.Circuit(2).ry(sf.param("t0"), 0).ry(sf.param("t1"), 1).cnot(0, 1)
obs = [([3, 3], 1.0, 0.0)]

params = {"t0": 0.5, "t1": 1.2}
lr = 0.1

for step in range(100):
    state = sf.simulate(ansatz, params=params)
    energy = state.expectation(obs)

    dag = ansatz.to_ir()  # keep parameters symbolic — .bind() would erase them
    grads = state.grad(obs, dag, params)

    for k in params:
        params[k] -= lr * grads.get(k, 0.0)

    if step % 20 == 0:
        print(f"Step {step}: energy = {energy:.6f}")

On this page