Delay jax import

Author: jessegrabowski

pymc_extras/inference/find_map.py

@@ -15,7 +15,6 @@
 from pymc.initial_point import make_initial_point_fn
 from pymc.model.transform.optimization import freeze_dims_and_data
 from pymc.pytensorf import join_nonshared_inputs
-from pymc.sampling.jax import _replace_shared_variables
 from pymc.util import get_default_varnames
 from pytensor.compile import Function
 from pytensor.compile.mode import Mode
@@ -306,6 +305,8 @@ def scipy_optimize_funcs_from_loss(
     # computing jax gradients, we discard the function wrapper, so we can't handle shared variables --> rewrite them
     # away.
     if use_jax_gradients:
+        from pymc.sampling.jax import _replace_shared_variables
+
         [loss] = _replace_shared_variables([loss])
 
     compute_grad = use_grad and not use_jax_gradients