JAX library performs auto-differentiation on native Python code. As such implementation of the ReaxFF energy expressions (see Eq. ( )) in Python forms the core of JAX-ReaxFF. Once the individual energy expressions and the training error function are provided, JAX can easily calculate the gradient of the training error function with respect to the ReaxFF parameters to be optimized. As mentioned earlier, atomic forces can also be part of the training dataset, these can be calculated using the gradients of ReaxFF energy expressions with respect to atom positions, too.