feat: use new nuts-rs storage interface

Author: aseyboldt

Cargo.lock

@@ -10,7 +10,7 @@ license = "MIT"
 repository = "https://github.com/pymc-devs/nutpie"
 keywords = ["statistics", "bayes"]
 description = "Python wrapper for nuts-rs -- a NUTS sampler written in Rust."
-rust-version = "1.76"
+rust-version = "1.90"
 
 [features]
 extension-module = ["pyo3/extension-module"]
@@ -21,23 +21,28 @@ name = "_lib"
 crate-type = ["cdylib"]
 
 [dependencies]
-nuts-rs = "0.16.1"
+nuts-rs = { version = "0.16.1", features = ["zarr", "arrow"] }
 numpy = "0.26.0"
 rand = "0.9.0"
 thiserror = "2.0.3"
 rand_chacha = "0.9.0"
-rayon = "1.10.0"
-# Keep arrow in sync with nuts-rs requirements
-arrow = { version = "56.1.0", default-features = false, features = ["ffi"] }
+rayon = "1.11.0"
 anyhow = "1.0.72"
 itertools = "0.14.0"
-bridgestan = "2.6.1"
+bridgestan = "2.7.0"
 rand_distr = "0.5.0"
-smallvec = "1.14.0"
+smallvec = "1.15.0"
 upon = { version = "0.10.0", default-features = false, features = [] }
 time-humanize = { version = "0.1.3", default-features = false }
 indicatif = "0.18.0"
 tch = { version = "0.20.0", optional = true }
+pyo3-object_store = "0.6.0"
+# Keep zarrs crates in sync with nuts-rs requirements
+zarrs = { version = "0.22.2", features = ["async"] }
+zarrs_object_store = "0.5.0"
+tokio = { version = "1.47.1", features = ["rt", "rt-multi-thread"] }
+pyo3-arrow = "0.12.0"
+arrow = { version = "56.2.0", features = ["json"] }
 
 [dependencies.pyo3]
 version = "0.26.0"

Cargo.toml

@@ -17,6 +17,7 @@ classifiers = [
 
 dependencies = [
   "pyarrow >= 12.0.0",
+  "arro3-core >= 0.6.0",
   "pandas >= 2.0",
   "xarray >= 2025.01.2",
   "arviz >= 0.20.0",
@@ -28,12 +29,12 @@ Homepage = "https://pymc-devs.github.io/nutpie/"
 Repository = "https://github.com/pymc-devs/nutpie"
 
 [project.optional-dependencies]
-stan = ["bridgestan >= 2.6.1", "stanio >= 0.5.1"]
+stan = ["bridgestan >= 2.7.0", "stanio >= 0.5.1"]
 pymc = ["pymc >= 5.20.1", "numba >= 0.60.0"]
 pymc-jax = ["pymc >= 5.20.1", "jax >= 0.4.27"]
 nnflow = ["flowjax >= 17.1.0", "equinox >= 0.11.12"]
 dev = [
-  "bridgestan >= 2.6.1",
+  "bridgestan >= 2.7.0",
   "stanio >= 0.5.1",
   "pymc >= 5.20.1",
   "numba >= 0.60.0",
@@ -44,7 +45,7 @@ dev = [
   "pytest-arraydiff",
 ]
 all = [
-  "bridgestan >= 2.6.1",
+  "bridgestan >= 2.7.0",
   "stanio >= 0.5.1",
   "pymc >= 5.20.1",
   "numba >= 0.60.0",
@@ -76,7 +77,7 @@ features = ["pyo3/extension-module"]
 
 [tool.pytest.ini_options]
 markers = [
-    "flow: tests for normalizing flows",
-    "stan: tests for Stan models",
-    "pymc: tests for PyMC models",
+  "flow: tests for normalizing flows",
+  "stan: tests for Stan models",
+  "pymc: tests for PyMC models",
 ]

pyproject.toml

@@ -111,6 +111,7 @@ class CompiledPyMCModel(CompiledModel):
     _n_dim: int
     _shapes: dict[str, tuple[int, ...]]
     _coords: Optional[dict[str, Any]]
+    _transform_adapt_args: dict | None = None
 
     @property
     def n_dim(self):
@@ -146,13 +147,14 @@ def with_data(self, **updates):
             user_data=user_data,
         )
 
-    def _make_sampler(self, settings, init_mean, cores, progress_type):
+    def _make_sampler(self, settings, init_mean, cores, progress_type, store):
         model = self._make_model(init_mean)
         return _lib.PySampler.from_pymc(
             settings,
             cores,
             model,
             progress_type,
+            store,
         )
 
     def _make_model(self, init_mean):
@@ -164,24 +166,46 @@ def _make_model(self, init_mean):
             self,
         )
         logp_fn = _lib.LogpFunc(
-            self.n_dim,
             self.compiled_logp_func.address,
             self.user_data.ctypes.data,
             self,
         )
 
-        var_sizes = [prod(shape) for shape in self.shape_info[2]]
         var_names = self.shape_info[0]
 
+        coords = self._coords.copy() if self._coords is not None else {}
+        dim_sizes = {name: len(vals) for name, vals in coords.items()}
+        dims = self.dims.copy() if self.dims is not None else {}
+        var_types = ["float64"] * len(var_names)
+        var_shapes = self.shape_info[2]
+
+        variables = _lib.PyVariable.new_variables(
+            var_names, var_types, var_shapes, dim_sizes, dims
+        )
+
+        outer_kwargs = self._transform_adapt_args
+        if outer_kwargs is None:
+            outer_kwargs = {}
+
+        def make_adapter(*args, **kwargs):
+            from nutpie.transform_adapter import make_transform_adapter
+
+            return make_transform_adapter(**outer_kwargs)(*args, **kwargs, logp_fn=None)
+
         return _lib.PyMcModel(
-            self.n_dim,
             logp_fn,
             expand_fn,
+            variables,
+            self.n_dim,
+            dim_sizes,
+            coords,
             self.initial_point_func,
-            var_sizes,
-            var_names,
+            make_adapter,
         )
 
+    def with_transform_adapt(self, **kwargs):
+        return dataclasses.replace(self, _transform_adapt_args=kwargs)
+
 
 def update_user_data(user_data, user_data_storage):
     user_data = user_data[()]

python/nutpie/compile_pymc.py

@@ -52,13 +52,25 @@ def make_adapter(*args, **kwargs):
 
             return make_transform_adapter(**outer_kwargs)(*args, **kwargs, logp_fn=None)
 
-        model = _lib.StanModel(self.library, seed, data_json, make_adapter)
+        coords = self._coords
+        if coords is None:
+            coords = {}
+        coords = coords.copy()
+
+        dims = self.dims
+        if dims is None:
+            dims = {}
+        dims = dims.copy()
+        dim_sizes = {name: len(dim) for name, dim in coords.items()}
+
+        model = _lib.StanModel(
+            self.library, dim_sizes, dims, coords, seed, data_json, make_adapter
+        )
         coords = self._coords
         if coords is None:
             coords = {}
         else:
             coords = coords.copy()
-        coords["unconstrained_parameter"] = pd.Index(model.param_unc_names())
 
         return CompiledStanModel(
             _coords=coords,
@@ -93,13 +105,14 @@ def _make_model(self, init_mean):
             return self.with_data().model
         return self.model
 
-    def _make_sampler(self, settings, init_mean, cores, progress_type):
+    def _make_sampler(self, settings, init_mean, cores, progress_type, store):
         model = self._make_model(init_mean)
         return _lib.PySampler.from_stan(
             settings,
             cores,
             model,
             progress_type,
+            store,
         )
 
     @property

python/nutpie/compile_stan.py

@@ -19,6 +19,7 @@ class PyFuncModel(CompiledModel):
     _shared_data: dict[str, Any]
     _n_dim: int
     _variables: list[_lib.PyVariable]
+    _dim_sizes: dict[str, int]
     _coords: dict[str, Any]
     _raw_logp_fn: Callable | None
     _transform_adapt_args: dict | None = None
@@ -47,13 +48,14 @@ def with_data(self, **updates):
     def with_transform_adapt(self, **kwargs):
         return dataclasses.replace(self, _transform_adapt_args=kwargs)
 
-    def _make_sampler(self, settings, init_mean, cores, progress_type):
+    def _make_sampler(self, settings, init_mean, cores, progress_type, store):
         model = self._make_model(init_mean)
         return _lib.PySampler.from_pyfunc(
             settings,
             cores,
             model,
             progress_type,
+            store,
         )
 
     def _make_model(self, init_mean):
@@ -85,6 +87,8 @@ def make_adapter(*args, **kwargs):
             make_expand_func,
             self._variables,
             self.n_dim,
+            dim_sizes=self._dim_sizes,
+            coords=self._coords,
             init_point_func=self._make_initial_points,
             transform_adapter=make_adapter,
         )
@@ -105,30 +109,30 @@ def from_pyfunc(
     make_transform_adapter=None,
     raw_logp_fn=None,
 ):
-    variables = []
-    for name, shape, dtype in zip(
-        expanded_names, expanded_shapes, expanded_dtypes, strict=True
-    ):
-        shape = _lib.TensorShape(list(shape))
-        if dtype == np.float64:
-            dtype = _lib.ExpandDtype.float64_array(shape)
-        elif dtype == np.float32:
-            dtype = _lib.ExpandDtype.float32_array(shape)
-        elif dtype == np.int64:
-            dtype = _lib.ExpandDtype.int64_array(shape)
-        variables.append(_lib.PyVariable(name, dtype))
-
     if coords is None:
         coords = {}
     if dims is None:
         dims = {}
     if shared_data is None:
         shared_data = {}
 
+    coords = coords.copy()
+
+    dim_sizes = {k: len(v) for k, v in coords.items()}
+    shapes = [tuple(shape) for shape in expanded_shapes]
+    variables = _lib.PyVariable.new_variables(
+        expanded_names,
+        [str(dtype) for dtype in expanded_dtypes],
+        shapes,
+        dim_sizes,
+        dims,
+    )
+
     return PyFuncModel(
         _n_dim=ndim,
         dims=dims,
         _coords=coords,
+        _dim_sizes=dim_sizes,
         _make_logp_func=make_logp_fn,
         _make_expand_func=make_expand_fn,
         _make_initial_points=make_initial_point_fn,