Hessian for Implicit Differentiation

This repository has two purposes:

• Provides an implementation for Hessian-vector products in implicit differentiable programming
• Shows more examples of differentiable finite elements based on JAX-FEM

General picture

Differentiable programming breaks the boundary between deep learning and differentiable physics.

This repository is based on JAX-FEM to solve differentiable physics problems, providing second-order derivative information in the form of Hessian-vector products.

Quick start

Refer to simple.ipynb for a simple illustrative example.

Installation

Works with JAX-FEM version 0.0.9.

Examples

E1: Source field identification

Goal: Change the source term to match observed data.

Predicted solutions gradually match the reference data.

E2: Boundary force identification

Goal: Change the boundary traction force to match observed displacement.

Predicted displacements gradually match the reference displacement.

E3: Thermal-mechanical control

Goal: Change the boundary temperature to achieve desired deformation.

Predicted displacements gradually match the reference displacement.

E4: Shape optimization

Goal: Rotate the square-shaped holes for better beam stiffness.

Compliance minimization by changing orientations of the square holes.

Paper

Please consider citing the following article if you find the code is useful.

@article{xue2026implicit,
  title={Implicit differentiation with second-order derivatives and benchmarks in finite-element-based differentiable physics},
  author={Xue, Tianju},
  journal={Computer Physics Communications},
  pages={110102},
  year={2026},
  publisher={Elsevier}
}