pymc-extras 0.2.2__tar.gz → 0.2.4__tar.gz

{pymc_extras-0.2.2 → pymc_extras-0.2.4}/PKG-INFO

@@ -1,11 +1,11 @@
 Metadata-Version: 2.2
 Name: pymc-extras
-Version: 0.2.2
+Version: 0.2.4
 Summary: A home for new additions to PyMC, which may include unusual probability distribitions, advanced model fitting algorithms, or any code that may be inappropriate to include in the pymc repository, but may want to be made available to users.
 Home-page: http://github.com/pymc-devs/pymc-extras
 Maintainer: PyMC Developers
 Maintainer-email: pymc.devs@gmail.com
-License: Apache License, Version 2.0
+License: Apache-2.0
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
@@ -20,8 +20,9 @@ Classifier: Operating System :: OS Independent
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: pymc>=5.19.1
+Requires-Dist: pymc>=5.21.1
 Requires-Dist: scikit-learn
+Requires-Dist: better-optimize
 Provides-Extra: dask-histogram
 Requires-Dist: dask[complete]; extra == "dask-histogram"
 Requires-Dist: xhistogram; extra == "dask-histogram"

{pymc_extras-0.2.2 → pymc_extras-0.2.4}/pymc_extras/__init__.py

@@ -15,7 +15,9 @@ import logging
 
 from pymc_extras import gp, statespace, utils
 from pymc_extras.distributions import *
+from pymc_extras.inference.find_map import find_MAP
 from pymc_extras.inference.fit import fit
+from pymc_extras.inference.laplace import fit_laplace
 from pymc_extras.model.marginal.marginal_model import (
     MarginalModel,
     marginalize,

{pymc_extras-0.2.2 → pymc_extras-0.2.4}/pymc_extras/inference/find_map.py

@@ -1,9 +1,9 @@
 import logging
 
 from collections.abc import Callable
+from importlib.util import find_spec
 from typing import Literal, cast, get_args
 
-import jax
 import numpy as np
 import pymc as pm
 import pytensor
@@ -30,13 +30,29 @@ VALID_BACKENDS = get_args(GradientBackend)
 def set_optimizer_function_defaults(method, use_grad, use_hess, use_hessp):
     method_info = MINIMIZE_MODE_KWARGS[method].copy()
 
-    use_grad = use_grad if use_grad is not None else method_info["uses_grad"]
-    use_hess = use_hess if use_hess is not None else method_info["uses_hess"]
-    use_hessp = use_hessp if use_hessp is not None else method_info["uses_hessp"]
-
     if use_hess and use_hessp:
+        _log.warning(
+            'Both "use_hess" and "use_hessp" are set to True, but scipy.optimize.minimize never uses both at the '
+            'same time. When possible "use_hessp" is preferred because its is computationally more efficient. '
+            'Setting "use_hess" to False.'
+        )
         use_hess = False
 
+    use_grad = use_grad if use_grad is not None else method_info["uses_grad"]
+
+    if use_hessp is not None and use_hess is None:
+        use_hess = not use_hessp
+
+    elif use_hess is not None and use_hessp is None:
+        use_hessp = not use_hess
+
+    elif use_hessp is None and use_hess is None:
+        use_hessp = method_info["uses_hessp"]
+        use_hess = method_info["uses_hess"]
+        if use_hessp and use_hess:
+            # If a method could use either hess or hessp, we default to using hessp
+            use_hess = False
+
     return use_grad, use_hess, use_hessp
 
 
@@ -59,7 +75,7 @@ def get_nearest_psd(A: np.ndarray) -> np.ndarray:
         The nearest positive semi-definite matrix to the input matrix.
     """
     C = (A + A.T) / 2
-    eigval, eigvec = np.linalg.eig(C)
+    eigval, eigvec = np.linalg.eigh(C)
     eigval[eigval < 0] = 0
 
     return eigvec @ np.diag(eigval) @ eigvec.T
@@ -97,7 +113,7 @@ def _create_transformed_draws(H_inv, slices, out_shapes, posterior_draws, model,
     return f_untransform(posterior_draws)
 
 
-def _compile_jax_gradients(
+def _compile_grad_and_hess_to_jax(
     f_loss: Function, use_hess: bool, use_hessp: bool
 ) -> tuple[Callable | None, Callable | None]:
     """
@@ -122,6 +138,8 @@ def _compile_jax_gradients(
     f_hessp: Callable | None
         The compiled hessian-vector product function, or None if use_hessp is False.
     """
+    import jax
+
     f_hess = None
     f_hessp = None
 
@@ -152,7 +170,7 @@ def _compile_jax_gradients(
     return f_loss_and_grad, f_hess, f_hessp
 
 
-def _compile_functions(
+def _compile_functions_for_scipy_optimize(
     loss: TensorVariable,
     inputs: list[TensorVariable],
     compute_grad: bool,
@@ -177,7 +195,7 @@ def _compile_functions(
     compute_hessp: bool
         Whether to compile a function that computes the Hessian-vector product of the loss function.
     compile_kwargs: dict, optional
-        Additional keyword arguments to pass to the ``pm.compile_pymc`` function.
+        Additional keyword arguments to pass to the ``pm.compile`` function.
 
     Returns
     -------
@@ -193,19 +211,19 @@ def _compile_functions(
     if compute_grad:
         grads = pytensor.gradient.grad(loss, inputs)
         grad = pt.concatenate([grad.ravel() for grad in grads])
-        f_loss_and_grad = pm.compile_pymc(inputs, [loss, grad], **compile_kwargs)
+        f_loss_and_grad = pm.compile(inputs, [loss, grad], **compile_kwargs)
     else:
-        f_loss = pm.compile_pymc(inputs, loss, **compile_kwargs)
+        f_loss = pm.compile(inputs, loss, **compile_kwargs)
         return [f_loss]
 
     if compute_hess:
         hess = pytensor.gradient.jacobian(grad, inputs)[0]
-        f_hess = pm.compile_pymc(inputs, hess, **compile_kwargs)
+        f_hess = pm.compile(inputs, hess, **compile_kwargs)
 
     if compute_hessp:
         p = pt.tensor("p", shape=inputs[0].type.shape)
         hessp = pytensor.gradient.hessian_vector_product(loss, inputs, p)
-        f_hessp = pm.compile_pymc([*inputs, p], hessp[0], **compile_kwargs)
+        f_hessp = pm.compile([*inputs, p], hessp[0], **compile_kwargs)
 
     return [f_loss_and_grad, f_hess, f_hessp]
 
@@ -240,7 +258,7 @@ def scipy_optimize_funcs_from_loss(
     gradient_backend: str, default "pytensor"
         Which backend to use to compute gradients. Must be one of "jax" or "pytensor"
     compile_kwargs:
-        Additional keyword arguments to pass to the ``pm.compile_pymc`` function.
+        Additional keyword arguments to pass to the ``pm.compile`` function.
 
     Returns
     -------
@@ -265,6 +283,8 @@ def scipy_optimize_funcs_from_loss(
         )
 
     use_jax_gradients = (gradient_backend == "jax") and use_grad
+    if use_jax_gradients and not find_spec("jax"):
+        raise ImportError("JAX must be installed to use JAX gradients")
 
     mode = compile_kwargs.get("mode", None)
     if mode is None and use_jax_gradients:
@@ -285,7 +305,7 @@ def scipy_optimize_funcs_from_loss(
     compute_hess = use_hess and not use_jax_gradients
     compute_hessp = use_hessp and not use_jax_gradients
 
-    funcs = _compile_functions(
+    funcs = _compile_functions_for_scipy_optimize(
         loss=loss,
         inputs=[flat_input],
         compute_grad=compute_grad,
@@ -301,7 +321,7 @@ def scipy_optimize_funcs_from_loss(
 
     if use_jax_gradients:
         # f_loss here is f_loss_and_grad; the name is unchanged to simplify the return values
-        f_loss, f_hess, f_hessp = _compile_jax_gradients(f_loss, use_hess, use_hessp)
+        f_loss, f_hess, f_hessp = _compile_grad_and_hess_to_jax(f_loss, use_hess, use_hessp)
 
     return f_loss, f_hess, f_hessp
 

{pymc_extras-0.2.2 → pymc_extras-0.2.4}/pymc_extras/inference/laplace.py

@@ -16,6 +16,7 @@
 import logging
 
 from functools import reduce
+from importlib.util import find_spec
 from itertools import product
 from typing import Literal
 
@@ -231,7 +232,7 @@ def add_data_to_inferencedata(
     return idata
 
 
-def fit_mvn_to_MAP(
+def fit_mvn_at_MAP(
     optimized_point: dict[str, np.ndarray],
     model: pm.Model | None = None,
     on_bad_cov: Literal["warn", "error", "ignore"] = "ignore",
@@ -276,6 +277,9 @@ def fit_mvn_to_MAP(
     inverse_hessian: np.ndarray
         The inverse Hessian matrix of the log-posterior evaluated at the MAP estimate.
     """
+    if gradient_backend == "jax" and not find_spec("jax"):
+        raise ImportError("JAX must be installed to use JAX gradients")
+
     model = pm.modelcontext(model)
     compile_kwargs = {} if compile_kwargs is None else compile_kwargs
     frozen_model = freeze_dims_and_data(model)
@@ -344,8 +348,10 @@ def sample_laplace_posterior(
 
     Parameters
     ----------
-    mu
-    H_inv
+    mu: RaveledVars
+        The MAP estimate of the model parameters.
+    H_inv: np.ndarray
+        The inverse Hessian matrix of the log-posterior evaluated at the MAP estimate.
     model : Model
         A PyMC model
     chains : int
@@ -384,9 +390,7 @@ def sample_laplace_posterior(
             constrained_rvs, replace={unconstrained_vector: batched_values}
         )
 
-        f_constrain = pm.compile_pymc(
-            inputs=[batched_values], outputs=batched_rvs, **compile_kwargs
-        )
+        f_constrain = pm.compile(inputs=[batched_values], outputs=batched_rvs, **compile_kwargs)
         posterior_draws = f_constrain(posterior_draws)
 
     else:
@@ -472,15 +476,17 @@ def fit_laplace(
         and 1).
 
         .. warning::
-            This argumnet should be considered highly experimental. It has not been verified if this method produces
+            This argument should be considered highly experimental. It has not been verified if this method produces
             valid draws from the posterior. **Use at your own risk**.
 
     gradient_backend: str, default "pytensor"
         The backend to use for gradient computations. Must be one of "pytensor" or "jax".
     chains: int, default: 2
-        The number of sampling chains running in parallel.
+        The number of chain dimensions to sample. Note that this is *not* the number of chains to run in parallel,
+        because the Laplace approximation is not an MCMC method. This argument exists to ensure that outputs are
+        compatible with the ArviZ library.
     draws: int, default: 500
-        The number of samples to draw from the approximated posterior.
+        The number of samples to draw from the approximated posterior. Totals samples will be chains * draws.
     on_bad_cov : str, one of 'ignore', 'warn', or 'error', default: 'ignore'
         What to do when ``H_inv`` (inverse Hessian) is not positive semi-definite.
         If 'ignore' or 'warn', the closest positive-semi-definite matrix to ``H_inv`` (in L1 norm) will be returned.
@@ -547,11 +553,12 @@ def fit_laplace(
         **optimizer_kwargs,
     )
 
-    mu, H_inv = fit_mvn_to_MAP(
+    mu, H_inv = fit_mvn_at_MAP(
         optimized_point=optimized_point,
         model=model,
         on_bad_cov=on_bad_cov,
         transform_samples=fit_in_unconstrained_space,
+        gradient_backend=gradient_backend,
         zero_tol=zero_tol,
         diag_jitter=diag_jitter,
         compile_kwargs=compile_kwargs,

{pymc_extras-0.2.2 → pymc_extras-0.2.4}/pymc_extras/inference/pathfinder/importance_sampling.py

@@ -20,7 +20,7 @@ class ImportanceSamplingResult:
     samples: NDArray
     pareto_k: float | None = None
     warnings: list[str] = field(default_factory=list)
-    method: str = "none"
+    method: str = "psis"
 
 
 def importance_sampling(
@@ -28,7 +28,7 @@ def importance_sampling(
     logP: NDArray,
     logQ: NDArray,
     num_draws: int,
-    method: Literal["psis", "psir", "identity", "none"] | None,
+    method: Literal["psis", "psir", "identity"] | None,
     random_seed: int | None = None,
 ) -> ImportanceSamplingResult:
     """Pareto Smoothed Importance Resampling (PSIR)
@@ -44,8 +44,15 @@ def importance_sampling(
         log probability values of proposal distribution, shape (L, M)
     num_draws : int
         number of draws to return where num_draws <= samples.shape[0]
-    method : str, optional
-        importance sampling method to use. Options are "psis" (default), "psir", "identity", "none. Pareto Smoothed Importance Sampling (psis) is recommended in many cases for more stable results than Pareto Smoothed Importance Resampling (psir). identity applies the log importance weights directly without resampling. none applies no importance sampling weights and returns the samples as is of size num_draws_per_path * num_paths.
+    method : str, None, optional
+        Method to apply sampling based on log importance weights (logP - logQ).
+        Options are:
+        "psis" : Pareto Smoothed Importance Sampling (default)
+                Recommended for more stable results.
+        "psir" : Pareto Smoothed Importance Resampling
+                Less stable than PSIS.
+        "identity" : Applies log importance weights directly without resampling.
+        None : No importance sampling weights. Returns raw samples of size (num_paths, num_draws_per_path, N) where N is number of model parameters. Other methods return samples of size (num_draws, N).
     random_seed : int | None
 
     Returns
@@ -71,11 +78,11 @@ def importance_sampling(
     warnings = []
     num_paths, _, N = samples.shape
 
-    if method == "none":
+    if method is None:
         warnings.append(
             "Importance sampling is disabled. The samples are returned as is which may include samples from failed paths with non-finite logP or logQ values. It is recommended to use importance_sampling='psis' for better stability."
         )
-        return ImportanceSamplingResult(samples=samples, warnings=warnings)
+        return ImportanceSamplingResult(samples=samples, warnings=warnings, method=method)
     else:
         samples = samples.reshape(-1, N)
         logP = logP.ravel()
@@ -91,17 +98,16 @@ def importance_sampling(
             _warnings.filterwarnings(
                 "ignore", category=RuntimeWarning, message="overflow encountered in exp"
             )
-            if method == "psis":
-                replace = False
-                logiw, pareto_k = az.psislw(logiw)
-            elif method == "psir":
-                replace = True
-                logiw, pareto_k = az.psislw(logiw)
-            elif method == "identity":
-                replace = False
-                pareto_k = None
-            else:
-                raise ValueError(f"Invalid importance sampling method: {method}")
+            match method:
+                case "psis":
+                    replace = False
+                    logiw, pareto_k = az.psislw(logiw)
+                case "psir":
+                    replace = True
+                    logiw, pareto_k = az.psislw(logiw)
+                case "identity":
+                    replace = False
+                    pareto_k = None
 
     # NOTE: Pareto k is normally bad for Pathfinder even when the posterior is close to the NUTS posterior or closer to NUTS than ADVI.
     # Pareto k may not be a good diagnostic for Pathfinder.

{pymc_extras-0.2.2 → pymc_extras-0.2.4}/pymc_extras/inference/pathfinder/pathfinder.py

@@ -60,6 +60,7 @@ from pytensor.graph import Apply, Op, vectorize_graph
 from pytensor.tensor import TensorConstant, TensorVariable
 from rich.console import Console, Group
 from rich.padding import Padding
+from rich.progress import BarColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn
 from rich.table import Table
 from rich.text import Text
 
@@ -155,7 +156,7 @@ def convert_flat_trace_to_idata(
     postprocessing_backend: Literal["cpu", "gpu"] = "cpu",
     inference_backend: Literal["pymc", "blackjax"] = "pymc",
     model: Model | None = None,
-    importance_sampling: Literal["psis", "psir", "identity", "none"] = "psis",
+    importance_sampling: Literal["psis", "psir", "identity"] | None = "psis",
 ) -> az.InferenceData:
     """convert flattened samples to arviz InferenceData format.
 
@@ -180,7 +181,7 @@ def convert_flat_trace_to_idata(
         arviz inference data object
     """
 
-    if importance_sampling == "none":
+    if importance_sampling is None:
         # samples.ndim == 3 in this case, otherwise ndim == 2
         num_paths, num_pdraws, N = samples.shape
         samples = samples.reshape(-1, N)
@@ -219,7 +220,7 @@ def convert_flat_trace_to_idata(
         fn.trust_input = True
         result = fn(*list(trace.values()))
 
-        if importance_sampling == "none":
+        if importance_sampling is None:
             result = [res.reshape(num_paths, num_pdraws, *res.shape[2:]) for res in result]
 
     elif inference_backend == "blackjax":
@@ -1188,7 +1189,7 @@ class MultiPathfinderResult:
     elbo_argmax: NDArray | None = None
     lbfgs_status: Counter = field(default_factory=Counter)
     path_status: Counter = field(default_factory=Counter)
-    importance_sampling: str = "none"
+    importance_sampling: str | None = "psis"
     warnings: list[str] = field(default_factory=list)
     pareto_k: float | None = None
 
@@ -1257,7 +1258,7 @@ class MultiPathfinderResult:
     def with_importance_sampling(
         self,
         num_draws: int,
-        method: Literal["psis", "psir", "identity", "none"] | None,
+        method: Literal["psis", "psir", "identity"] | None,
         random_seed: int | None = None,
     ) -> Self:
         """perform importance sampling"""
@@ -1395,7 +1396,7 @@ def _get_status_warning(mpr: MultiPathfinderResult) -> list[str]:
 
     path_status_message = {
         PathStatus.ELBO_ARGMAX_AT_ZERO: "ELBO argmax at zero refers to the first iteration during LBFGS. A high occurrence suggests the model's default initial point + jitter is may be too close to the mean posterior and a poor exploration of the parameter space. Consider increasing jitter if this occurence is high relative to the number of paths.",
-        PathStatus.INVALID_LOGP: "Invalid logP values occur when a path's logP values are not finite. The failed path is not included in samples when importance sampling is used. Consider reparameterizing the model or adjusting the pathfinder arguments if this occurence is high relative to the number of paths.",
+        PathStatus.ELBO_ARGMAX_AT_ZERO: "ELBO argmax at zero refers to the first iteration during LBFGS. A high occurrence suggests the model's default initial point + jitter values are concentrated in high-density regions in the target distribution and may result in poor exploration of the parameter space. Consider increasing jitter if this occurrence is high relative to the number of paths.",
         PathStatus.INVALID_LOGQ: "Invalid logQ values occur when a path's logQ values are not finite. The failed path is not included in samples when importance sampling is used. Consider reparameterizing the model or adjusting the pathfinder arguments if this occurence is high relative to the number of paths.",
     }
 
@@ -1423,7 +1424,7 @@ def multipath_pathfinder(
     num_elbo_draws: int,
     jitter: float,
     epsilon: float,
-    importance_sampling: Literal["psis", "psir", "identity", "none"] | None,
+    importance_sampling: Literal["psis", "psir", "identity"] | None,
     progressbar: bool,
     concurrent: Literal["thread", "process"] | None,
     random_seed: RandomSeed,
@@ -1459,8 +1460,14 @@ def multipath_pathfinder(
         Amount of jitter to apply to initial points (default is 2.0). Note that Pathfinder may be highly sensitive to the jitter value. It is recommended to increase num_paths when increasing the jitter value.
     epsilon: float
         value used to filter out large changes in the direction of the update gradient at each iteration l in L. Iteration l is only accepted if delta_theta[l] * delta_grad[l] > epsilon * L2_norm(delta_grad[l]) for each l in L. (default is 1e-8).
-    importance_sampling : str, optional
-        importance sampling method to use which applies sampling based on the log importance weights equal to logP - logQ. Options are "psis" (default), "psir", "identity", "none". Pareto Smoothed Importance Sampling (psis) is recommended in many cases for more stable results than Pareto Smoothed Importance Resampling (psir). identity applies the log importance weights directly without resampling. none applies no importance sampling weights and returns the samples as is of size (num_paths, num_draws_per_path, N) where N is the number of model parameters, otherwise sample size is (num_draws, N).
+    importance_sampling : str, None, optional
+        Method to apply sampling based on log importance weights (logP - logQ).
+        "psis" : Pareto Smoothed Importance Sampling (default)
+                Recommended for more stable results.
+        "psir" : Pareto Smoothed Importance Resampling
+                Less stable than PSIS.
+        "identity" : Applies log importance weights directly without resampling.
+        None : No importance sampling weights. Returns raw samples of size (num_paths, num_draws_per_path, N) where N is number of model parameters. Other methods return samples of size (num_draws, N).
     progressbar : bool, optional
         Whether to display a progress bar (default is False). Setting this to True will likely increase the computation time.
     random_seed : RandomSeed, optional
@@ -1482,12 +1489,6 @@ def multipath_pathfinder(
         The result containing samples and other information from the Multi-Path Pathfinder algorithm.
     """
 
-    valid_importance_sampling = ["psis", "psir", "identity", "none", None]
-    if importance_sampling is None:
-        importance_sampling = "none"
-    if importance_sampling.lower() not in valid_importance_sampling:
-        raise ValueError(f"Invalid importance sampling method: {importance_sampling}")
-
     *path_seeds, choice_seed = _get_seeds_per_chain(random_seed, num_paths + 1)
 
     pathfinder_config = PathfinderConfig(
@@ -1521,12 +1522,20 @@ def multipath_pathfinder(
     results = []
     compute_start = time.time()
     try:
-        with CustomProgress(
+        desc = f"Paths Complete: {{path_idx}}/{num_paths}"
+        progress = CustomProgress(
+            "[progress.description]{task.description}",
+            BarColumn(),
+            "[progress.percentage]{task.percentage:>3.0f}%",
+            TimeRemainingColumn(),
+            TextColumn("/"),
+            TimeElapsedColumn(),
             console=Console(theme=default_progress_theme),
             disable=not progressbar,
-        ) as progress:
-            task = progress.add_task("Fitting", total=num_paths)
-            for result in generator:
+        )
+        with progress:
+            task = progress.add_task(desc.format(path_idx=0), completed=0, total=num_paths)
+            for path_idx, result in enumerate(generator, start=1):
                 try:
                     if isinstance(result, Exception):
                         raise result
@@ -1552,7 +1561,14 @@ def multipath_pathfinder(
                             lbfgs_status=LBFGSStatus.LBFGS_FAILED,
                         )
                     )
-                progress.update(task, advance=1)
+                finally:
+                    # TODO: display LBFGS and Path Status in real time
+                    progress.update(
+                        task,
+                        description=desc.format(path_idx=path_idx),
+                        completed=path_idx,
+                        refresh=True,
+                    )
     except (KeyboardInterrupt, StopIteration) as e:
         # if exception is raised here, MultiPathfinderResult will collect all the successful results and report the results. User is free to abort the process earlier and the results will still be collected and return az.InferenceData.
         if isinstance(e, StopIteration):
@@ -1606,7 +1622,7 @@ def fit_pathfinder(
     num_elbo_draws: int = 10,  # K
     jitter: float = 2.0,
     epsilon: float = 1e-8,
-    importance_sampling: Literal["psis", "psir", "identity", "none"] = "psis",
+    importance_sampling: Literal["psis", "psir", "identity"] | None = "psis",
     progressbar: bool = True,
     concurrent: Literal["thread", "process"] | None = None,
     random_seed: RandomSeed | None = None,
@@ -1646,8 +1662,15 @@ def fit_pathfinder(
         Amount of jitter to apply to initial points (default is 2.0). Note that Pathfinder may be highly sensitive to the jitter value. It is recommended to increase num_paths when increasing the jitter value.
     epsilon: float
         value used to filter out large changes in the direction of the update gradient at each iteration l in L. Iteration l is only accepted if delta_theta[l] * delta_grad[l] > epsilon * L2_norm(delta_grad[l]) for each l in L. (default is 1e-8).
-    importance_sampling : str, optional
-        importance sampling method to use which applies sampling based on the log importance weights equal to logP - logQ. Options are "psis" (default), "psir", "identity", "none". Pareto Smoothed Importance Sampling (psis) is recommended in many cases for more stable results than Pareto Smoothed Importance Resampling (psir). identity applies the log importance weights directly without resampling. none applies no importance sampling weights and returns the samples as is of size (num_paths, num_draws_per_path, N) where N is the number of model parameters, otherwise sample size is (num_draws, N).
+    importance_sampling : str, None, optional
+        Method to apply sampling based on log importance weights (logP - logQ).
+        Options are:
+        "psis" : Pareto Smoothed Importance Sampling (default)
+                Recommended for more stable results.
+        "psir" : Pareto Smoothed Importance Resampling
+                Less stable than PSIS.
+        "identity" : Applies log importance weights directly without resampling.
+        None : No importance sampling weights. Returns raw samples of size (num_paths, num_draws_per_path, N) where N is number of model parameters. Other methods return samples of size (num_draws, N).
     progressbar : bool, optional
         Whether to display a progress bar (default is True). Setting this to False will likely reduce the computation time.
     random_seed : RandomSeed, optional
@@ -1674,6 +1697,15 @@ def fit_pathfinder(
     """
 
     model = modelcontext(model)
+
+    valid_importance_sampling = {"psis", "psir", "identity", None}
+
+    if importance_sampling is not None:
+        importance_sampling = importance_sampling.lower()
+
+    if importance_sampling not in valid_importance_sampling:
+        raise ValueError(f"Invalid importance sampling method: {importance_sampling}")
+
     N = DictToArrayBijection.map(model.initial_point()).data.shape[0]
 
     if maxcor is None:

{pymc_extras-0.2.2 → pymc_extras-0.2.4}/pymc_extras/model/marginal/marginal_model.py

@@ -19,7 +19,8 @@ from pymc.model.fgraph import (
     model_free_rv,
     model_from_fgraph,
 )
-from pymc.pytensorf import collect_default_updates, compile_pymc, constant_fold, toposort_replace
+from pymc.pytensorf import collect_default_updates, constant_fold, toposort_replace
+from pymc.pytensorf import compile as compile_pymc
 from pymc.util import RandomState, _get_seeds_per_chain
 from pytensor import In, Out
 from pytensor.compile import SharedVariable

{pymc_extras-0.2.2 → pymc_extras-0.2.4}/pymc_extras/statespace/core/compile.py

@@ -30,7 +30,7 @@ def compile_statespace(
 
     inputs = list(pytensor.graph.basic.explicit_graph_inputs(outputs))
 
-    _f = pm.compile_pymc(inputs, outputs, on_unused_input="ignore", **compile_kwargs)
+    _f = pm.compile(inputs, outputs, on_unused_input="ignore", **compile_kwargs)
 
     def f(*, draws=1, **params):
         if isinstance(steps, pt.Variable):

{pymc_extras-0.2.2 → pymc_extras-0.2.4}/pymc_extras/statespace/core/statespace.py

@@ -28,7 +28,6 @@ from pymc_extras.statespace.filters import (
 )
 from pymc_extras.statespace.filters.distributions import (
     LinearGaussianStateSpace,
-    MvNormalSVD,
     SequenceMvNormal,
 )
 from pymc_extras.statespace.filters.utilities import stabilize
@@ -707,7 +706,7 @@ class PyMCStateSpace:
         with pymc_model:
             for param_name in self.param_names:
                 param = getattr(pymc_model, param_name, None)
-                if param:
+                if param is not None:
                     found_params.append(param.name)
 
         missing_params = list(set(self.param_names) - set(found_params))
@@ -746,7 +745,7 @@ class PyMCStateSpace:
         with pymc_model:
             for data_name in data_names:
                 data = getattr(pymc_model, data_name, None)
-                if data:
+                if data is not None:
                     found_data.append(data.name)
 
         missing_data = list(set(data_names) - set(found_data))
@@ -2233,7 +2232,9 @@ class PyMCStateSpace:
             if shock_trajectory is None:
                 shock_trajectory = pt.zeros((n_steps, self.k_posdef))
                 if Q is not None:
-                    init_shock = MvNormalSVD("initial_shock", mu=0, cov=Q, dims=[SHOCK_DIM])
+                    init_shock = pm.MvNormal(
+                        "initial_shock", mu=0, cov=Q, dims=[SHOCK_DIM], method="svd"
+                    )
                 else:
                     init_shock = pm.Deterministic(
                         "initial_shock",