Merge upstream/main into main

Author: opherdonchin

.github/workflows/test.yml

@@ -32,7 +32,7 @@ jobs:
           - tests/statespace/filters/test_kalman_filter.py
           - tests/statespace --ignore tests/statespace/core/test_statespace.py --ignore tests/statespace/filters/test_kalman_filter.py
           - tests/distributions
-          - tests --ignore tests/model --ignore tests/statespace --ignore tests/distributions
+          - tests --ignore tests/model --ignore tests/statespace --ignore tests/distributions --ignore tests/pathfinder
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:

conda-envs/environment-test.yml

@@ -1,10 +1,10 @@
-name: pymc-extras-test
+name: pymc-extras
 channels:
 - conda-forge
 - nodefaults
 dependencies:
-  - pymc>=5.24.1
-  - pytensor>=2.31.4
+  - pymc>=5.26.1
+  - pytensor>=2.35.1
   - scikit-learn
   - better-optimize>=0.1.5
   - dask<2025.1.1

pymc_extras/inference/laplace_approx/find_map.py

@@ -168,6 +168,7 @@ def find_MAP(
     jitter_rvs: list[TensorVariable] | None = None,
     progressbar: bool = True,
     include_transformed: bool = True,
+    freeze_model: bool = True,
     gradient_backend: GradientBackend = "pytensor",
     compile_kwargs: dict | None = None,
     compute_hessian: bool = False,
@@ -210,6 +211,10 @@ def find_MAP(
         Whether to display a progress bar during optimization. Defaults to True.
     include_transformed: bool, optional
         Whether to include transformed variable values in the returned dictionary. Defaults to True.
+    freeze_model: bool, optional
+        If True, freeze_dims_and_data will be called on the model before compiling the loss functions. This is
+        sometimes necessary for JAX, and can sometimes improve performance by allowing constant folding. Defaults to
+        True.
     gradient_backend: str, default "pytensor"
         Which backend to use to compute gradients. Must be one of "pytensor" or "jax".
     compute_hessian: bool
@@ -229,11 +234,13 @@ def find_MAP(
         Results of Maximum A Posteriori (MAP) estimation, including the optimized point, inverse Hessian, transformed
         latent variables, and optimizer results.
     """
-    model = pm.modelcontext(model) if model is None else model
-    frozen_model = freeze_dims_and_data(model)
     compile_kwargs = {} if compile_kwargs is None else compile_kwargs
+    model = pm.modelcontext(model) if model is None else model
 
-    initial_params = _make_initial_point(frozen_model, initvals, random_seed, jitter_rvs)
+    if freeze_model:
+        model = freeze_dims_and_data(model)
+
+    initial_params = _make_initial_point(model, initvals, random_seed, jitter_rvs)
 
     do_basinhopping = method == "basinhopping"
     minimizer_kwargs = optimizer_kwargs.pop("minimizer_kwargs", {})
@@ -251,8 +258,8 @@ def find_MAP(
     )
 
     f_fused, f_hessp = scipy_optimize_funcs_from_loss(
-        loss=-frozen_model.logp(),
-        inputs=frozen_model.continuous_value_vars + frozen_model.discrete_value_vars,
+        loss=-model.logp(),
+        inputs=model.continuous_value_vars + model.discrete_value_vars,
         initial_point_dict=DictToArrayBijection.rmap(initial_params),
         use_grad=use_grad,
         use_hess=use_hess,
@@ -316,12 +323,10 @@ def find_MAP(
     }
 
     idata = map_results_to_inference_data(
-        map_point=optimized_point, model=frozen_model, include_transformed=include_transformed
+        map_point=optimized_point, model=model, include_transformed=include_transformed
     )
 
-    idata = add_fit_to_inference_data(
-        idata=idata, mu=raveled_optimized, H_inv=H_inv, model=frozen_model
-    )
+    idata = add_fit_to_inference_data(idata=idata, mu=raveled_optimized, H_inv=H_inv, model=model)
 
     idata = add_optimizer_result_to_inference_data(
         idata=idata, result=optimizer_result, method=method, mu=raveled_optimized, model=model

pymc_extras/inference/laplace_approx/laplace.py

@@ -168,9 +168,13 @@ def _unconstrained_vector_to_constrained_rvs(model):
     unconstrained_vector.name = "unconstrained_vector"
 
     # Redo the names list to ensure it is sorted to match the return order
-    names = [*constrained_names, *unconstrained_names]
+    constrained_rvs_and_names = [(rv, name) for rv, name in zip(constrained_rvs, constrained_names)]
+    value_rvs_and_names = [
+        (rv, name) for rv, name in zip(value_rvs, names) for name in unconstrained_names
+    ]
+    # names = [*constrained_names, *unconstrained_names]
 
-    return names, constrained_rvs, value_rvs, unconstrained_vector
+    return constrained_rvs_and_names, value_rvs_and_names, unconstrained_vector
 
 
 def model_to_laplace_approx(
@@ -182,8 +186,11 @@ def model_to_laplace_approx(
 
     # temp_chain and temp_draw are a hack to allow sampling from the Laplace approximation. We only have one mu and cov,
     # so we add batch dims (which correspond to chains and draws). But the names "chain" and "draw" are reserved.
-    names, constrained_rvs, value_rvs, unconstrained_vector = (
-        _unconstrained_vector_to_constrained_rvs(model)
+
+    # The model was frozen during the find_MAP procedure. To ensure we're operating on the same model, freeze it again.
+    frozen_model = freeze_dims_and_data(model)
+    constrained_rvs_and_names, _, unconstrained_vector = _unconstrained_vector_to_constrained_rvs(
+        frozen_model
     )
 
     coords = model.coords | {
@@ -204,12 +211,13 @@ def model_to_laplace_approx(
         )
 
         cast_to_var = partial(type_cast, Variable)
+        constrained_rvs, constrained_names = zip(*constrained_rvs_and_names)
         batched_rvs = vectorize_graph(
             type_cast(list[Variable], constrained_rvs),
             replace={cast_to_var(unconstrained_vector): cast_to_var(laplace_approximation)},
         )
 
-        for name, batched_rv in zip(names, batched_rvs):
+        for name, batched_rv in zip(constrained_names, batched_rvs):
             batch_dims = ("temp_chain", "temp_draw")
             if batched_rv.ndim == 2:
                 dims = batch_dims
@@ -285,6 +293,7 @@ def fit_laplace(
     jitter_rvs: list[pt.TensorVariable] | None = None,
     progressbar: bool = True,
     include_transformed: bool = True,
+    freeze_model: bool = True,
     gradient_backend: GradientBackend = "pytensor",
     chains: int = 2,
     draws: int = 500,
@@ -328,6 +337,10 @@ def fit_laplace(
     include_transformed: bool, default True
         Whether to include transformed variables in the output. If True, transformed variables will be included in the
         output InferenceData object. If False, only the original variables will be included.
+    freeze_model: bool, optional
+        If True, freeze_dims_and_data will be called on the model before compiling the loss functions. This is
+        sometimes necessary for JAX, and can sometimes improve performance by allowing constant folding. Defaults to
+        True.
     gradient_backend: str, default "pytensor"
         The backend to use for gradient computations. Must be one of "pytensor" or "jax".
     chains: int, default: 2
@@ -376,6 +389,9 @@ def fit_laplace(
     optimizer_kwargs = {} if optimizer_kwargs is None else optimizer_kwargs
     model = pm.modelcontext(model) if model is None else model
 
+    if freeze_model:
+        model = freeze_dims_and_data(model)
+
     idata = find_MAP(
         method=optimize_method,
         model=model,
@@ -387,6 +403,7 @@ def fit_laplace(
         jitter_rvs=jitter_rvs,
         progressbar=progressbar,
         include_transformed=include_transformed,
+        freeze_model=False,
         gradient_backend=gradient_backend,
         compile_kwargs=compile_kwargs,
         compute_hessian=True,

pymc_extras/inference/pathfinder/pathfinder.py

@@ -22,7 +22,7 @@
 from collections.abc import Callable, Iterator
 from dataclasses import asdict, dataclass, field, replace
 from enum import Enum, auto
-from typing import Literal, TypeAlias
+from typing import Literal, Self, TypeAlias
 
 import arviz as az
 import filelock
@@ -60,9 +60,6 @@
 from rich.table import Table
 from rich.text import Text
 
-# TODO: change to typing.Self after Python versions greater than 3.10
-from typing_extensions import Self
-
 from pymc_extras.inference.laplace_approx.idata import add_data_to_inference_data
 from pymc_extras.inference.pathfinder.importance_sampling import (
     importance_sampling as _importance_sampling,
@@ -533,7 +530,7 @@ def bfgs_sample_sparse(
 
     # qr_input: (L, N, 2J)
     qr_input = inv_sqrt_alpha_diag @ beta
-    (Q, R), _ = pytensor.scan(fn=pt.nlinalg.qr, sequences=[qr_input], allow_gc=False)
+    (Q, R), _ = pytensor.scan(fn=pt.linalg.qr, sequences=[qr_input], allow_gc=False)
 
     IdN = pt.eye(R.shape[1])[None, ...]
     IdN += IdN * REGULARISATION_TERM

pymc_extras/model/marginal/graph_analysis.py

@@ -6,8 +6,8 @@
 from pymc import SymbolicRandomVariable
 from pymc.model.fgraph import ModelVar
 from pymc.variational.minibatch_rv import MinibatchRandomVariable
-from pytensor.graph import Variable, ancestors
-from pytensor.graph.basic import io_toposort
+from pytensor.graph.basic import Variable
+from pytensor.graph.traversal import ancestors, io_toposort
 from pytensor.tensor import TensorType, TensorVariable
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise

pymc_extras/statespace/core/compile.py

@@ -28,7 +28,7 @@ def compile_statespace(
         x0, P0, c, d, T, Z, R, H, Q, steps=steps, sequence_names=sequence_names
     )
 
-    inputs = list(pytensor.graph.basic.explicit_graph_inputs(outputs))
+    inputs = list(pytensor.graph.traversal.explicit_graph_inputs(outputs))
 
     _f = pm.compile(inputs, outputs, on_unused_input="ignore", **compile_kwargs)
 

pymc_extras/statespace/filters/kalman_filter.py

@@ -200,7 +200,7 @@ def build_graph(
         self.n_endog = Z_shape[-2]
 
         data, a0, P0, *params = self.check_params(data, a0, P0, c, d, T, Z, R, H, Q)
-
+        data = pt.specify_shape(data, (data.type.shape[0], self.n_endog))
         sequences, non_sequences, seq_names, non_seq_names = split_vars_into_seq_and_nonseq(
             params, PARAM_NAMES
         )
@@ -658,7 +658,7 @@ def update(self, a, P, y, d, Z, H, all_nan_flag):
         # Construct upper-triangular block matrix A = [[chol(H), Z @ L_pred],
         #                                              [0,           L_pred]]
         # The Schur decomposition of this matrix will be B (upper triangular). We are
-        # more insterested in B^T:
+        # more interested in B^T:
         # Structure of B^T = [[chol(F),     0              ],
         #                    [K @ chol(F), chol(P_filtered)]
         zeros = pt.zeros((self.n_states, self.n_endog))

pymc_extras/statespace/filters/kalman_smoother.py

@@ -1,8 +1,6 @@
 import pytensor
 import pytensor.tensor as pt
 
-from pytensor.tensor.nlinalg import matrix_dot
-
 from pymc_extras.statespace.filters.utilities import (
     quad_form_sym,
     split_vars_into_seq_and_nonseq,
@@ -105,7 +103,7 @@ def smoother_step(self, *args):
         a_hat, P_hat = self.predict(a, P, T, R, Q)
 
         # Use pinv, otherwise P_hat is singular when there is missing data
-        smoother_gain = matrix_dot(pt.linalg.pinv(P_hat, hermitian=True), T, P).T
+        smoother_gain = (pt.linalg.pinv(P_hat, hermitian=True) @ T @ P).mT
         a_smooth_next = a + smoother_gain @ (a_smooth - a_hat)
 
         P_smooth_next = P + quad_form_sym(smoother_gain, P_smooth - P_hat)

pymc_extras/statespace/models/DFM.py

@@ -5,7 +5,7 @@
 import pytensor.tensor as pt
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace
-from pymc_extras.statespace.models.utilities import make_default_coords
+from pymc_extras.statespace.models.utilities import make_default_coords, validate_names
 from pymc_extras.statespace.utils.constants import (
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
@@ -224,9 +224,7 @@ def __init__(
         self,
         k_factors: int,
         factor_order: int,
-        k_endog: int | None = None,
         endog_names: Sequence[str] | None = None,
-        k_exog: int | None = None,
         exog_names: Sequence[str] | None = None,
         shared_exog_states: bool = False,
         exog_innovations: bool = False,
@@ -249,19 +247,11 @@ def __init__(
             and are modeled as a white noise process, i.e., :math:`f_t = \varepsilon_{f,t}`.
             Therefore, the state vector will include one state per factor and "factor_ar" will not exist.
 
-        k_endog : int, optional
-            Number of observed time series. If not provided, the number of observed series will be inferred from `endog_names`.
-            At least one of `k_endog` or `endog_names` must be provided.
-
         endog_names : list of str, optional
-            Names of the observed time series. If not provided, default names will be generated as `endog_1`, `endog_2`, ..., `endog_k` based on `k_endog`.
-            At least one of `k_endog` or `endog_names` must be provided.
-
-        k_exog : int, optional
-            Number of exogenous variables. If not provided, the model will not have exogenous variables.
+            Names of the observed time series.
 
         exog_names : Sequence[str], optional
-            Names of the exogenous variables. If not provided, but `k_exog` is specified, default names will be generated as `exog_1`, `exog_2`, ..., `exog_k`.
+            Names of the exogenous variables.
 
         shared_exog_states: bool, optional
             Whether exogenous latent states are shared across the observed states. If True, there will be only one set of exogenous latent
@@ -289,13 +279,8 @@ def __init__(
 
         """
 
-        if k_endog is None and endog_names is None:
-            raise ValueError("Either k_endog or endog_names must be provided.")
-        if k_endog is None:
-            k_endog = len(endog_names)
-        if endog_names is None:
-            endog_names = [f"endog_{i}" for i in range(k_endog)]
-
+        validate_names(endog_names, var_name="endog_names", optional=False)
+        k_endog = len(endog_names)
         self.endog_names = endog_names
         self.k_endog = k_endog
         self.k_factors = k_factors
@@ -304,17 +289,17 @@ def __init__(
         self.error_var = error_var
         self.error_cov_type = error_cov_type
 
-        if k_exog is None and exog_names is None:
-            self.k_exog = 0
-        else:
+        if exog_names is not None:
             self.shared_exog_states = shared_exog_states
             self.exog_innovations = exog_innovations
-            if k_exog is None:
-                k_exog = len(exog_names) if exog_names is not None else 0
-            elif exog_names is None:
-                exog_names = [f"exog_{i}" for i in range(k_exog)] if k_exog > 0 else None
+            validate_names(
+                exog_names, var_name="exog_names", optional=True
+            )  # Not sure if this adds anything
+            k_exog = len(exog_names)
             self.k_exog = k_exog
             self.exog_names = exog_names
+        else:
+            self.k_exog = 0
 
         self.k_exog_states = self.k_exog * self.k_endog if not shared_exog_states else self.k_exog
         self.exog_flag = self.k_exog > 0