From 387c11634cc4c2976d7d744776d4d6de77e2de18 Mon Sep 17 00:00:00 2001
From: Daniel Weindl <daniel.weindl@helmholtz-munich.de>
Date: Wed, 3 Apr 2024 22:55:05 +0200
Subject: [PATCH] Import arviz only for type checking

Related to #1354
---
 pypesto/sample/pymc.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/pypesto/sample/pymc.py b/pypesto/sample/pymc.py
index c6b37e3fa..fc0d08d84 100644
--- a/pypesto/sample/pymc.py
+++ b/pypesto/sample/pymc.py
@@ -3,9 +3,8 @@
 from __future__ import annotations
 
 import logging
-from typing import Union
+from typing import TYPE_CHECKING, Union
 
-import arviz as az
 import numpy as np
 import pymc
 import pytensor.tensor as pt
@@ -16,6 +15,9 @@
 from ..result import McmcPtResult
 from .sampler import Sampler, SamplerImportError
 
+if TYPE_CHECKING:
+    import arviz as az
+
 logger = logging.getLogger(__name__)
 
 # implementation based on: