Welcome to diffopt’s documentation!
DiffOpt is a collection of tools that are useful for parallelizable
optimization of large-parameter, memory-intensive, and/or differentiable
models implemented in Jax. It is composed of the following three subpackages:
(1) multigrad enables enables you to define a data-parallelized (over MPI)
loss function and compute its gradient, (2) kdescent performs stochastic
gradient descent over mini-batched KDE statistics, and (3) multiswarm is an
MPI-parallelized implementation of Particle Swarm Optimization (PSO).
The code is open-source and available on
GitHub.