add support for rosenbrock method

[poonai]

Hello,

I'm new to JAX and numerical computing, and willing to invest the time to learn by implementing numerical methods. After this, I plan to add support for DAE and additional rosenbrock methods. I would appreciate your guidance on getting this PR merged.

Thanks

[poonai]

@patrick-kidger I've addressed your comments. Please review them when you get a chance. In the meantime, I'll start reading about other methods.

Thanks