Metadata-Version: 2.4
Name: triton-qml
Version: 0.0.1
Summary: Fast and memory-efficient classical simulation of QML via Triton
Author-email: Yoshiaki Kawase <ykawase@g.ecc.u-tokyo.ac.jp>
Project-URL: Homepage, https://github.com/puyokw/triton_qml
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1.0
Requires-Dist: triton>=2.1.0
Provides-Extra: dev
Requires-Dist: parameterized; extra == "dev"
Requires-Dist: pandas; extra == "dev"
Dynamic: license-file

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

[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
[![PyTorch](https://img.shields.io/badge/PyTorch-2.1+-ee4c2c.svg)](https://pytorch.org/)
[![Triton](https://img.shields.io/badge/Triton-optimized-blueviolet.svg)](https://github.com/openai/triton)

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**](https://arxiv.org/abs/2603.02804)

If you are looking for the codes used in the paper, please switch to the [**paper branch**](https://github.com/puyokw/triton_qml/tree/paper)

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
```bash
git clone https://github.com/puyokw/triton_qml.git
cd triton_qml
pip install . 
```

## Usage

```python
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:

```bibtex
@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](LICENSE) file for details.
