diff --git a/bletl/growth.py b/bletl/growth.py
index 8f818e5..6ffa99e 100644
--- a/bletl/growth.py
+++ b/bletl/growth.py
@@ -2,11 +2,11 @@
 import logging
 from typing import Dict, Optional, Sequence, Tuple, Union
 
-import arviz
 import calibr8
 import numpy
 import pymc as pm
 import pytensor.tensor as pt
+import xarray as xa
 from packaging import version
 
 _log = logging.getLogger(__file__)
@@ -94,8 +94,8 @@ def theta_map(self) -> Dict[str, numpy.ndarray]:
         return self._theta_map
 
     @property
-    def idata(self) -> Optional[arviz.InferenceData]:
-        """ArviZ InferenceData object of the MCMC trace."""
+    def idata(self) -> Optional[xa.DataTree]:
+        """MCMC trace obtained from running :func:`GrowthRateResult.sample`."""
         return self._idata
 
     @property
@@ -114,7 +114,7 @@ def mu_mcmc(self) -> Optional[numpy.ndarray]:
         if not self.idata:
             return None
         assert hasattr(self.idata, "posterior")
-        return self.idata.posterior.mu_t.stack(sample=("chain", "draw")).values.T
+        return self.idata["posterior"].mu_t.stack(sample=("chain", "draw")).values.T
 
     @property
     def x_mcmc(self) -> Optional[numpy.ndarray]:
@@ -122,7 +122,7 @@ def x_mcmc(self) -> Optional[numpy.ndarray]:
         if self.idata is None:
             return None
         assert hasattr(self.idata, "posterior")
-        return self.idata.posterior["X"].stack(sample=("chain", "draw")).T
+        return self.idata["posterior"]["X"].stack(sample=("chain", "draw")).T
 
     def sample(self, **kwargs) -> None:
         """Runs MCMC sampling with default settings on the growth model.
@@ -134,7 +134,6 @@ def sample(self, **kwargs) -> None:
         """
         if importlib.util.find_spec("nutpie"):
             sample_kwargs = dict(
-                return_inferencedata=True,
                 target_accept=0.95,
                 nuts_sampler="nutpie",
                 init="adapt_diag",
@@ -144,7 +143,6 @@ def sample(self, **kwargs) -> None:
             )
         else:
             sample_kwargs = dict(
-                return_inferencedata=True,
                 target_accept=0.95,
                 init="adapt_diag",
                 initvals=self.theta_map,