Skip to content

Added support for JAX for Neumann and Robin BC #2015

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Added support for JAX for Neumann and Robin BC #2015

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions deepxde/gradients/gradients_reverse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,7 @@ def __call__(self, i=None, j=None):
elif backend_name == "jax":
ndim_y = bkd.ndim(self.ys[0])
if ndim_y == 3:
raise NotImplementedError(
"Reverse-mode autodiff doesn't support 3D output"
)
raise NotImplementedError("Reverse-mode autodiff doesn't support 3D output")
Copy link
Owner

@lululxvi lululxvi Aug 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't modify this.


# Compute J[i, :]
if i not in self.J:
Expand Down Expand Up @@ -138,6 +136,8 @@ def __call__(self, ys, xs, component=0, i=0, j=0):
key = (id(ys[0]), id(xs), component)
if key not in self.Hs:
self.Hs[key] = Hessian(ys, xs, component=component)
if backend_name == "jax":
return self.Hs[key](i, j)[0]
return self.Hs[key](i, j)

def clear(self):
Expand Down
11 changes: 9 additions & 2 deletions deepxde/icbc/boundary_conditions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,12 @@ def collocation_points(self, X):
return self.filter(X)

def normal_derivative(self, X, inputs, outputs, beg, end):
dydx = grad.jacobian(outputs, inputs, i=self.component, j=None)[beg:end]
if backend_name == "jax":
dydx = grad.jacobian(
(outputs, self.func), inputs, i=self.component, j=None
Copy link
Owner

@lululxvi lululxvi Aug 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why self.func? Here we should compute the derivative of network outputs wrt network inputs. The self.func is the function in the BC.

Copy link
Author

@Pheithar Pheithar Aug 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Inside gradients, the Jacobian calculation needs an array with 2 elements, where the first element is the same as in the other backends (what we want to compute the gradient to), but the second element needs to be a function that computes the gradient (the BC in this case). Previously, only the output was fed to the grad.jacobian, and that was producing an error. It should have been an out of index error, in line 50 of deepxde/gradients/gradients_reverse.py and line 76 of deepxde/gradients/gradients_forward.py, but JAX does not throw an error in those instances.

Therefore, in those lines, the error was different, which made it difficult to pinpoint, and my conclusion was that the function to compute the gradient against needed to be passed forward to the Jacobian calculation.

elif backend_name == "jax":
    tangent = jax.numpy.zeros(self.dim_x).at[j].set(1)
    grad_fn = lambda x: jax.jvp(self.ys[1], (x,), (tangent,))[1]
    self.J[j] = (jax.vmap(grad_fn)(self.xs), grad_fn)

This is what was causing the error. When computing self.ys[1], instead of out of bounds, it was throwing a different error. (I believe it was something along the lines of TypeError: 'jaxlib._jax.ArrayImpl' object is not callable)

Copy link
Owner

@lululxvi lululxvi Aug 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, for JAX, we need another element of the function. My question is why it is self.func? Should it be the function of the network forward pass?

Copy link
Author

@Pheithar Pheithar Aug 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think so. We want to compute the boundaries conditions. That part of the code is called in the NeumannBC to compute the loss term, and if I am not mistaken, self.func is $f(X)$, and we want our loss term for Neumann boundaries to be $\mathcal{L}_{\text{Neumann}} = \Delta f(X) - Y$. (I hope I am explaining myself here).

In short, what I mean is that we do not want to compute the gradients w.r.t the network, as what we want is to compute a term for the loss, not the gradient of the loss that can be used for backpropagation.

Copy link
Owner

@lululxvi lululxvi Aug 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Neumann BC is dy/dn = f(x), where dy/dn dy/dx in the n direction.
https://en.wikipedia.org/wiki/Neumann_boundary_condition

)[0][beg:end]
else:
dydx = grad.jacobian(outputs, inputs, i=self.component, j=None)[beg:end]
n = self.boundary_normal(X, beg, end, None)
return bkd.sum(dydx * n, 1, keepdims=True)

Expand Down Expand Up @@ -282,7 +287,9 @@ def collocation_points(self, X):

def error(self, X, inputs, outputs, beg, end, aux_var=None):
if self.batch_size is not None:
return self.func(inputs, outputs, X)[beg:end] - self.values[self.batch_indices]
return (
Copy link
Owner

@lululxvi lululxvi Aug 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't modify this.

self.func(inputs, outputs, X)[beg:end] - self.values[self.batch_indices]
)
return self.func(inputs, outputs, X)[beg:end] - self.values


Expand Down
Loading