Fix bug introduced by code refactor

Author: jessegrabowski

python/nutpie/compile_pymc.py

@@ -61,19 +61,20 @@ def rv_dict_to_flat_array_wrapper(
 
     @wraps(fn)
     def seeded_array_fn(seed: SeedType = None):
-        inital_value_dict = fn(seed)
-        total_size = sum(np.prod(shape) for shape in shapes)
+        initial_value_dict = fn(seed)
+        total_size = sum(np.prod(shape).astype(int) for shape in shapes)
         flat_array = np.empty(total_size, dtype="float64", order="C")
         cursor = 0
 
         for name, shape in zip(names, shapes):
-            initial_value = inital_value_dict[name]
+            initial_value = initial_value_dict[name]
             n = int(np.prod(initial_value.shape))
             if initial_value.shape != shape:
                 raise ValueError(
                     f"Size of initial value for {name} is {initial_value.shape}, "
                     f"expected {shape}"
                 )
+
             flat_array[cursor : cursor + n] = initial_value.ravel().astype("float64")
             cursor += n
 
@@ -144,16 +145,16 @@ def with_data(self, **updates):
             user_data=user_data,
         )
 
-    def _make_sampler(self, settings, init_mean, cores, progress_type):
-        model = self._make_model(init_mean)
+    def _make_sampler(self, settings, cores, progress_type):
+        model = self._make_model()
         return _lib.PySampler.from_pymc(
             settings,
             cores,
             model,
             progress_type,
         )
 
-    def _make_model(self, init_mean):
+    def _make_model(self):
         expand_fn = _lib.ExpandFunc(
             self.n_dim,
             self.n_expanded,
@@ -169,14 +170,15 @@ def _make_model(self, init_mean):
         )
 
         var_sizes = [prod(shape) for shape in self.shape_info[2]]
+        var_names = self.shape_info[0]
 
         return _lib.PyMcModel(
             self.n_dim,
             logp_fn,
             expand_fn,
             self.initial_point_func,
             var_sizes,
-            self.shape_info[0],
+            var_names,
         )
 
 
@@ -472,7 +474,7 @@ def compile_pymc_model(
         overrides=overrides,
         default_strategy=default_strategy,
         jitter_rvs=jitter_rvs,
-        return_transformed=False,
+        return_transformed=True,
     )
 
     if backend.lower() == "numba":

python/nutpie/compiled_pyfunc.py

@@ -73,7 +73,7 @@ def from_pyfunc(
     ndim: int,
     make_logp_fn: Callable,
     make_expand_fn: Callable,
-    make_initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]],
+    make_initial_point_fn: Callable[[SeedType], np.ndarray],
     expanded_dtypes: list[np.dtype],
     expanded_shapes: list[tuple[int, ...]],
     expanded_names: list[str],
@@ -102,19 +102,13 @@ def from_pyfunc(
     if shared_data is None:
         shared_data = {}
 
-    from nutpie.compile_pymc import rv_dict_to_flat_array_wrapper
-
-    initial_point_fn = rv_dict_to_flat_array_wrapper(
-        make_initial_point_fn, names=expanded_names, shapes=expanded_shapes
-    )
-
     return PyFuncModel(
         _n_dim=ndim,
         dims=dims,
         _coords=coords,
         _make_logp_func=make_logp_fn,
         _make_expand_func=make_expand_fn,
-        _make_initial_points=initial_point_fn,
+        _make_initial_points=make_initial_point_fn,
         _variables=variables,
         _shared_data=shared_data,
     )