We present the use of the JAX machine learning ecosystem to develop an expressive environment for implementing material constitutive models and their subsequent integration into finite element solvers, e.g. those based on domain-specific languages like FEniCS and more traditional UMAT-conforming solvers such as ABAQUS and COMSOL. By casting constitutive laws in a differentiable programming framework, these models can be treated as parameterized operators in the same formal sense as modern machine-learning models, enabling modular composition and seamless extension to hybrid physics- and data-driven formulations. In particular, we introduce a JAX-based domain-specific language for defining complex material models and demonstrate the use of automatic differentiation to systematically derive consistent tangent operators. Leveraging JAX’s static directed-acyclic-graph representation of complete constitutive models, this approach enables whole-model differentiation and compilation even in the presence of complex outer algorithms, such as ODE-based internal-variable evolution and nonlinear iteration schemes, as provided by JAX-based libraries like Diffrax and Equinox. We further discuss efficient execution on GPU/TPU hardware, and the integration of these models into finite element simulation workflows. This project has received funding from the European Research Council (ERC) under the European Union’s Horizon Europe programme, Consolidator Grant No. 101229452 (Automatix). This research was funded in whole, or in part, by the Luxembourg National Research Fund (FNR), grant reference PRIDE/21/16747448/MATHCODA.