jaxlib
, the C++ support library for jax
.jax
Python package.jaxlib
¶jaxlib
with pip¶jaxlib
from a prebuilt wheel using pip:jaxlib
from source¶jaxlib
from source, you must also install some prerequisites:numpy
, scipy
, six
, wheel
.six
package is required for during the jaxlib build only, and is notrequired at install time.pip
:jaxlib
with CUDA support, you can run:pythonbuild/build.py--help
for configuration options, including ways tospecify the paths to CUDA and CUDNN, which you must have installed. Herepython
should be the name of your Python 3 interpreter; on some systems, youmay need to use python3
instead. By default, the wheel is written to thedist/
subdirectory of the current directory.jaxlib
without CUDA GPU support (CPU only), drop the --enable_cuda
:jaxlib
from source on Windows¶bazel
, patch
and realpath
areaccessible. Activate the conda environment. The following command builds withCUDA enabled, adjust it to whatever suitable for you:--bazel_options='--copt=/Z7'
.jax
¶jaxlib
has been installed, you can install jax
by running:gitpull
from the JAXrepository root, and rebuild by running build.py
or upgrading jaxlib
ifnecessary. You shouldn’t have to reinstall jax
because pipinstall-e
sets up symbolic links from site-packages into the repository.pytest-xdist
, which can run tests inparallel. First, install pytest-xdist
and pytest-benchmark
by runningpipinstallpytest-xdistpytest-benchmark
.Then, from the repository root directory run:--test_targets
flag. For example, you can run allthe tests of jax.numpy.pad
using:mypy
to check the type hints. To check types locally the same wayas Travis checks them:pandoc
in order to regenerate the notebooks.See Install Pandoc,or using Miniconda whichI have used successfully on the Mac: condainstall-cconda-forgepandoc
.If you do not want to install pandoc
then you should regenerate the documentationwithout the notebooks.docs/build/html/index.html
.Runallcells
thenDownloadipynb
. You may want to test that it executes properly, using sphinx-build
asexplained above.master
branch.For each code version, the building process is driven by the.readthedocs.yml
and the docs/conf.py
configuration files.test-docs
branch. That branch is also built automatically, and you cansee the generated documentation here.