Enhanced JAX compatibility#1490

tttc3 wants to merge 5 commits intolululxvi:masterlululxvi/deepxde:masterfrom

tttc3:jax-compatibilitytttc3/deepxde:jax-compatibilityCopy head branch name to clipboard