The code in this repository is based on the papers Differential geometry and stochastic dynamics with deep learning numerics arXiv:1712.08364 and Computational Anatomy in Theano arXiv:1706.07690.

The code is a reimplementation of the Theano Geometry library https://bitbucket.org/stefansommer/jaxgeometry/ replacing Theano with Jax https://github.com/google/jax.

The source repository is at https://bitbucket.org/stefansommer/jaxgeometry/

Please contact Stefan Sommer [email protected]

Please use Python 3.X.

Install with

pip install jaxdifferentialgeometry

Check out the source with git sandiInstall required packages:

pip install -r requirements.txt

Use e.g. a Python 3 virtualenv:

virtualenv -p python3 .
source bin/activate
pip install -r requirements.txt

If you don't use a virtual environment, make sure that you are actually using Python 3, e.g. use pip3 instead of pip.

Alternatively, use conda:

conda install -c conda-forge jaxlib
conda install -c conda-forge jax

and similarly for the remaining requirements in requirements.txt.

After cloning the source repository, start jupyter notebook

PYTHONPATH=$(pwd)/src jupyter notebook

Your browser should now open and you can find the example Jax Geometry notebooks in the examples folder.

Some good discussions about the architectural differences between autodiff frameworks: https://www.assemblyai.com/blog/why-you-should-or-shouldnt-be-using-jax-in-2022/ and http://www.stochasticlifestyle.com/engineering-trade-offs-in-automatic-differentiation-from-tensorflow-and-pytorch-to-jax-and-julia/