Metadata-Version: 2.4
Name: triton-qml
Version: 0.1.0
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
[![PyPI version](https://badge.fury.io/py/triton-qml.svg)](https://badge.fury.io/py/triton-qml)
[![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
You can easily install our library using `pip`:
```bash
pip install triton-qml
```
You can also install it from source:
```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 = 4
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)

# applying Fused CZ gates
# applying CZ(0,1), CZ(1,2), CZ(2,3), and CZ(3,1)
state_vec = simulator.apply_fused_cz_gate(state_vec, [[0,1],[1,2],[2,3],[3,1]])

# --- Fused 1q and 2q Pauli rotation gates ---(up to 4qubit wires)
state_vec = simulator.apply_1q2q_pauli_rotation_gate_sequence(
    state_vec, 
    [
        ['rx', 0, theta[3]],
        ['ry', 1, theta[4]],
        ['rz', 2, theta[5]],
        ['rxx', [0,1], theta[6]],
        ['ryy', [2,3], theta[7]],
        ['rzz', [1,2], theta[8]]
    ]
)

# --- Fused gate ---(up to 4qubit wires)
# this gate can use
# rx, ry, rz, 
# rxx, ryy, rzz, 
# crx, cry, crz, cp, 
# x, y, z, s, sdag, t, tdag, h, sqrtx, sqrtxdag, 
# cnot(cx), cz, swap,
# ccx
# this function often takes longer compile time than "apply_1q2q_pauli_rotation_gate_sequence", 
# especially during the backward path. 
state_vec = simulator.apply_fused_gates(
    state_vec, 
    [
        ["rx", 0, theta[9]],
        ["ryy", [0,1], theta[10]], 
        ["crz", [2,3], theta[11]], 
        ["x", 3], 
        ["cnot", [1,2]],
        ["ccx", [3,2,1]]
    ]
)

# fused cnot gate is also available 
# applying CX(0,1), CX(1,2), CX(2,3), and CX(3,1)
state_vec = simulator.apply_fused_cnot_gate(state_vec, [[0,1],[1,2],[2,3],[3,1]])

# 3. Calculate Expectation Value (Loss)
obs = "XYZI" # X_0 Y_1 Z_2 I_3
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.
