Add JAX backend support

[bonneted]

No description provided.

[bonneted]

I've got the same final parameters as the ones in the paper :

Epoch	E	tp	ti	td	k	Rm	a1	C1	C2	C4	C5	Ub	U0	Um	Rg	alpha	beta
1010000	0.20086897812616378	5.984280323717105	101.07973763650493	11.9777607693349	0.008332535739606049	208.50174547489897	6.593664003041278	301.0031941572703	34.69381362708656	79.07846369570885	25.952992080032526	71.5795274689672	3.99920640001627	89.31458070648316	179.61120094478127	7.538939453290313	1.778194872565763

[lululxvi]

Is this ready to merge?

[bonneted]

The code relies on this DeepXDE PR that isn't merged yet: lululxvi/deepxde#1775
I've just updated the code of DeepXDE PR, please let me know if you'd like any more changes.

I've tested the code again and it's working as intended.
So whenever the other PR is accepted, this one should be ready to merge.

[lululxvi]

Is this ready to merge?

[bonneted]

Yes, I just changed epochs to iterations.

I've run the code again with dde master branch, it works as expected. It took 20min on a RTX 4070 GPU with jax backend.
Here are the final values:

Epoch	E	tp	ti	td	k	Rm	a1	C1	C2	C4	C5	Ub	U0	Um	Rg	alpha	beta
1010000	0.2212	5.882	140.0	11.72	0.008384	376.2	6.405	428.2	33.55	43.03	22.20	60.63	0.9840	77.95	155.9	13.50	2.273

Some values are worse than the original paper, but this is probably due to the stochasticity of the run.