pybamm.JaxSolver docstring improvement

docs/conf.py

@@ -281,6 +281,7 @@
     "numpy": ("https://numpy.org/doc/stable", None),
     "scipy": ("https://docs.scipy.org/doc/scipy", None),
     "matplotlib": ("https://matplotlib.org/stable/", None),
+    "jax": ("https://jax.readthedocs.io/en/stable", None),
 }
 
 # -- sphinxcontrib-bibtex configuration --------------------------------------

pybamm/solvers/jax_solver.py

@@ -30,9 +30,9 @@ class JaxSolver(pybamm.BaseSolver):
 
     Parameters
     ----------
-    method: str
-        'RK45' (default) uses jax.experimental.odeint
-        'BDF' uses custom jax_bdf_integrate (see jax_bdf_integrate.py for details)
+    method: str, optional (see :func:`jax.experimental.odeint` for details)
+        * 'RK45' (default) uses jax.experimental.odeint
+        * 'BDF' uses custom jax_bdf_integrate (see :func:`jax_bdf_integrate.py` for details)
     root_method: str, optional
         Method to use to calculate consistent initial conditions. By default this uses
         the newton chord method internal to the jax bdf solver, otherwise choose from