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 https://github.com/google/objax.git
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

# If you have CUDA installed, specify your installed CUDA version.
CUDA_VERSION=11.0
pip install -f https://storage.googleapis.com/jax-releases/jax_releases.html jaxlib==`python3 -c 'import jaxlib; print(jaxlib.__version__)'`+cuda`echo $CUDA_VERSION | sed s:\\\.::g`

Running tests and linter

Run linter:

./tests/run_linter.sh

Run tests:

./tests/run_tests.sh

Running a single test:

CUDA_VISIBLE_DEVICES= python3 -m unittest tests/jit.py