print("RESULTS")
print(f"NumPy:          {numpy_time:.4f} seconds")
print(f"JAX (no JIT):   {jax_time:.4f} seconds")
print(f"JAX (with JIT): {jit_time:.4f} seconds")
print(f"\nSpeedup (JIT vs NumPy): {numpy_time / jit_time:.1f}x")

        __
        __