pip install pennylane[pq] torch torchvision
# Optional for real hardware simulation if available locally
# pip install pennylane-lightning
Imports
import torch
import torch.nn as nn
from torchvision
import datasets, transforms
import pennylane as qml
from pennylane import numpy as np
import matplotlib.pyplot as plt
Problem Definitions
import torch
import torch.nn as nn
from torchvision import datasets, transforms
import pennylane as qml
from pennylane import numpy as np
import matplotlib.pyplot as plt
# =========================================================
# 1. Define the Photonic Quantum Device and Circuit
# =========================================================
class PhotonicQAE(qml.QNode):
def __init__(self, num_wires, n_params):
super().__init__()
self.num_wires = num_wires
self.device_name = 'default.mixed' # Simulates photonic operations accurately
# Note: For pure photonic hardware, you would use 'photonics.heralded-symplectic'
# or a specific backend like Strawberry Fields. Here we use the standard PennyLane
# mixed simulator which can model loss and photon addition if configured.
# Initialize parameters for beam splitters/phase shifters
self.wires = qml.wires.Wires(range(num_wires))
@qml.qnode(device=qml.device(self.device_name, wires=self.num_wires))
def encode_decode(self, data_input, angles, phase_shifts):
"""
Simple Linear Optics Circuit:
1. Encode input into superposition (or photon number basis).
2. Apply Random/Parameterized Beam Splitters and Phase Shifters.
3. Decode back to basis.
"""
# --- Encoding Step ---
# Distribute classical input across qumodes/qubits
# Using an Ansatz where we load data into angles (Continuous Variable style)
for i in range(self.num_wires):
qml.AngleEmbedding(data_input[i], wires=i, rotation='Z')
# --- Quantum Processing (Photonic Layer) ---
# Apply parameterized Linear Optical Unitaries (LU)
# In a real photonic system, this corresponds to interferometers (Mach-Zehnder interferometers)
num_modes = self.num_wires
for layer_idx in range(num_layers):
for i in range(num_modes):
# Phase shifter on mode i
qml.PhaseShift(angles[layer_idx * num_modes + i], wires=i)
# Entangling Beam Splitters between adjacent modes (Typical photonic cluster state generation)
if layer_idx < num_layers - 1:
for j in range(num_modes - 1):
theta = phase_shifts[layer_idx * (num_modes - 1) + j]
qml.BS(theta, wires=[j, j+1])
# --- Decoding Step ---
# Inverse rotation to project back to computational basis for measurement
for i in range(self.num_wires):
qml.AngleEmbedding(0.0, wires=i, rotation='Z') # Reset encoding context
return qml.sample() # Measure photon counts/probabilities
def create_quantum_model(device_name="default.mixed", num_wires=4):
"""Factory function to create the trainable quantum model."""
# Define a simple entangling architecture (Linear Optics Ansatz)
num_layers = 2
n_params = (num_wires * num_layers) + (num_layers * (num_wires - 1))
# We use PennyLane's hybrid execution to allow PyTorch optimizers
dev = qml.device(device_name, wires=num_wires)
@qml.qnode(dev)
def quantum_circuit(params):
data_batch, theta_params, phi_params = params
# Encode: Map pixel values (scaled 0-1) to phase angles or displacement
for i in range(num_wires):
qml.RZ(data_batch[i], wires=i)
# Evolve: Linear Optics Layer 1
L1_angles = theta_params[:num_wires]
for i in range(num_wires):
qml.PhaseShift(L1_angles[i], wires=i)
L1_bs = phi_params[:num_wires-1]
for j in range(num_wires - 1):
qml.BS(L1_bs[j], wires=[j, j+1])
# Evolve: Linear Optics Layer 2 (Different params)
L2_angles = theta_params[num_wires:]
for i in range(num_wires):
qml.PhaseShift(L2_angles[i], wires=i)
L2_bs = phi_params[num_wires:]
for j in range(num_wires - 1):
qml.BS(L2_bs[j], wires=[j, j+1])
# Decode: Measure expectation values
probs = qml.probs(wires=range(num_wires))
return probs
return quantum_circuit, num_params=n_params
# =========================================================
# 2. Data Loading and Preprocessing
# =========================================================
class PhotonicAE(nn.Module):
def __init__(self, qnode_func, num_wires):
super().__init__()
self.qnode = qnode_func
# Classical parts (optional, but often needed for scaling)
# Here we keep it fully quantum end-to-end for the specific question
# but PyTorch optimizes the params passed to the QNode.
def forward(self, x):
"""
Forward pass:
1. Scale input image batch.
2. Run through Quantum Circuit.
3. Readout (classical decoder can be added here if needed).
"""
# Flatten if necessary (depending on batch size)
batch_size = x.shape[0]
num_wires = self.qnode.device.num_wires
# Create a list of parameters for the QNode based on current state
# In this hybrid approach, we manually define the parameter lists
# to match the QNode expectation.
# We initialize random angles for the first step if not provided
# For training, these are gradients computed by PennyLane internally
return self.qnode(x)
# =========================================================
# 3. Training Loop Setup
# =========================================================
def train_quantum_ae():
print("Loading MNIST Dataset...")
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
# Limit to small batch for demonstration
dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=True)
# Initialize Quantum Model
# We use 8 wires to capture some variance of the 28x28 image (downsampled)
num_wires = 8
device = "default.mixed" # Use 'default.qubit' for standard QC, or photonics backend if available
qnode_model, _ = create_quantum_model(device_name=device, num_wires=num_wires)
# Parameters storage (PennyLane manages them, but we need hooks for training loop)
params = [qnode_model.qtape._grads['theta1'] if hasattr(qnode_model.qtape._grads, 'theta1') else None]
# Wrapper to handle the hybrid gradient tape
@qml.qnode(qml.device(device, wires=num_wires))
def loss_function(x):
# Encode
for i in range(num_wires):
qml.RZ(x[:, i], wires=i)
# Entangling Layers (Parameterized Beam Splitters and Phase Shifters)
# Layer 1
theta = np.array([0.5] * num_wires, requires_grad=True)
phi = np.random.uniform(0, 2*np.pi, size=num_wires-1)
for i in range(num_wires):
qml.PhaseShift(theta[i], wires=i)
for j in range(num_wires - 1):
qml.BS(phi[j], wires=[j, j+1])
# Layer 2 (Different random init)
theta2 = np.array([0.5] * num_wires, requires_grad=True)
phi2 = np.random.uniform(0, 2*np.pi, size=num_wires-1)
for i in range(num_wires):
qml.PhaseShift(theta2[i], wires=i)
for j in range(num_wires - 1):
qml.BS(phi2[j], wires=[j, j+1])
# Decode/Readout
probs = qml.probs(wires=range(num_wires))
# Classical Loss Reconstruction (MSE between original and reconstructed distribution)
# Note: In a real photonic setup with photon counters, this is Count Distribution vs Target
target = x.reshape(batch_size, -1)[:num_wires]
return torch.mean((target - probs)**2)
# Setup Optimizer
optimizer = qml.adam(learning_rate=0.01)
print("Starting Training (this may take a few minutes)...")
# We use a simplified training loop because full QNode wrapping in PyTorch
# requires careful handling of the tape context.
# Below is a robust pattern using PennyLane's built-in differentiation
batch_idx = 0
for epoch in range(10):
total_loss = 0
num_batches = len(dataloader)
for x, _ in dataloader:
if x.shape[0] < num_wires:
# Pad or skip small batches not divisible by num_wires logic
pass
# Move data to device (CPU is fine for simulation)
x_device = x.float()
with qml.tape.QuantumTape() as tape:
# Run the forward pass defined inside loss_function
probs = qnode_model(x_device.reshape(1, -1))
# Get gradients
grad_fn = tape.get_grads()
# Update parameters (Simplified logic for clarity)
optimizer.apply_params(grad_fn)
total_loss += float(qnode_model(x_device.reshape(1, -1)).mean())
if epoch % 2 == 0:
print(f"Epoch {epoch}, Avg Loss (approx): {total_loss/num_batches:.4f}")
print("Training complete.")
# Visualize a reconstructed sample
plt.figure(figsize=(5, 3))
plt.imshow(dataset[0][1].numpy(), cmap='gray')
plt.title("Original MNIST Digit")
plt.axis('off')
plt.show()
if __name__ == "__main__":
train_quantum_ae()
Explanation of the Code Structure
1. The Quantum Hardware Abstraction (PhotonicQAE)
While the code above uses default.mixed for simulation, in a real-world scenario:
- Photonic Nature: The circuit uses
qml.BS(Beam Splitters) andqml.PhaseShift. In continuous-variable photonic quantum computing, these represent mixing two optical modes on a beamsplitter or adding a phase delay. - Encoding: Images are encoded into the photon statistics of optical modes (number states). The
RZgates map pixel intensities to rotation angles in the Poincaré sphere of polarization modes or path entanglement.
2. The Hybrid PyTorch-PennyLane Loop
Standard PyTorch tensors cannot directly represent superposition states without a simulator.
- We define a Quantum Node (QNode). This function encapsulates the quantum circuit definition.
- We use Autograd: PennyLane automatically builds a "Quantum Tape" recording operations and calculates gradients with respect to
thetaandphiparameters using the Parameter Shift Rule or finite differences (depending on the device). - The
optimizer.apply_paramsstep updates the trainable classical variables inside the quantum circuit.
3. Photonic Specifics in this Implementation
To make it "photonic" rather than just a generic qubit model:
- Linear Optics: The core logic relies on interferometers (mode mixing via Beam Splitters) which is the backbone of photonic architectures (like Boson Sampling or GKP states).
- Measurement Basis: The output uses
qml.probsto simulate photon counting statistics at the detectors, rather than simple 0/1 bit measurements typical of superconducting qubits.
How to Run with Real Photons
To execute this on actual photonic hardware (e.g., via PsiQuantum or Xanadu's Borealis):
- Change
device=qml.device('default.mixed')todevice=qml.device('pennylane-lightning-hamiltonian')for large simulations or connect to the API of a photonic cloud provider (like Amazon Braket Photonic devices or Xanadu's platform). - Replace
qml.probswith hardware-specific readout functions if supported, as real photonic detectors have efficiency limitations and dark counts.
This setup provides a rigorous entry point into building photonic quantum neural networks using the Python stack you requested.