Quax

A high-performance quantum information science library built on top of JAX.

Documentation

Check out the docs.

Features

• Quantum Objects: States, gates and superoperators objects are defined to allow natural manipulations and operations.
• Clear typing: Standard objects such as Unitaries, Chois and Density matrices are all typed, clarifying the nature of various objects.
• Standard operators composition/application @, tensor products |, scalar multiplication * and powers * are defined on all quantum objects.
• Automatic promotion Pure states and operators are automatically promoted to mixed states and superoperators when appropriate.
• Quantum instruments Mid-circuit measurements are modeled with QuantumInstrument, supporting confusion/transition matrices, composition, and tensor products.
• Qudit support Operations on d-dimensional qudits are supported.
• Batch support Operating on batches or ensembles of states is supported for straightforward parallelization.

Installation

pip install rigetti-quax

Quick Example

import jax
import quax as qx

# Create a quantum state
state = qx.zero_state_vector(dims=(2,))

# Apply a unitary operation
U = qx.random_unitary(dims=(2,), key=jax.random.key(0))
final_state = U @ state

# Convert between representations
choi = qx.unitary_to_choi(U)
pauli_liouville = qx.to_pauli_liouville(choi)

Acknowledgements

Quax draws inspiration, educational material, and some code from forest-benchmarking, Rigetti's open-source library for quantum characterization, verification, and validation. We gratefully acknowledge the forest-benchmarking contributors for their foundational work on superoperator representations, quantum channel conventions, and distance metrics.

License

Copyright 2026 Rigetti & Co, LLC. Licensed under Apache License 2.0.