Time-progressive FEM

Hello JAX-FEM,
I'm a deep modelling engineer myself and I've ran into an issue that the current time-dependant (multibody, ones that progress over time) open-source solvers don't have GPU support, only CPU. (Project Chrono FEM does not support GPU)

I will have to write a multibody solver myself, but I'm an AI engineer, not a physics sim PhD.

Would JAX-FEM be suitable as a part of an existing code base? Would anybody help implement it?

Cheers,
Maksym Riabov