Finalize inital point function inside _make_funcs

Author: jessegrabowski

python/nutpie/compile_pymc.py

@@ -66,7 +66,7 @@ def seeded_array_fn(seed: SeedType = None):
         flat_array = np.empty(total_size, dtype="float64", order="C")
         cursor = 0
 
-        for name, shape in zip(names, shapes):
+        for name, shape in zip(names, shapes, strict=True):
             initial_value = initial_value_dict[name]
             n = int(np.prod(initial_value.shape))
             if initial_value.shape != shape:
@@ -217,7 +217,7 @@ def make_user_data(shared_vars, shared_data):
 
 def _compile_pymc_model_numba(
     model: "pm.Model",
-    initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]],
+    pymc_initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]],
     **kwargs,
 ) -> CompiledPyMCModel:
     if find_spec("numba") is None:
@@ -234,8 +234,15 @@ def _compile_pymc_model_numba(
         n_expanded,
         logp_fn_pt,
         expand_fn_pt,
+        initial_point_fn,
         shape_info,
-    ) = _make_functions(model, mode="NUMBA", compute_grad=True, join_expanded=True)
+    ) = _make_functions(
+        model,
+        mode="NUMBA",
+        compute_grad=True,
+        join_expanded=True,
+        pymc_initial_point_fn=pymc_initial_point_fn,
+    )
 
     expand_fn = expand_fn_pt.vm.jit_fn
     logp_fn = logp_fn_pt.vm.jit_fn
@@ -282,17 +289,15 @@ def _compile_pymc_model_numba(
         expand_numba = numba.cfunc(c_sig_expand, **kwargs)(expand_numba_raw)
 
     dims, coords = _prepare_dims_and_coords(model, shape_info)
-    initial_point_fn_array = rv_dict_to_flat_array_wrapper(
-        initial_point_fn, names=shape_info[0], shapes=shape_info[-1]
-    )
+
     return CompiledPyMCModel(
         _n_dim=n_dim,
         dims=dims,
         _coords=coords,
         _shapes={name: tuple(shape) for name, _, shape in zip(*shape_info)},
         compiled_logp_func=logp_numba,
         compiled_expand_func=expand_numba,
-        initial_point_func=initial_point_fn_array,
+        initial_point_func=initial_point_fn,
         shared_data=shared_data,
         user_data=user_data,
         n_expanded=n_expanded,
@@ -331,7 +336,7 @@ def _compile_pymc_model_jax(
     model,
     *,
     gradient_backend=None,
-    initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]],
+    pymc_initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]],
     **kwargs,
 ):
     if find_spec("jax") is None:
@@ -353,12 +358,14 @@ def _compile_pymc_model_jax(
         _,
         logp_fn_pt,
         expand_fn_pt,
+        initial_point_fn,
         shape_info,
     ) = _make_functions(
         model,
         mode="JAX",
         compute_grad=gradient_backend == "pytensor",
         join_expanded=False,
+        pymc_initial_point_fn=pymc_initial_point_fn,
     )
 
     logp_fn = logp_fn_pt.vm.jit_fn
@@ -409,15 +416,11 @@ def expand(x, **shared):
 
     dims, coords = _prepare_dims_and_coords(model, shape_info)
 
-    initial_point_fn_array = rv_dict_to_flat_array_wrapper(
-        initial_point_fn, names=shape_info[0], shapes=shape_info[-1]
-    )
-
     return from_pyfunc(
         ndim=n_dim,
         make_logp_fn=make_logp_func,
         make_expand_fn=make_expand_func,
-        make_initial_point_fn=initial_point_fn_array,
+        make_initial_point_fn=initial_point_fn,
         expanded_dtypes=dtypes,
         expanded_shapes=shapes,
         expanded_names=names,
@@ -484,13 +487,13 @@ def compile_pymc_model(
         if gradient_backend == "jax":
             raise ValueError("Gradient backend cannot be jax when using numba backend")
         return _compile_pymc_model_numba(
-            model, initial_point_fn=initial_point_fn, **kwargs
+            model=model, pymc_initial_point_fn=initial_point_fn, **kwargs
         )
     elif backend.lower() == "jax":
         return _compile_pymc_model_jax(
-            model,
+            model=model,
             gradient_backend=gradient_backend,
-            initial_point_fn=initial_point_fn,
+            pymc_initial_point_fn=initial_point_fn,
             **kwargs,
         )
     else:
@@ -527,9 +530,19 @@ def _compute_shapes(model):
 
 
 def _make_functions(
-    model, *, mode, compute_grad, join_expanded
+    model: "pm.Model",
+    *,
+    mode: Literal["JAX", "NUMBA"],
+    compute_grad: bool,
+    join_expanded: bool,
+    pymc_initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]],
 ) -> tuple[
-    int, int, Callable, Callable, tuple[list[str], list[slice], list[tuple[int, ...]]]
+    int,
+    int,
+    Callable,
+    Callable,
+    Callable,
+    tuple[list[str], list[slice], list[tuple[int, ...]]],
 ]:
     """
     Compile functions required by nuts-rs from a given PyMC model.
@@ -546,6 +559,8 @@ def _make_functions(
     join_expanded: bool
         Whether to join the expanded variables into a single array. If False, the expanded variables will be returned
         as a list of arrays.
+    pymc_initial_point_fn: Callable
+        Initial point function created by pymc.initial_point.make_initial_point_fn
 
     Returns
     -------
@@ -558,6 +573,8 @@ def _make_functions(
         and the gradient, otherwise only the logp is returned.
     expand_fn_pt: Callable
         Compiled pytensor function that computes the remaining variables for the trace
+    initial_point_fn: Callable
+        Python function that takes a random seed and returns a flat array of initial values
     param_data: tuple of lists
         Tuple containing data necessary to unravel a flat array of model variables back into a ragged list of arrays.
         The first list contains the names of the variables, the second list contains the slices that correspond to the
@@ -607,6 +624,10 @@ def _make_functions(
 
     num_free_vars = count
 
+    initial_point_fn = rv_dict_to_flat_array_wrapper(
+        pymc_initial_point_fn, names=joined_names, shapes=joined_shapes
+    )
+
     joined = pt.TensorType("float64", shape=(num_free_vars,))(
         name="_unconstrained_point"
     )
@@ -673,6 +694,7 @@ def _make_functions(
         num_expanded,
         logp_fn_pt,
         expand_fn_pt,
+        initial_point_fn,
         (all_names, all_slices, all_shapes),
     )