Add way to set backend fn random generators

.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,9 +72,11 @@
 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
+from pymc.util import get_random_generator
 
 HAS_MCB = False
 try:
@@ -102,11 +104,13 @@ def _init_trace(
     model: Model,
     trace_vars: list[TensorVariable] | None = None,
     initial_point: PointType | None = None,
+    rng: np.random.Generator | None = None,
 ) -> BaseTrace:
     """Initialize a trace backend for a chain."""
+    rng_ = get_random_generator(rng)
     strace: BaseTrace
     if trace is None:
-        strace = NDArray(model=model, vars=trace_vars, test_point=initial_point)
+        strace = NDArray(model=model, vars=trace_vars, test_point=initial_point, rng=rng_)
     elif isinstance(trace, BaseTrace):
         if len(trace) > 0:
             raise ValueError("Continuation of traces is no longer supported.")
@@ -120,22 +124,37 @@ 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,
+    rng: np.random.Generator | None = None,
 ) -> 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,
+            rng=rng,
+        )
+        return None, backend.straces
     if HAS_MCB and isinstance(backend, Backend):
         return init_chain_adapters(
             backend=backend,
             chains=chains,
             initial_point=initial_point,
             step=step,
             model=model,
+            rng=rng,
         )
 
     assert backend is None or isinstance(backend, BaseTrace)
@@ -148,7 +167,8 @@ def init_traces(
             model=model,
             trace_vars=trace_vars,
             initial_point=initial_point,
+            rng=rng_,
         )
-        for chain_number in range(chains)
+        for chain_number, rng_ in enumerate(get_random_generator(rng).spawn(chains))
     ]
     return None, traces

pymc/backends/base.py

@@ -34,7 +34,7 @@
 
 from pymc.backends.report import SamplerReport
 from pymc.model import modelcontext
-from pymc.pytensorf import compile
+from pymc.pytensorf import compile, copy_function_with_new_rngs
 from pymc.util import get_var_name
 
 logger = logging.getLogger(__name__)
@@ -159,6 +159,7 @@ def __init__(
         fn=None,
         var_shapes=None,
         var_dtypes=None,
+        rng=None,
     ):
         model = modelcontext(model)
 
@@ -177,6 +178,8 @@ def __init__(
                 on_unused_input="ignore",
             )
             fn.trust_input = True
+        if rng is not None:
+            fn = copy_function_with_new_rngs(fn=fn, rng=rng)
 
         # Get variable shapes. Most backends will need this
         # information.

pymc/backends/mcbackend.py

@@ -29,7 +29,7 @@
 
 from pymc.backends.base import IBaseTrace
 from pymc.model import Model
-from pymc.pytensorf import PointFunc
+from pymc.pytensorf import PointFunc, copy_function_with_new_rngs
 from pymc.step_methods.compound import (
     BlockedStep,
     CompoundStep,
@@ -38,6 +38,7 @@
     flat_statname,
     flatten_steps,
 )
+from pymc.util import get_random_generator
 
 _log = logging.getLogger(__name__)
 
@@ -96,7 +97,11 @@ class ChainRecordAdapter(IBaseTrace):
     """Wraps an McBackend ``Chain`` as an ``IBaseTrace``."""
 
     def __init__(
-        self, chain: mcb.Chain, point_fn: PointFunc, stats_bijection: StatsBijection
+        self,
+        chain: mcb.Chain,
+        point_fn: PointFunc,
+        stats_bijection: StatsBijection,
+        rng: np.random.Generator | None = None,
     ) -> None:
         # Assign attributes required by IBaseTrace
         self.chain = chain.cmeta.chain_number
@@ -107,8 +112,11 @@ def __init__(
             for sstats in stats_bijection._stat_groups
         ]
 
+        self._rng = rng
         self._chain = chain
         self._point_fn = point_fn
+        if rng is not None:
+            self._point_fn = copy_function_with_new_rngs(self._point_fn, rng)
         self._statsbj = stats_bijection
         super().__init__()
 
@@ -257,6 +265,7 @@ def init_chain_adapters(
     initial_point: Mapping[str, np.ndarray],
     step: CompoundStep | BlockedStep,
     model: Model,
+    rng: np.random.Generator | None,
 ) -> tuple[mcb.Run, list[ChainRecordAdapter]]:
     """Create an McBackend metadata description for the MCMC run.
 
@@ -286,7 +295,8 @@ def init_chain_adapters(
             chain=run.init_chain(chain_number=chain_number),
             point_fn=point_fn,
             stats_bijection=statsbj,
+            rng=rng_,
         )
-        for chain_number in range(chains)
+        for chain_number, rng_ in enumerate(get_random_generator(rng).spawn(chains))
     ]
     return run, adapters