Solver for pseudo-2D model of Li-ion battery based on Finite difference Method. Incorporation of JAX for the nonlinear solver.