Add ZarrTrace and fix non-determinism with Generators in sample

.github/workflows/tests.yml

@@ -115,6 +115,7 @@ jobs:
 
           - |
             tests/backends/test_mcbackend.py
+            tests/backends/test_zarr.py
             tests/distributions/test_truncated.py
             tests/logprob/test_abstract.py
             tests/logprob/test_basic.py
@@ -240,6 +241,7 @@ jobs:
 
           - |
             tests/backends/test_arviz.py
+            tests/backends/test_zarr.py
             tests/variational/test_updates.py
       fail-fast: false
     runs-on: ${{ matrix.os }}

conda-envs/environment-dev.yml

@@ -19,6 +19,7 @@ dependencies:
 - scipy>=1.4.1
 - typing-extensions>=3.7.4
 - threadpoolctl>=3.1.0
+- zarr>=2.5.0,<3
 # Extra dependencies for dev, testing and docs build
 - ipython>=7.16
 - jax

conda-envs/environment-docs.yml

@@ -17,6 +17,7 @@ dependencies:
 - scipy>=1.4.1
 - typing-extensions>=3.7.4
 - threadpoolctl>=3.1.0
+- zarr>=2.5.0,<3
 # Extra dependencies for docs build
 - ipython>=7.16
 - jax

conda-envs/environment-jax.yml

@@ -10,6 +10,7 @@ dependencies:
 - cachetools>=4.2.1
 - cloudpickle
 - h5py>=2.7
+- zarr>=2.5.0,<3
 # Jaxlib version must not be greater than jax version!
 - blackjax>=1.2.2
 - jax>=0.4.28

conda-envs/environment-test.yml

@@ -21,6 +21,7 @@ dependencies:
 - scipy>=1.4.1
 - typing-extensions>=3.7.4
 - threadpoolctl>=3.1.0
+- zarr>=2.5.0,<3
 # Extra dependencies for testing
 - ipython>=7.16
 - pre-commit>=2.8.0

conda-envs/windows-environment-dev.yml

@@ -20,6 +20,7 @@ dependencies:
 - scipy>=1.4.1
 - typing-extensions>=3.7.4
 - threadpoolctl>=3.1.0
+- zarr>=2.5.0,<3
 # Extra dependencies for dev, testing and docs build
 - ipython>=7.16
 - myst-nb<=1.0.0

conda-envs/windows-environment-test.yml

@@ -23,6 +23,7 @@ dependencies:
 - scipy>=1.4.1
 - typing-extensions>=3.7.4
 - threadpoolctl>=3.1.0
+- zarr>=2.5.0,<3
 # Extra dependencies for testing
 - ipython>=7.16
 - pre-commit>=2.8.0

docs/source/api/backends.rst

@@ -20,3 +20,5 @@ Internal structures
    NDArray
    base.BaseTrace
    base.MultiTrace
+   zarr.ZarrTrace
+   zarr.ZarrChain

docs/source/conf.py

@@ -309,6 +309,7 @@
     "python": ("https://docs.python.org/3/", None),
     "scipy": ("https://docs.scipy.org/doc/scipy/", None),
     "xarray": ("https://docs.xarray.dev/en/stable/", None),
+    "zarr": ("https://zarr.readthedocs.io/en/stable/", None),
 }
 
 

pymc/backends/__init__.py

@@ -72,6 +72,7 @@
 from pymc.backends.arviz import predictions_to_inference_data, to_inference_data
 from pymc.backends.base import BaseTrace, IBaseTrace
 from pymc.backends.ndarray import NDArray
+from pymc.backends.zarr import ZarrTrace
 from pymc.blocking import PointType
 from pymc.model import Model
 from pymc.step_methods.compound import BlockedStep, CompoundStep
@@ -120,15 +121,27 @@ def _init_trace(
 
 def init_traces(
     *,
-    backend: TraceOrBackend | None,
+    backend: TraceOrBackend | ZarrTrace | None,
     chains: int,
     expected_length: int,
     step: BlockedStep | CompoundStep,
     initial_point: PointType,
     model: Model,
     trace_vars: list[TensorVariable] | None = None,
+    tune: int = 0,
 ) -> tuple[RunType | None, Sequence[IBaseTrace]]:
     """Initialize a trace recorder for each chain."""
+    if isinstance(backend, ZarrTrace):
+        backend.init_trace(
+            chains=chains,
+            draws=expected_length - tune,
+            tune=tune,
+            step=step,
+            model=model,
+            vars=trace_vars,
+            test_point=initial_point,
+        )
+        return None, backend.straces
     if HAS_MCB and isinstance(backend, Backend):
         return init_chain_adapters(
             backend=backend,