Revert "Minor fix of blackjax import in fit_pathfinder function (pymc-devs#443)"

This reverts commit 1fff560.

Author: aphc14

pymc_extras/inference/fit.py

@@ -11,10 +11,9 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-import arviz as az
 
 
-def fit(method: str, **kwargs) -> az.InferenceData:
+def fit(method, **kwargs):
     """
     Fit a model with an inference algorithm
 

pymc_extras/inference/pathfinder/pathfinder.py

@@ -21,9 +21,11 @@
 from collections.abc import Callable, Iterator
 from dataclasses import asdict, dataclass, field, replace
 from enum import Enum, auto
+from importlib.util import find_spec
 from typing import Literal, TypeAlias
 
 import arviz as az
+import blackjax
 import filelock
 import jax
 import numpy as np
@@ -1734,8 +1736,8 @@ def fit_pathfinder(
         )
         pathfinder_samples = mp_result.samples
     elif inference_backend == "blackjax":
-        import blackjax
-
+        if find_spec("blackjax") is None:
+            raise RuntimeError("Need BlackJAX to use `pathfinder`")
         if version.parse(blackjax.__version__).major < 1:
             raise ImportError("fit_pathfinder requires blackjax 1.0 or above")