Quax
A high-performance quantum information science library built on top of JAX.
Documentation
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 jax.numpy as jnp
import quax as qx
# Create a 2-qubit state
psi_0: qx.StateVector = qx.zero_state_vector(dims=(2, 2,))
# Apply gates to prepare a pure Bell state
psi = (qx.gates.H | qx.gates.I) @ psi_0
psi = qx.gates.CNOT @ psi
# Generate an ensemble of Lindbladians with varying strenght
L: qx.Lindbladian = qx.lindbladians.amplitude_damping(jnp.linspace(0.0, 0.1, 10))
# Combine the Lindbladians with the CNOT unitary to create an ensemble of quantum channels
# The unitary is automatically promoted to a superoperator to support the addition of the Lindbladian
noisy_cnot: qx.SuperOp = qx.gates.CNOT + (L | L)
# Apply the noise operation to create a noisy Bell states
rho: qx.DensityMatrix = noisy_cnot @ (qx.gates.H | qx.gates.I) @ psi_0
# Compute the fidelity between the noisy Bell state and the ideal Bell state
fidelity: jax.Array = qx.fidelity(rho, psi)
# Now let's apply a qutrit gate to the ensemble of noisy Bell states.
rho = (qx.gates.TRX12(jnp.pi) | qx.gates.I) @ rho
# We can estimate the 𝜆8 observable
obesrvables: jax.Array = qx.estimate(rho, observable=(qx.gates.GELLMANN8) | qx.gates.I)
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.
1 maintainer
Metadata
1 maintainer
Release files for rigetti-quax 0.7.4
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| rigetti_quax-0.7.4.tar.gz | 96.3 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| rigetti_quax-0.7.4-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 211.5 kB
Release files / rigetti_quax-0.7.4.tar.gz
| Download URL | rigetti_quax-0.7.4.tar.gz |
|---|---|
| Size | 96.3 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
4750c2686fd1accd8bc4297d088aa32a51ef531f89cde5cadc2905d2fef05fc8
|
|
BLAKE2b-256 checksum How to use checksums |
bf0f462aa8068b2cd8587ef6322d6a17e4923b44f99f88153f62f1e8e8f8936f
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
poetry/2.4.3 CPython/3.12.14 Linux/6.1.186-228.374.amzn2023.x86_64
|
Release files / rigetti_quax-0.7.4-py3-none-any.whl
| Download URL | rigetti_quax-0.7.4-py3-none-any.whl |
|---|---|
| Size | 115.2 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
ab106f83269ba4c60b0820b0ac3dbcd887a3b3a31864c3d390e0e87eb5bb486d
|
|
BLAKE2b-256 checksum How to use checksums |
2ec5a9d86900d1859967ce5b61feec88e64d3530cc6a2922ee2ef31f5b0b0fc3
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
poetry/2.4.3 CPython/3.12.14 Linux/6.1.186-228.374.amzn2023.x86_64
|
