Merge pull request #2729 from jeromtom/Issue#2728-Update-the-docstring-of-install_jax

Fix the API docs of install_jax

Author: valentinsulzer

docs/source/api/util.rst

@@ -22,6 +22,8 @@ Utility functions
 
 .. autofunction:: pybamm.get_parameters_filepath
 
+.. autofunction:: pybamm.install_jax
+
 .. autofunction:: pybamm.have_jax
 
 .. autofunction:: pybamm.is_jax_compatible

pybamm/util.py

@@ -336,13 +336,13 @@ def install_jax(arguments=None):  # pragma: no cover
     """
     Install compatible versions of jax, jaxlib.
 
-    Command Line Interface:
-    -----------------------
-    >>> pybamm_install_jax
+    Command Line Interface::
 
-    optional arguments:
-    -h, --help   show help message
-    -f, --force  force install compatible versions of jax and jaxlib
+        $ pybamm_install_jax
+
+    |   optional arguments:
+    |    -h, --help   show help message
+    |    -f, --force  force install compatible versions of jax and jaxlib
     """
     parser = argparse.ArgumentParser(description="Install jax and jaxlib")
     parser.add_argument(