Delay jax import

Author: jessegrabowski

pymc_extras/inference/find_map.py

@@ -3,7 +3,6 @@
 from collections.abc import Callable
 from typing import Literal, cast, get_args
 
-import jax
 import numpy as np
 import pymc as pm
 import pytensor
@@ -138,6 +137,13 @@ def _compile_grad_and_hess_to_jax(
     f_hessp: Callable | None
         The compiled hessian-vector product function, or None if use_hessp is False.
     """
+    try:
+        import jax
+    except ImportError:
+        raise ImportError(
+            "You don't have jax installed -- it is required if gradient_backend='jax'"
+        )
+
     f_hess = None
     f_hessp = None