The power of JAX to accelerate MDO.

gemseo-jax is a GEMSEO plugin based on JAX, a Python library for high-performance array computing.

The class JAXDiscipline wraps a JAX function, with built-in automatic differentiation. This class provides useful functionalities:

• filtering of the Jacobian computation graph for specific inputs/outputs,
• jit compiling function and jacobian call for lowering cost of re-evaluation,
• performing pre-run's to trigger and log compilation times.

AutoJAXDiscipline is a special JAXDiscipline inferring the input names, output names and default input values from the signature of the JAX function, in the manner of AutoPyDiscipline.

JAXChain is a JAXDiscipline allowing to assemble a series of JaxDisciplines and execute them all in JAX. This is useful to avoid meaningless JAX-to/from-NumPy convertions.

Documentation

How to get the docs?

Bugs/Questions

Please use the gitlab issue tracker at to submit bugs or questions.

License

The gemseo-jax source code is distributed under the GNU LGPL v3.0 license. A copy of it can be found in the LICENSE.txt file. The GNU LGPL v3.0 license is an exception to the GNU GPL v3.0 license. A copy of the GNU GPL v3.0 license can be found in the LICENSES folder. The gemseo-jax product depends on other software which have various licenses. The list of dependencies with their licenses is given in the CREDITS.md file.

Contributors

• Ian Costa-Alves
• François Gallard
• Matthias De Lozzo
• Antoine DECHAUME