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

Metadata

Release files for rigetti-quax 0.7.3

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for rigetti-quax 0.7.3
File	Size	Uploaded
rigetti_quax-0.7.3.tar.gz	95.3 kB	Sep 18, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for rigetti-quax 0.7.3
File	Interpreter	ABI	Platform	Reset
rigetti_quax-0.7.3-py3-none-any.whl 113.6 kB Sep 18, 2026	Python 3	none	any	Details

Total release size: 208.9 kB

Release files / rigetti_quax-0.7.3.tar.gz

Download URL	rigetti_quax-0.7.3.tar.gz
Size	95.3 kB
Tags	Source
SHA-256 checksum How to use checksums	`c0f45b9cdf1984b1a9c6f8946649dbe32668720e123c4ffd4e8f6d1742d47f06`
BLAKE2b-256 checksum How to use checksums	`d38ee4ad05d8b45776597d2757b01195fe88e8788605d27df0b7dc3ad2c4f5c1`
Upload date	Sep 18, 2026
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.3-py3-none-any.whl

Download URL	rigetti_quax-0.7.3-py3-none-any.whl
Size	113.6 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`55d54620034adf3577eeaa513fe773868889fea5e780d3ffcbd839c315a5143a`
BLAKE2b-256 checksum How to use checksums	`33daea49fdc64e5e084f93354339bc6b4444a0033ee8aa8bd4b0cf22f913c5a8`
Upload date	Sep 18, 2026
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`