serket (pypi): JAX NN library.
Keywords: python|machine-learning|pytorch|jax|finite-difference|neural-network