Revision 66b182b735503db3a7bbc956b2fc1d30d897505b authored by Matthew Johnson on 06 December 2019, 15:44:12 UTC, committed by Matthew Johnson on 06 December 2019, 19:12:59 UTC
1 parent 1da950d
developer.rst
Building from source
====================
First, obtain the JAX source code.
.. code-block:: shell
git clone https://github.com/google/jax
cd jax
There are two steps to building JAX: building ``jaxlib`` and installing ``jax``.
If you're only modifying Python portions of JAX, you may be able to install
``jaxlib`` from pip or a prebuilt wheel and skip to installing ``jax`` from
source.
To build ``jaxlib``, you must also install some prerequisites:
* a C++ compiler (g++ or clang)
* Numpy
* Scipy
* Cython
On Ubuntu 18.04 or Debian you can install the necessary prerequisites with:
.. code-block:: shell
sudo apt-get install g++ python python3-dev python3-numpy python3-scipy cython3
If you are building on a Mac, make sure XCode and the XCode command line tools
are installed.
You can also install the necessary Python dependencies using ``pip``:
.. code-block:: shell
pip install numpy scipy cython
To build ``jaxlib`` with CUDA support, you can run
.. code-block:: shell
python build/build.py --enable_cuda
pip install -e build # installs jaxlib (includes XLA)
See ``python build/build.py --help`` for configuration options, including ways to
specify the paths to CUDA and CUDNN, which you must have installed. The build
also depends on NumPy, and a compiler toolchain corresponding to that of
Ubuntu 16.04 or newer.
To build ``jaxlib`` without CUDA GPU support (CPU only), drop the ``--enable_cuda``:
.. code-block:: shell
python build/build.py
pip install -e build # installs jaxlib (includes XLA)
Once ``jaxlib`` has been installed, you can install ``jax`` by running
.. code-block:: shell
pip install -e . # installs jax
To upgrade to the latest version from GitHub, just run ``git pull`` from the JAX
repository root, and rebuild by running ``build.py`` or upgrading ``jaxlib`` if
necessary. You shouldn't have to reinstall because ``pip install -e`` sets up
symbolic links from site-packages into the repository.
Running the tests
=================
To run all the JAX tests, we recommend using ``pytest-xdist``, which can run tests in
parallel. First, install ``pytest-xdist`` by running ``pip install pytest-xdist``.
Then, from the repository root directory run
.. code-block:: shell
pytest -n auto tests
JAX generates test cases combinatorially, and you can control the number of
cases that are generated and checked for each test (default is 10). The automated tests
currently use 25:
.. code-block:: shell
JAX_NUM_GENERATED_CASES=25 pytest -n auto tests
The automated tests also run the tests with default 64-bit floats and ints:
.. code-block:: shell
JAX_ENABLE_X64=1 JAX_NUM_GENERATED_CASES=25 pytest -n auto tests
You can run a more specific set of tests using
`pytest <https://docs.pytest.org/en/latest/usage.html#specifying-tests-selecting-tests>`_'s
built-in selection mechanisms, or alternatively you can run a specific test
file directly to see more detailed information about the cases being run:
.. code-block:: shell
python tests/lax_numpy_test.py --num_generated_cases=5
The Colab notebooks are tested for errors as part of the documentation build.
Update documentation
====================
To rebuild the documentation, install several packages:
.. code-block:: shell
pip install -r docs/requirements.txt
You must also install ``pandoc`` in order to regenerate the notebooks.
See `Install Pandoc <https://pandoc.org/installing.html>`_. On Mac, I had success with
the miniconda installer, then ``conda install -c conda-forge pandoc``.
If you do not want to install ``pandoc`` then you should regenerate the documentation
without the notebooks.
You run at top-level one of the following commands:
.. code-block:: shell
sphinx-build -b html docs docs/build/html # with the notebooks
sphinx-build -b html -D nbsphinx_execute=never docs docs/build/html # without the notebooks
You can then see the generated documentation in
``docs/build/html/index.html``.
Update notebooks
----------------
Open the notebook with http://colab.research.google.com, update it, ``Run all cells`` then
``Download ipynb``. You may want to test that it executes properly, using ``sphinx-build`` as
explained above.
Documentation building on readthedocs.io
----------------------------------------
JAX's auto-generated documentations is at `jax.readthedocs.io <https://jax.readthedocs.io/>`_.
The documentation building is controlled for the entire project by the
`readthedocs JAX settings <https://readthedocs.org/dashboard/jax>`_. The current settings
trigger a documentation build as soon as code is pushed to the GitHub ``master`` branch.
For each code version, the building process is driven by the
``.readthedocs.yml`` and the ``docs/conf.py`` configuration files.
For each automated documentation build you can see the
`documentation build logs <https://readthedocs.org/projects/jax/builds/>`_.
If you want to test the documentation generation on Readthedocs, you can push code to the ``test-docs``
branch. That branch is also built automatically, and you can
see the generated documentation `here <https://jax.readthedocs.io/en/test-docs/>`_.
For a local test, I was able to do it in a fresh directory by replaying the commands
I saw in the Readthedocs logs:
.. code-block:: shell
mkvirtualenv jax-docs # A new virtualenv
mkdir jax-docs # A new directory
cd jax-docs
git clone --no-single-branch --depth 50 https://github.com/google/jax
cd jax
git checkout --force origin/test-docs
git clean -d -f -f
python -m pip install --upgrade --no-cache-dir pip
python -m pip install --upgrade --no-cache-dir -I Pygments==2.3.1 setuptools==41.0.1 docutils==0.14 mock==1.0.1 pillow==5.4.1 alabaster>=0.7,<0.8,!=0.7.5 commonmark==0.8.1 recommonmark==0.5.0 'sphinx<2' 'sphinx-rtd-theme<0.5' 'readthedocs-sphinx-ext<1.1'
python -m pip install --exists-action=w --no-cache-dir -r docs/requirements.txt
python `which sphinx-build` -T -E -b html -d _build/doctrees-readthedocs -D language=en . _build/html
![swh spinner](/static/img/swh-spinner.gif)
Computing file changes ...