pybamm.JaxSolver docstring improvement (#3994)

* pybamm.JaxSolver docstring improvement * Update conf.py to add `intersphinx_mapping` for `jax` * Update jax_solver.py to fix docstring rendering * Update docs/conf.py Co-authored-by: Agriya Khetarpal <[email protected]> * Update pybamm/solvers/jax_solver.py * Update docs/conf.py * Apply suggestions from code review * Update pybamm/solvers/jax_solver.py Co-authored-by: Arjun Verma <[email protected]> * Update pybamm/solvers/jax_solver.py Co-authored-by: Agriya Khetarpal <[email protected]> * Update docs/conf.py Co-authored-by: Eric G. Kratz <[email protected]> --------- Co-authored-by: Agriya Khetarpal <[email protected]> Co-authored-by: Arjun Verma <[email protected]> Co-authored-by: Eric G. Kratz <[email protected]>
pybamm-team · May 31, 2024 · 16da4ca · 16da4ca
1 parent ce1d87a
commit 16da4ca
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pybamm/solvers/jax_solver.py b/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 `jax.experimental.ode.odeint` for details)
+        * 'RK45' (default) uses jax.experimental.ode.odeint
+        * 'BDF' uses custom jax_bdf_integrate (see `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