2. GPU Acceleration

PyBEST can offload selected tensor contractions to an NVIDIA GPU using either CuPy or PyTorch. Both libraries must be installed separately — they are not included in the default PyBEST installation.

If you use GPU acceleration in PyBEST, please cite [kriebel2024] and [dobrowolska2026].

Note

GPU acceleration applies to tensordot-based contractions only. All other contractions are still executed on the CPU. If the arrays do not fit in GPU memory despite automatic batching, PyBEST falls back to the CPU.

2.1. Environment variables

All GPU settings are configured via environment variables before starting Python.

2.1.1. Selecting a backend

To enable a backend, set the corresponding variable to a non-empty value. To disable it, the variable must be empty or unset.

# Enable CuPy backend
export PYBEST_CUPY_AVAIL=1

# Disable CuPy backend (variable must be empty)
export PYBEST_CUPY_AVAIL=

# Enable PyTorch backend
export PYBEST_PYTORCH_AVAIL=1

# Disable PyTorch backend (variable must be empty)
export PYBEST_PYTORCH_AVAIL=

Both variables can be set at the same time; CuPy will be preferred when both are available.

2.1.2. Selecting a batching strategy

For Cholesky-based contractions that do not fit in GPU memory, PyBEST splits the arrays into smaller batches. Two strategies are available. Exactly one variable must be set to a non-empty value; the other must be empty.

# C-splitting: batch along the Cholesky index (default)
export PYBEST_C_SPLITTING=1
export PYBEST_X_SPLITTING=

# X-splitting: batch along the auxiliary (x) index
export PYBEST_C_SPLITTING=
export PYBEST_X_SPLITTING=1

If neither variable is set, or both are set simultaneously, PyBEST defaults to C-splitting.