Update loops to use jax calls and JIT some functions#27

chahak13 wants to merge 21 commits into

geoelements:maingeoelements/diffmpm:mainfrom

chahak13:jax_apply_loopschahak13/diff-mpm:jax_apply_loopsCopy head branch name to clipboard