feat: Use support_point as default init for pymc

Author: aseyboldt

python/nutpie/compile_pymc.py

@@ -435,9 +435,12 @@ def compile_pymc_model(
     *,
     backend: Literal["numba", "jax"] = "numba",
     gradient_backend: Literal["pytensor", "jax"] = "pytensor",
-    overrides: dict[Union["Variable", str], np.ndarray | float | int] | None = None,
+    initial_points: dict[Union["Variable", str], np.ndarray | float | int]
+    | None = None,
     jitter_rvs: set["TensorVariable"] | None = None,
-    default_strategy: Literal["support_point", "prior"] = "prior",
+    default_initialization_strategy: Literal[
+        "support_point", "prior"
+    ] = "support_point",
     **kwargs,
 ) -> CompiledModel:
     """Compile necessary functions for sampling a pymc model.
@@ -455,10 +458,10 @@ def compile_pymc_model(
         The set (or list or tuple) of random variables for which a U(-1, +1)
         jitter should be added to the initial value. Only available for
         variables that have a transform or real-valued support.
-    default_strategy : str
+    default_initialization_strategy : str
         Which of { "support_point", "prior" } to prefer if the initval setting
         for an RV is None.
-    overrides : dict
+    initial_points : dict
         Initial value (strategies) to use instead of what's specified in
         `Model.initial_values`.
     Returns
@@ -475,13 +478,13 @@ def compile_pymc_model(
             "and restart your kernel in case you are in an interactive session."
         )
 
-    if default_strategy == "support_point" and jitter_rvs is None:
+    if default_initialization_strategy == "support_point" and jitter_rvs is None:
         jitter_rvs = set(model.free_RVs)
 
     initial_point_fn = make_initial_point_fn(
         model=model,
-        overrides=overrides,
-        default_strategy=default_strategy,
+        overrides=initial_points,
+        default_strategy=default_initialization_strategy,
         jitter_rvs=jitter_rvs,
         return_transformed=True,
     )

python/nutpie/compiled_pyfunc.py

@@ -8,14 +8,14 @@
 from nutpie import _lib
 from nutpie.sample import CompiledModel
 
-SeedType = int | float | np.random.Generator | None
+SeedType = int
 
 
 @dataclass(frozen=True)
 class PyFuncModel(CompiledModel):
     _make_logp_func: Callable
     _make_expand_func: Callable
-    _make_initial_points: Callable[[SeedType], np.ndarray]
+    _make_initial_points: Callable[[SeedType], np.ndarray] | None
     _shared_data: dict[str, Any]
     _n_dim: int
     _variables: list[_lib.PyVariable]
@@ -73,14 +73,14 @@ def from_pyfunc(
     ndim: int,
     make_logp_fn: Callable,
     make_expand_fn: Callable,
-    make_initial_point_fn: Callable[[SeedType], np.ndarray],
     expanded_dtypes: list[np.dtype],
     expanded_shapes: list[tuple[int, ...]],
     expanded_names: list[str],
     *,
     coords: dict[str, Any] | None = None,
     dims: dict[str, tuple[str, ...]] | None = None,
     shared_data: dict[str, Any] | None = None,
+    make_initial_point_fn: Callable[[SeedType], np.ndarray] | None,
 ):
     variables = []
     for name, shape, dtype in zip(

tests/test_pymc.py

@@ -38,6 +38,19 @@ def test_pymc_model_float32(backend, gradient_backend):
         trace.posterior.a  # noqa: B018
 
 
+@parameterize_backends
+def test_pymc_model_no_prior(backend, gradient_backend):
+    with pm.Model() as model:
+        a = pm.Flat("a")
+        pm.Normal("b", mu=a, observed=0.0)
+
+    compiled = nutpie.compile_pymc_model(
+        model, backend=backend, gradient_backend=gradient_backend
+    )
+    trace = nutpie.sample(compiled, chains=1)
+    trace.posterior.a  # noqa: B018
+
+
 @parameterize_backends
 def test_blocking(backend, gradient_backend):
     with pm.Model() as model: