4 %IMP currently has rudimentary support
for running on a graphics
5 processing unit (GPU) or similar systems such as Tensor Processing Units (TPUs).
6 This support uses the [JAX](https:
8 To use the JAX support in optimization, first install the JAX library,
9 for example with `pip install jax`. Then set up the system as per usual and
10 replace any calls to IMP::core::MonteCarlo::optimize or
11 IMP::atom::MolecularDynamics::optimize with `_optimize_jax()`. If using
12 PMI, JAX can be activated using the ``use_jax`` argument to
13 IMP::pmi::macros::ReplicaExchange.
15 The JAX code is still in active development and many caveats apply:
17 - Only a small number of scoring function terms and optimizers currently
18 have JAX implementations. Trying to use others will result in a
19 Python NotImplementedError exception.
20 - Some IMP::ScoreState (aka constraint) classes do not yet work - this
21 includes common applications such as close pair containers.
23 To add JAX support for a particular IMP::Restraint, IMP::PairScore,
24 IMP::core::MonteCarloMover, or IMP::OptimizerState,
25 implement the `_get_jax()` method. See the
26 [IMP.example module](https:
29 Note that the JAX code will also run on a CPU. In some circumstances the
30 JAX code will run faster than the native %IMP C++ code on a CPU, so it may
31 be worth benchmarking both approaches.
33 Note that %IMP also has some very basic C++ support for NVIDIA GPUs using the
34 CUDA toolkit. This is unlikely to be further developed, however.
35 To build %IMP from source code with CUDA support (there are currently no
36 prebuilt %IMP binaries that use CUDA), ensure that the `nvcc` compiler
37 from NVIDIA's [CUDA toolkit](https:
38 is available, and add `-DIMP_CUDA` to your
39 [CMake invocation](@ref cmake_config).
