Add ZarrTrace

Author: lucianopaz

.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),
 }