PDE Demos with JAX Backend

[jessebett]

Hello,

I have a question about current support for the JAX backend. In particular I am attempting to follow the demos of PDE problems with a JAX backend.

I noticed that the "Complete code" section at the bottom of the demo pages include as the first line:
"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, paddle"""
Is the absence of jax in this list meant to indicate that JAX backend is not supported on the PDE examples?

If I try to follow a PDE demo naively, only setting the DDE_BACKEND=jax and making minimal modification to the demo, I encounter an error.
For example, the Poisson equation in 1D with Dirichlet boundary conditions,
with the pde defined by the function

import jax.numpy as np  

def pde(x, y):
    dy_xx = dde.grad.hessian(y, x)
    return -dy_xx - np.pi ** 2 * np.sin(np.pi * x)

(note the modification from the demo, that jax.numpy.sin replaces tf.sin)

If I attempt to run this code I get an error:
raise ValueError("The component of y is missing.")

Investigating the Trace further, it seems that this error is being thrown in deepxde/gradients.py at:

... in hessian
   return hessian._Hessians(ys, xs, component=component, i=i, j=j, grad_y=grad_y)
... in __call__
   self.Hs[key] = Hessian(y, xs, component=component, grad_y=grad_y)
... in __init__
    raise ValueError("The component of y is missing.")

Does this error indicate that dde.grad.hessian is unsupported with the JAX backend?
Alternatively, is there a way I can modify the call dy_xx = dde.grad.hessian(y, x) that is different for the JAX backend?

Thanks!

[pescap]

Same error here. I was currently trying to fix that. Sounds like jax backend was tested for 1st order derivatives (see here: #655):

deepxde/examples/pinn_forward/ode_system.py

Lines 15 to 19 in 0736d1d

	# Backend jax
	# y_val, y_fn = y
	# y1, y2 = y_val[:, 0:1], y_val[:, 1:]
	# dy1_x, _ = dde.grad.jacobian(y, x, i=0)
	# dy2_x, _ = dde.grad.jacobian(y, x, i=1)

but not for 2nd order systems, e.g, https://github.com/lululxvi/deepxde/blob/master/examples/pinn_forward/ode_2nd.py.

[pescap]

Hi, use:

dy_xx, _ = dde.grad.hessian(y, x, component =0)

[jessebett]

Thanks! I did not notice the component argument.

[pescap]

• Lshape is ok (see added Lshape for jax #971)
• I am experiencing troubles with BCs involving derivatives:

deepxde/examples/pinn_forward/ode_2nd.py

Line 28 in 0736d1d

return dde.grad.jacobian(outputs, inputs, i=0, j=None) - 2

outputs is already in the form y_val and not as y. I think that it is what causes the error here:

deepxde/deepxde/gradients.py

Line 29 in 0736d1d

self.dim_y = ys[0].shape[1]

[pescap]

Note: jax is really fast 👍