Molecular systems used for force field training tend to have a small number of atoms compared to regular MD runs. Using a software designed for running simulations with thousands of atoms (such as PuReMD or LAMMPS/ReaxFF) to run several small scale simulations introduces overheads. Optimizations in these software (such as optimized sparse solvers for atomic charges, fast neighbor list generation algorithms and distributed computation) would actually increase the overall run-time for small systems and result in unnecessarily complex code. Even though vanilla Python code tends to be slower than optimized Fortran or C code, when the auto-diff functionality, benefits of targeting small geometries and the just-in-compiled XLA support (discussed in Section 3.3 are considered, the advantages of JAX-ReaxFF outweighs the performance loss from not using Fortran or C. While gradient-based optimization functionality is straight-forward to achieve using JAX as described above, there are a number of important considerations to realize an efficient (from a runtime point-of-view) and a scalable (from a memory utilization perspective) parameter optimization framework. Fig. gives an overview of the task-flow in JAX-ReaxFF. After the neighbor list and interactions lists are calculated for the input geometries (Section 3.2, we cluster the inputs based on the size of their interaction lists and align them properly in memory to ensure efficient single instruction multiple data (SIMD) parallelization (Section 3.3). After these preparation steps, the main optimization loop is executed until convergence or the maximum number of optimization steps are reached (which typically takes only tens of iterations). During the parameter optimization loop, some molecules might require energy minimization before the final calculation to prevent overfitting and also to tune the parameters based on more likely states of the atom clusters as the lower energy states are more likely to be observed. Hence, the main optimization loop contains a "gradient-based optimization" step followed by a "geometry optimization" step. We discuss each step in more detail in the ensuing subsections.