Installation¶

Virtual Environment Setup¶

For optimal setup, we recommend using UV package manager for creating and managing virtual environments. See details at uv installation guidelines.

Make a new virtual environment using UV,

uv venv -p 3.12
source .venv/bin/activate

Replace 3.12 with your desired Python version. After activating the virtual environment, you can proceed with the installation of GWKokab.

Stable Release¶

GWKokab is available on PyPI and can be easily installed using pip. Setup virtual environment by following details in the Virtual Environment Setup section. Then, depending on your hardware configuration, run one of the following commands to install the stable release:

CPU 🐢

It is against the philosophy of GWKokab to use it with CPU support only, but if you really want to:

uv pip install -U "gwkokab[cpu]"

This option is mainly intended for development purposes. See JAX’s official documentation for the extend of supported CPU platforms and features.

Nvidia-GPU 🚀

We highly recommend using GWKokab with GPU support for better performance. Before installation, appropriate Nvidia drivers and CUDA-toolkit must be ensured on your system.

For CUDA-toolkit 12:

uv pip install -U "gwkokab[cuda12]"

or for CUDA-toolkit 13:

uv pip install -U "gwkokab[cuda13]"

See JAX’s official documentation for more details on supported Nvidia GPUs and drivers.

TPU ☄️

TPU support is available for GWKokab users.

uv pip install -U "gwkokab[tpu]"

See JAX’s official documentation for more details on using JAX with TPUs.

Nightly Build¶

For bleeding-edge features and the latest updates, you can install the nightly build of GWKokab directly from the source. This is recommended for advanced users and developers who want to stay up-to-date with the latest changes.

You must have GNU Make and UV installed on your system. Clone the repository and navigate to the project directory:

git clone https://github.com/kokabsc/gwkokab.git
cd gwkokab

Setup virtual environment by following details in the Virtual Environment Setup section. Then, depending on your hardware configuration, run one of the following commands to install the nightly build:

CPU 🐢

Linux with x86_64 is the sweet spot for CPU support with JAX.

make install PIP_FLAGS=--upgrade EXTRA=cpu

See JAX’s official documentation for the extend of supported CPU platforms and features.