Development setup

This section describes some basic setup to start developing and extending Objax.

Environment setup

First of all you need to install all necessary dependencies. We recommend to setup a separate virtualenv to work on Objax, it could be done with following commands on Ubuntu or similar Linux distribution:

# Install virtualenv if you haven't done so already
sudo apt install python3-dev python3-virtualenv python3-tk imagemagick virtualenv pandoc
# Create a virtual environment (for example ~/.venv/objax, you can use your name here)
virtualenv -p python3 --system-site-packages ~/.venv/objax
# Start the virtual environment
. ~/.venv/objax/bin/activate

# Clone objax git repository, if you haven't.
git clone
cd objax

# Install python dependencies.
pip install --upgrade -r requirements.txt
pip install --upgrade -r tests/requirements.txt
pip install --upgrade -r docs/requirements.txt
pip install --upgrade -r examples/requirements.txt
pip install flake8

# jaxlib releases require CUDA 11.2 or newer
JAX_VERSION=`python3 -c 'import jax; print(jax.__version__)'`
pip uninstall -y jaxlib
pip install -f $RELEASE_URL jax[cuda]==$JAX_VERSION

Running tests and linter

Run linter:


Run tests:


Running a single test:

CUDA_VISIBLE_DEVICES= python3 -m unittest tests/