JAX

JAX was mentioned in #439 as a possible way of automating something like FEniCS’s derivative for calculating Jacobians in nonlinear problems; however, one user noted that they

could never manage to install jax on a windows machine in the past.

A successful reimplementation of ex10 which currently has a handwritten Jacobian

https://github.com/kinnala/scikit-fem/blob/85b83ec91a13b87e160326246246ab86d5ac63ae/docs/examples/ex10.py#L15-L19

using JAX was demonstrated in #440 .

[kinnala]

I’d prefer not to actually depend on something like JAX.

On the other hand, trying it out revealed already how an interface utilizing these autodiff tools should look like. So maybe we can eventually provide some additional utilities so that it’s easier to use JAX or autograd or similar tools in scikit-fem.

PS. Does it work in Docker?

[gdmcbain]

I tried installing JAX under Windows 10 using Miniconda3:

conda create -n jax
conda activate jax
conda install pip
pip install jax

but that didn’t work;

python -c "import jax"

raises

ModuleNotFoundError: No module named ‘jaxlib’

Then, seeing that jax is in conda-forge, I tried:

conda deactivate
conda env remove -n jax
conda create -n jax
conda activate jax
conda install -c conda-forge jax

but that failed with

Collecting package metadata (current_repodata.json): done
Solving environment: failed with initial frozen solve. Retrying with flexible solve.
Solving environment: failed with repodata from current_repodata.json, will retry with next repodata source.
Collecting package metadata (repodata.json): done
Solving environment: failed with initial frozen solve. Retrying with flexible solve.
Solving environment: -
Found conflicts! Looking for incompatible packages.
This can take several minutes.  Press CTRL-C to abort.
failed

UnsatisfiableError: