bqskit-qfactor-jax is a Python package that implements circuit instantiation using the QFactor and QFactor-Sample algorithms on GPUs to accelerate BQSKit. It uses JAX as an abstraction layer of the GPUs, seamlessly utilizing JIT compilation and GPU parallelism.

bqskit-qfactor-jax is available for Python 3.8+ on Linux.

First, install JAX with GPU support, you may refer to JAX's installation instructions. For users working on Perlmutter please use the following modules before installing JAX in your environment:

module load cudnn/8.9.3_cuda12
module load nccl/2.18.3-cu12

Next, install this package with pip:

pip install bqskit-qfactor-jax

Please set the environment variable XLA_PYTHON_CLIENT_PREALLOCATE=False when using this package. Also, if you encounter OOM issues consider setting XLA_PYTHON_CLIENT_ALLOCATOR=platform.

Please look at the examples for basic usage, especially at performance comparison between QFactor and QFactor-Sample.

When using several workers on the same GPU, we recommend using Nvidia's MPS. You may initiate it using the command line

nvidia-cuda-mps-control -d

You can disable it by running this command line:

echo quit | nvidia-cuda-mps-control

If you are using QFactor-JAX please cite:
Kukliansky, Alon, et al. "QFactor: A Domain-Specific Optimizer for Quantum Circuit Instantiation." 2023 IEEE International Conference on Quantum Computing and Engineering (QCE). Vol. 1. IEEE, 2023. Link.

If you are using QFactor-Sample please cite:
Kukliansky, Alon, et al. "Leveraging Quantum Machine Learning Generalization to Significantly Speed-up Quantum Compilation" arXiv preprint arXiv:2405.12866 (2024).

The software in this repository is licensed under a BSD free software license and can be used in source or binary form for any purpose as long as the simple licensing requirements are followed. See the LICENSE file for more information.

Quantum Fast Circuit Optimizer (QFactor) JAX implementation Copyright (c) 2024, U.S. Federal Government and the Government of Israel. All rights reserved.