Another important advantage of JAX is its architectural portability enabled by the XLA technology used under the hood. Hence, JAX-ReaxFF can run efficiently on various architectures, including graphics procesing units (GPU) and tensor processing units (TPU), through automatic thread parallelization and vector processing. As we demonstrate through extensive tests, JAX-ReaxFF can reduce the overall training time by up to three orders of magnitude compared to the existing global optimization schemes, while achieving similar (or better) fitness scores and yielding accurate simulation results.