Repository: pymc-devs/nutpie (language: python)

Issue:
`__version__` is missing



Discussion / hints:

Fix patch:
diff --git a/Cargo.toml b/Cargo.toml
index 287a2f6..77db80d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,7 +17,7 @@ default = ["extension-module"]
 simd_support = ["nuts-rs/simd_support"]
 
 [lib]
-name = "nutpie"
+name = "_lib"
 crate-type = ["cdylib"]
 
 [dependencies]
diff --git a/pyproject.toml b/pyproject.toml
index 3a5f093..4c252d0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,6 +3,8 @@ requires = ["maturin>=1.1,<2.0"]
 build-backend = "maturin"
 
 [tool.maturin]
+module-name = "nutpie._lib"
+python-source = "python"
 features = ["pyo3/extension-module"]
 
 [project]
diff --git a/nutpie/__init__.py b/python/nutpie/__init__.py
similarity index 72%
rename from nutpie/__init__.py
rename to python/nutpie/__init__.py
index 996a709..e9c6f60 100644
--- a/nutpie/__init__.py
+++ b/python/nutpie/__init__.py
@@ -1,4 +1,5 @@
-from . import nutpie as lib
+from nutpie import _lib
+from nutpie.sample import sample
 
 try:
     from .compile_pymc import compile_pymc_model
@@ -16,4 +17,7 @@ def compile_stan_model(*args, **kwargs):
         raise ImportError("Missing dependencies for stan models. Install bridgestan.")
 
 
-from .sample import sample
+__version__ = _lib.__version__
+
+
+__all__ = ["sample", "compile_pymc_model", "compile_stan_model"]
diff --git a/nutpie/compile_pymc.py b/python/nutpie/compile_pymc.py
similarity index 94%
rename from nutpie/compile_pymc.py
rename to python/nutpie/compile_pymc.py
index bb5bd44..2342b41 100644
--- a/nutpie/compile_pymc.py
+++ b/python/nutpie/compile_pymc.py
@@ -1,5 +1,4 @@
 import dataclasses
-import functools
 import itertools
 from dataclasses import dataclass
 from math import prod
@@ -17,8 +16,8 @@
 from numba.cpython.unsafe.tuple import alloca_once, tuple_setitem
 from numpy.typing import NDArray
 
-from . import lib
-from .sample import CompiledModel
+from nutpie import _lib
+from nutpie.sample import CompiledModel
 
 
 @numba.extending.intrinsic
@@ -84,17 +83,17 @@ def with_data(self, **updates):
 
     def _make_sampler(self, settings, init_mean, chains, cores, seed):
         model = self._make_model(init_mean)
-        return lib.PySampler.from_pymc(settings, chains, cores, model, seed)
+        return _lib.PySampler.from_pymc(settings, chains, cores, model, seed)
 
     def _make_model(self, init_mean):
-        expand_fn = lib.ExpandFunc(
+        expand_fn = _lib.ExpandFunc(
             self.n_dim,
             self.n_expanded,
             self.compiled_expand_func.address,
             self.user_data.ctypes.data,
             self,
         )
-        logp_fn = lib.LogpFunc(
+        logp_fn = _lib.LogpFunc(
             self.n_dim,
             self.compiled_logp_func.address,
             self.user_data.ctypes.data,
@@ -103,7 +102,7 @@ def _make_model(self, init_mean):
 
         var_sizes = [prod(shape) for shape in self.shape_info[2]]
 
-        return lib.PyMcModel(
+        return _lib.PyMcModel(
             self.n_dim,
             logp_fn,
             expand_fn,
@@ -307,8 +306,6 @@ def _make_functions(model):
     }
 
     (logp, grad) = pytensor.graph_replace([logp, grad], replacements)
-    # (logp, grad) = pytensor.graph.rewrite_graph(logp, include=["canonicalize", "stabilize"])
-    # grad = pytensor.gradient.grad(logp, joined)
 
     # We should avoid compiling the function, and optimize only
     logp_fn_pt = pytensor.compile.function.function(
@@ -318,12 +315,6 @@ def _make_functions(model):
     logp_fn = logp_fn_pt.vm.jit_fn
 
     # Make function that computes remaining variables for the trace
-    trace_vars = {
-        name: var
-        for (name, var) in model.named_vars.items()
-        if var not in model.observed_RVs + model.potentials
-    }
-    remaining_names = [name for name in trace_vars if name not in joined_names]
     remaining_rvs = [
         var for var in model.unobserved_value_vars if var.name not in joined_names
     ]
@@ -352,9 +343,6 @@ def _make_functions(model):
         mode=pytensor.compile.NUMBA,
     )
     expand_fn = expand_fn_pt.vm.jit_fn
-    # expand_fn = numba.njit(expand_fn, fastmath=True, error_model="numpy")
-    # Trigger a compile
-    # expand_fn(np.zeros(num_free_vars), *[var.get_value() for var in expand_fn_pt.get_shared()])
 
     return (
         num_free_vars,
@@ -428,8 +416,6 @@ def impl_extract_array(user_data, index):
 
         name, ndim, base_shape, dtype = shared_metadata[index]
 
-        ndim_range = tuple(range(ndim))
-
         def impl(user_data, index):
             data_ptr = address_as_void_pointer(user_data["data"][name][()])
             data = numba.carray(data_ptr, int(user_data["size"][name][()]), dtype)
diff --git a/nutpie/compile_stan.py b/python/nutpie/compile_stan.py
similarity index 95%
rename from nutpie/compile_stan.py
rename to python/nutpie/compile_stan.py
index 2b5236b..8cdcefd 100644
--- a/nutpie/compile_stan.py
+++ b/python/nutpie/compile_stan.py
@@ -8,7 +8,7 @@
 import pandas as pd
 from numpy.typing import NDArray
 
-from nutpie import lib
+from nutpie import _lib
 from nutpie.sample import CompiledModel
 
 
@@ -41,7 +41,7 @@ def with_data(self, *, seed=None, **updates):
         else:
             data_json = None
 
-        model = lib.StanModel(self.library, seed, data_json)
+        model = _lib.StanModel(self.library, seed, data_json)
         coords = self._coords
         if coords is None:
             coords = {}
@@ -81,7 +81,7 @@ def _make_model(self, init_mean):
 
     def _make_sampler(self, settings, init_mean, chains, cores, seed):
         model = self._make_model(init_mean)
-        return lib.PySampler.from_stan(settings, chains, cores, model, seed)
+        return _lib.PySampler.from_stan(settings, chains, cores, model, seed)
 
     @property
     def n_dim(self):
@@ -145,7 +145,7 @@ def compile_stan_model(
         so_path = bridgestan.compile_model(model_path, make_args=make_args)
         # Set necessary library loading paths
         bridgestan.compile.windows_dll_path_setup()
-        library = lib.StanLibrary(so_path)
+        library = _lib.StanLibrary(so_path)
     finally:
         try:
             if cleanup:
diff --git a/nutpie/sample.py b/python/nutpie/sample.py
similarity index 98%
rename from nutpie/sample.py
rename to python/nutpie/sample.py
index b517f49..8a1c710 100644
--- a/nutpie/sample.py
+++ b/python/nutpie/sample.py
@@ -6,9 +6,8 @@
 import numpy as np
 import pandas as pd
 import pyarrow
-import xarray as xr
 
-from . import lib
+from nutpie import _lib
 
 
 @dataclass(frozen=True)
@@ -193,14 +192,14 @@ def sample(
         Return the raw trace object (an apache arrow structure)
         instead of converting to arviz.
     **kwargs
-        Pass additional arguments to nutpie.lib.PySamplerArgs
+        Pass additional arguments to nutpie._lib.PySamplerArgs
 
     Returns
     -------
     trace : arviz.InferenceData
         An ArviZ ``InferenceData`` object that contains the samples.
     """
-    settings = lib.PySamplerArgs()
+    settings = _lib.PySamplerArgs()
     settings.num_tune = tune
     settings.num_draws = draws
 
diff --git a/src/lib.rs b/src/lib.rs
index a058f4d..b56963e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,4 +3,4 @@ mod sampler;
 mod stan;
 mod wrapper;
 
-pub use wrapper::nutpie;
+pub use wrapper::_lib;
diff --git a/src/wrapper.rs b/src/wrapper.rs
index e652d4a..ccf602d 100644
--- a/src/wrapper.rs
+++ b/src/wrapper.rs
@@ -327,7 +327,7 @@ fn export_array(py: Python<'_>, name: String, data: Box<dyn Array>) -> PyResult<
 
 /// A Python module implemented in Rust.
 #[pymodule]
-pub fn nutpie(_py: Python, m: &PyModule) -> PyResult<()> {
+pub fn _lib(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<PySamplerArgs>()?;
     m.add_class::<PySampleStats>()?;
     m.add_class::<PySampler>()?;
@@ -336,5 +336,6 @@ pub fn nutpie(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<ExpandFunc>()?;
     m.add_class::<StanLibrary>()?;
     m.add_class::<StanModel>()?;
+    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
     Ok(())
 }


Test patch:
diff --git a/tests/test_pymc.py b/tests/test_pymc.py
index b0b2bc1..8623bcd 100644
--- a/tests/test_pymc.py
+++ b/tests/test_pymc.py
@@ -27,7 +27,7 @@ def test_trafo():
 def test_det():
     with pm.Model() as model:
         a = pm.Uniform("a", shape=2)
-        b = pm.Deterministic("b", 2 * a)
+        pm.Deterministic("b", 2 * a)
 
     compiled = nutpie.compile_pymc_model(model)
     trace = nutpie.sample(compiled, chains=1)
@@ -51,4 +51,4 @@ def test_pymc_model_shared():
 
     compiled3 = compiled.with_data(mu=0.5, sigma=3 * np.ones(4))
     with pytest.raises(ValueError):
-        trace3 = nutpie.sample(compiled3, chains=1)
+        nutpie.sample(compiled3, chains=1)