triton-qml 0.0.1

Fast and memory-efficient classical simulation of QML via Triton

Triton-QML: Fast and Memory-Efficient Classical Simulation of Quantum Machine Learning

This repository provides the official implementation of the paper:
Fast and memory-efficient classical simulation of quantum machine learning via forward and backward gate fusion

If you are looking for the codes used in the paper, please switch to the paper branch

This master branch is under development as a general-purpose quantum circuit simulator, with a strong focus on accelerating Quantum Machine Learning and Variational Quantum Algorithms.

Installation

git clone https://github.com/puyokw/triton_qml.git
cd triton_qml
pip install . 

Usage

import torch
from triton_qml.simulator import FunctionalQuantumSimulator

# Setup device and simulator
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
n_qubits = 3
batch_size = 100

# Initialize the Triton-based simulator (with memory saving mode enabled)
simulator = FunctionalQuantumSimulator( n_qubits, device=device)

# 1. Prepare initial state |0...0>
state_vec = simulator.create_initial_state(batch_size)
# |0...0> -> |+...+>
for i in range(n_qubits):
    state_vec = simulator.apply_H_gate(state_vec, i)

data = torch.randn(batch_size, 6, device=device)
theta = torch.randn(18, requires_grad=True, device=device)

# --- Standard single-qubit gates ---
state_vec = simulator.apply_Rx_gate(state_vec, theta[0], target_qubit=0)
state_vec = simulator.apply_Ry_gate(state_vec, theta[1], target_qubit=1)
state_vec = simulator.apply_Rz_gate(state_vec, theta[2], target_qubit=2)

# --- Fused CZ gate ---
controls =[i for i in range(n_qubits)]
targets =[(i+1) % n_qubits for i in range(n_qubits)]
state_vec = simulator.apply_fused_cz_gate(state_vec, controls, targets)

# --- Fused single-qubit gate ---
state_vec = simulator.apply_single_qubits_gate_sequence(state_vec, 0, ['rx', 'ry', 'rz', 'rx'], [theta[3], theta[4], theta[5], data[:,0]])
# --- Fused 2-qubit block ---
state_vec = simulator.apply_fused_two_qubit_block(
    state_vec,  1, ["rx", "ry", "rz", 'ry'], [theta[6], theta[7], theta[8], data[:,1]], 
                2, ["rx", "ry", "rz", 'rz'], [theta[9], theta[10], theta[11], data[:,2]]
)

# fused cnot gate is also available 
state_vec = simulator.apply_fused_cnot_gate(state_vec, controls, targets)

# --- Fused 3-qubit block ---
# you can also encode data
state_vec = simulator.apply_fused_three_qubit_block( 
    state_vec,  0, ["rx", "ry", "rz", "rx"], [theta[12], data[:,0], theta[13], data[:,3]], 
                1, ["rx", "ry", "rz", "ry", "rx"], [theta[14], data[:,1], theta[15], data[:,4], data[:,5]],
                2, ["rx", "ry", "rz"], [theta[16], data[:,2], theta[17]]
)

# 3. Calculate Expectation Value (Loss)
obs = "XYZ" # X_0 Y_1 Z_2
exp_val = simulator.calc_exp_val(state_vec, observable=obs)
loss = exp_val.sum()

loss.backward()
print("Gradient w.r.t theta:", theta.grad)

Citation

If you find this code useful, please consider citing our paper:

@article{kawase2026fast,
  title={Fast and memory-efficient classical simulation of quantum machine learning via forward and backward gate fusion},
  author={Yoshiaki Kawase},
  journal={arXiv preprint arXiv:2603.02804},
  year={2026}
}

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.