Add return type documentation for util.have_jax() and is_jax_compatib…

…le() (#4005)

* add return type for util.have_jax() and is_jax_compatible()

* add description

* Update pybamm/util.py

Co-authored-by: Agriya Khetarpal <74401230+agriyakhetarpal@users.noreply.github.com>

---------

Co-authored-by: Eric G. Kratz <kratman@users.noreply.github.com>
Co-authored-by: Agriya Khetarpal <74401230+agriyakhetarpal@users.noreply.github.com>

pybamm/util.py

@@ -260,7 +260,15 @@ def get_parameters_filepath(path):
 
 
 def have_jax():
-    """Check if jax and jaxlib are installed with the correct versions"""
+    """
+    Check if jax and jaxlib are installed with the correct versions
+
+    Returns
+    -------
+    bool
+        True if jax and jaxlib are installed with the correct versions, False if otherwise
+
+    """
     return (
         (importlib.util.find_spec("jax") is not None)
         and (importlib.util.find_spec("jaxlib") is not None)
@@ -269,7 +277,14 @@ def have_jax():
 
 
 def is_jax_compatible():
-    """Check if the available version of jax and jaxlib are compatible with PyBaMM"""
+    """
+    Check if the available versions of jax and jaxlib are compatible with PyBaMM
+
+    Returns
+    -------
+    bool
+        True if jax and jaxlib are compatible with PyBaMM, False if otherwise
+    """
     return importlib.metadata.distribution("jax").version.startswith(
         JAX_VERSION
     ) and importlib.metadata.distribution("jaxlib").version.startswith(JAXLIB_VERSION)