JAX-FEM with MLP

[leedz0210]

Hello Author,

Thank you for creating such a useful tool. I have a question regarding the integration of JAX-FEM with neural networks.

Question: Can JAX-FEM be incorporated as a module within a neural network to serve as part of the training process?

Background: I am currently using FEniCS, and when training with MLP (Multi-Layer Perceptron), I encounter the following error:

jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method array() was called on traced array with shape float32

This error occurs because FEniCS is not compatible with JAX's automatic differentiation system, as it tries to convert JAX traced arrays to NumPy arrays during execution.

My Questions:

Can JAX-FEM be directly used as part of a neural network's computational graph?
Does JAX-FEM support end-to-end differentiation through the FEM solver?
Do you have any related work or examples demonstrating this integration?
I would greatly appreciate any guidance or references you could provide.

Thank you very much!

[paulgekeler]

(I am not the author), but you can wrap the solver which makes it possible to differentiate an objective function with respect to some parameter and defined the forward and backward pass. See this example from the docs. Any additional parameters you might have need to be passed to Problem.set_params() to make them differentiable. Have a look at problem.py and the other examples in the docs.

[tianjuxue]

Perhaps this example will be of some relevance.

[Strawl]

@leedz0210
I have done so in the past. Here is some related work: link to gist.
The only downside of jax-fem is the lack of jit/vmap support. Feax would be another library written in jax, which has jit/vmap support, on the downside that the solvers implemented are less powerful compared to jax-fem's.

[leedz0210]

Thank you for your helpful answers. I will carefully consider them.Thank you again.