pymc-extras 0.6.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

pymc_extras/distributions/timeseries.py

@@ -196,21 +196,20 @@ class DiscreteMarkovChain(Distribution):
         state_rng = pytensor.shared(np.random.default_rng())
 
         def transition(*args):
-            *states, transition_probs, old_rng = args
+            old_rng, *states, transition_probs = args
             p = transition_probs[tuple(states)]
             next_rng, next_state = pm.Categorical.dist(p=p, rng=old_rng).owner.outputs
-            return next_state, {old_rng: next_rng}
+            return next_rng, next_state
 
-        markov_chain, state_updates = pytensor.scan(
+        state_next_rng, markov_chain = pytensor.scan(
             transition,
-            non_sequences=[P_, state_rng],
-            outputs_info=_make_outputs_info(n_lags, init_dist_),
+            outputs_info=[state_rng, *_make_outputs_info(n_lags, init_dist_)],
+            non_sequences=[P_],
             n_steps=steps_,
             strict=True,
+            return_updates=False,
         )
 
-        (state_next_rng,) = tuple(state_updates.values())
-
         discrete_mc_ = pt.moveaxis(pt.concatenate([init_dist_, markov_chain], axis=0), 0, -1)
 
         discrete_mc_op = DiscreteMarkovChainRV(
@@ -239,16 +238,17 @@ def discrete_mc_moment(op, rv, P, steps, init_dist, state_rng):
     n_lags = op.n_lags
 
     def greedy_transition(*args):
-        *states, transition_probs, old_rng = args
+        *states, transition_probs = args
         p = transition_probs[tuple(states)]
         return pt.argmax(p)
 
-    chain_moment, moment_updates = pytensor.scan(
+    chain_moment = pytensor.scan(
         greedy_transition,
-        non_sequences=[P, state_rng],
+        non_sequences=[P],
         outputs_info=_make_outputs_info(n_lags, init_dist),
         n_steps=steps,
         strict=True,
+        return_updates=False,
     )
     chain_moment = pt.concatenate([init_dist_moment, chain_moment])
     return chain_moment

pymc_extras/inference/dadvi/dadvi.py

@@ -3,25 +3,20 @@ import numpy as np
 import pymc
 import pytensor
 import pytensor.tensor as pt
-import xarray
 
+from arviz import InferenceData
 from better_optimize import basinhopping, minimize
 from better_optimize.constants import minimize_method
 from pymc import DictToArrayBijection, Model, join_nonshared_inputs
-from pymc.backends.arviz import (
-    PointFunc,
-    apply_function_over_dataset,
-    coords_and_dims_for_inferencedata,
-)
 from pymc.blocking import RaveledVars
-from pymc.util import RandomSeed, get_default_varnames
+from pymc.util import RandomSeed
 from pytensor.tensor.variable import TensorVariable
 
 from pymc_extras.inference.laplace_approx.idata import (
     add_data_to_inference_data,
     add_optimizer_result_to_inference_data,
 )
-from pymc_extras.inference.laplace_approx.laplace import unstack_laplace_draws
+from pymc_extras.inference.laplace_approx.laplace import draws_from_laplace_approx
 from pymc_extras.inference.laplace_approx.scipy_interface import (
     scipy_optimize_funcs_from_loss,
     set_optimizer_function_defaults,
@@ -193,16 +188,18 @@ def fit_dadvi(
     opt_var_params = result.x
     opt_means, opt_log_sds = np.split(opt_var_params, 2)
 
-    # Make the draws:
-    generator = np.random.default_rng(seed=random_seed)
-    draws_raw = generator.standard_normal(size=(n_draws, n_params))
-
-    draws = opt_means + draws_raw * np.exp(opt_log_sds)
-    draws_arviz = unstack_laplace_draws(draws, model, chains=1, draws=n_draws)
-
-    idata = dadvi_result_to_idata(
-        draws_arviz, model, include_transformed=include_transformed, progressbar=progressbar
+    posterior, unconstrained_posterior = draws_from_laplace_approx(
+        mean=opt_means,
+        standard_deviation=np.exp(opt_log_sds),
+        draws=n_draws,
+        model=model,
+        vectorize_draws=False,
+        return_unconstrained=include_transformed,
+        random_seed=random_seed,
     )
+    idata = InferenceData(posterior=posterior)
+    if include_transformed:
+        idata.add_groups(unconstrained_posterior=unconstrained_posterior)
 
     var_name_to_model_var = {f"{var_name}_mu": var_name for var_name in initial_point_dict.keys()}
     var_name_to_model_var.update(
@@ -283,69 +280,3 @@ def create_dadvi_graph(
     objective = -mean_log_density - entropy
 
     return var_params, objective
-
-
-def dadvi_result_to_idata(
-    unstacked_draws: xarray.Dataset,
-    model: Model,
-    include_transformed: bool = False,
-    progressbar: bool = True,
-):
-    """
-    Transforms the unconstrained draws back into the constrained space.
-
-    Parameters
-    ----------
-    unstacked_draws : xarray.Dataset
-        The draws to constrain back into the original space.
-
-    model : Model
-        The PyMC model the variables were derived from.
-
-    n_draws: int
-        The number of draws to return from the variational approximation.
-
-    include_transformed: bool
-        Whether or not to keep the unconstrained variables in the output.
-
-    progressbar: bool
-        Whether or not to show a progress bar during the transformation. Default is True.
-
-    Returns
-    -------
-    :class:`~arviz.InferenceData`
-        Draws from the original constrained parameters.
-    """
-
-    filtered_var_names = model.unobserved_value_vars
-    vars_to_sample = list(
-        get_default_varnames(filtered_var_names, include_transformed=include_transformed)
-    )
-    fn = pytensor.function(model.value_vars, vars_to_sample)
-    point_func = PointFunc(fn)
-
-    coords, dims = coords_and_dims_for_inferencedata(model)
-
-    transformed_result = apply_function_over_dataset(
-        point_func,
-        unstacked_draws,
-        output_var_names=[x.name for x in vars_to_sample],
-        coords=coords,
-        dims=dims,
-        progressbar=progressbar,
-    )
-
-    constrained_names = [
-        x.name for x in get_default_varnames(model.unobserved_value_vars, include_transformed=False)
-    ]
-    all_varnames = [
-        x.name for x in get_default_varnames(model.unobserved_value_vars, include_transformed=True)
-    ]
-    unconstrained_names = sorted(set(all_varnames) - set(constrained_names))
-
-    idata = az.InferenceData(posterior=transformed_result[constrained_names])
-
-    if unconstrained_names and include_transformed:
-        idata["unconstrained_posterior"] = transformed_result[unconstrained_names]
-
-    return idata

pymc_extras/inference/laplace_approx/laplace.py

@@ -16,9 +16,7 @@
 import logging
 
 from collections.abc import Callable
-from functools import partial
 from typing import Literal
-from typing import cast as type_cast
 
 import arviz as az
 import numpy as np
@@ -27,16 +25,18 @@ import pytensor
 import pytensor.tensor as pt
 import xarray as xr
 
+from arviz import dict_to_dataset
 from better_optimize.constants import minimize_method
 from numpy.typing import ArrayLike
+from pymc import Model
+from pymc.backends.arviz import coords_and_dims_for_inferencedata
 from pymc.blocking import DictToArrayBijection
 from pymc.model.transform.optimization import freeze_dims_and_data
-from pymc.pytensorf import join_nonshared_inputs
-from pymc.util import get_default_varnames
+from pymc.util import get_untransformed_name, is_transformed_name
 from pytensor.graph import vectorize_graph
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.optimize import minimize
-from pytensor.tensor.type import Variable
+from xarray import Dataset
 
 from pymc_extras.inference.laplace_approx.find_map import (
     _compute_inverse_hessian,
@@ -147,138 +147,175 @@ def get_conditional_gaussian_approximation(
     return pytensor.function(args, [x0, conditional_gaussian_approx])
 
 
-def _unconstrained_vector_to_constrained_rvs(model):
-    outputs = get_default_varnames(model.unobserved_value_vars, include_transformed=True)
-    constrained_names = [
-        x.name for x in get_default_varnames(model.unobserved_value_vars, include_transformed=False)
-    ]
-    names = [x.name for x in outputs]
+def unpack_last_axis(packed_input, packed_shapes):
+    if len(packed_shapes) == 1:
+        # Single case currently fails in unpack
+        return [pt.split_dims(packed_input, packed_shapes[0], axis=-1)]
 
-    unconstrained_names = [name for name in names if name not in constrained_names]
+    keep_axes = tuple(range(packed_input.ndim))[:-1]
+    return pt.unpack(packed_input, axes=keep_axes, packed_shapes=packed_shapes)
 
-    new_outputs, unconstrained_vector = join_nonshared_inputs(
-        model.initial_point(),
-        inputs=model.value_vars,
-        outputs=outputs,
-    )
-
-    constrained_rvs = [x for x, name in zip(new_outputs, names) if name in constrained_names]
-    value_rvs = [x for x in new_outputs if x not in constrained_rvs]
-
-    unconstrained_vector.name = "unconstrained_vector"
 
-    # Redo the names list to ensure it is sorted to match the return order
-    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]
+def draws_from_laplace_approx(
+    *,
+    mean,
+    covariance=None,
+    standard_deviation=None,
+    draws: int,
+    model: Model,
+    vectorize_draws: bool = True,
+    return_unconstrained: bool = True,
+    random_seed=None,
+    compile_kwargs: dict | None = None,
+) -> tuple[Dataset, Dataset | None]:
+    """
+    Generate draws from the Laplace approximation of the posterior.
 
-    return constrained_rvs_and_names, value_rvs_and_names, unconstrained_vector
+    Parameters
+    ----------
+    mean : np.ndarray
+        The mean of the Laplace approximation (MAP estimate).
+    covariance : np.ndarray, optional
+        The covariance matrix of the Laplace approximation.
+        Mutually exclusive with `standard_deviation`.
+    standard_deviation : np.ndarray, optional
+        The standard deviation of the Laplace approximation (diagonal approximation).
+        Mutually exclusive with `covariance`.
+    draws : int
+        The number of draws.
+    model : pm.Model
+        The PyMC model.
+    vectorize_draws : bool, default True
+        Whether to vectorize the draws.
+    return_unconstrained : bool, default True
+        Whether to return the unconstrained draws in addition to the constrained ones.
+    random_seed : int, optional
+        Random seed for reproducibility.
+    compile_kwargs: dict, optional
+        Optional compile kwargs
 
+    Returns
+    -------
+    tuple[Dataset, Dataset | None]
+        A tuple containing the constrained draws (trace) and optionally the unconstrained draws.
+
+    Raises
+    ------
+    ValueError
+        If neither `covariance` nor `standard_deviation` is provided,
+        or if both are provided.
+    """
+    # This function assumes that mean/covariance/standard_deviation are aligned with model.initial_point()
+    if covariance is None and standard_deviation is None:
+        raise ValueError("Must specify either covariance or standard_deviation")
+    if covariance is not None and standard_deviation is not None:
+        raise ValueError("Cannot specify both covariance and standard_deviation")
+    if compile_kwargs is None:
+        compile_kwargs = {}
 
-def model_to_laplace_approx(
-    model: pm.Model, unpacked_variable_names: list[str], chains: int = 1, draws: int = 500
-):
     initial_point = model.initial_point()
-    raveled_vars = DictToArrayBijection.map(initial_point)
-    raveled_shape = raveled_vars.data.shape[0]
-
-    # 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.
-
-    # 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
+    n = int(np.sum([np.prod(v.shape) for v in initial_point.values()]))
+    assert mean.shape == (n,)
+    if covariance is not None:
+        assert covariance.shape == (n, n)
+    elif standard_deviation is not None:
+        assert standard_deviation.shape == (n,)
+
+    vars_to_sample = [v for v in model.free_RVs + model.deterministics]
+    var_names = [v.name for v in vars_to_sample]
+
+    orig_constrained_vars = model.value_vars
+    orig_outputs = model.replace_rvs_by_values(vars_to_sample)
+    if return_unconstrained:
+        orig_outputs.extend(model.value_vars)
+
+    mu_pt = pt.vector("mu", shape=(n,), dtype=mean.dtype)
+    size = (draws,) if vectorize_draws else ()
+    if covariance is not None:
+        sigma_pt = pt.matrix("cov", shape=(n, n), dtype=covariance.dtype)
+        laplace_approximation = pm.MvNormal.dist(mu=mu_pt, cov=sigma_pt, size=size, method="svd")
+    else:
+        sigma_pt = pt.vector("sigma", shape=(n,), dtype=standard_deviation.dtype)
+        laplace_approximation = pm.Normal.dist(mu=mu_pt, sigma=sigma_pt, size=(*size, n))
+
+    constrained_vars = unpack_last_axis(
+        laplace_approximation,
+        [initial_point[v.name].shape for v in orig_constrained_vars],
+    )
+    outputs = vectorize_graph(
+        orig_outputs, replace=dict(zip(orig_constrained_vars, constrained_vars))
     )
 
-    coords = model.coords | {
-        "temp_chain": np.arange(chains),
-        "temp_draw": np.arange(draws),
-        "unpacked_variable_names": unpacked_variable_names,
-    }
-
-    with pm.Model(coords=coords, model=None) as laplace_model:
-        mu = pm.Flat("mean_vector", shape=(raveled_shape,))
-        cov = pm.Flat("covariance_matrix", shape=(raveled_shape, raveled_shape))
-        laplace_approximation = pm.MvNormal(
-            "laplace_approximation",
-            mu=mu,
-            cov=cov,
-            dims=["temp_chain", "temp_draw", "unpacked_variable_names"],
-            method="svd",
-        )
-
-        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(constrained_names, batched_rvs):
-            batch_dims = ("temp_chain", "temp_draw")
-            if batched_rv.ndim == 2:
-                dims = batch_dims
-            elif name in model.named_vars_to_dims:
-                dims = (*batch_dims, *model.named_vars_to_dims[name])
-            else:
-                dims = (*batch_dims, *[f"{name}_dim_{i}" for i in range(batched_rv.ndim - 2)])
-                initval = initial_point.get(name, None)
-                dim_shapes = initval.shape if initval is not None else batched_rv.type.shape[2:]
-                laplace_model.add_coords(
-                    {name: np.arange(shape) for name, shape in zip(dims[2:], dim_shapes)}
-                )
-
-            pm.Deterministic(name, batched_rv, dims=dims)
-
-    return laplace_model
-
-
-def unstack_laplace_draws(laplace_data, model, chains=2, draws=500):
-    """
-    The `model_to_laplace_approx` function returns a model with a single MvNormal distribution, draws from which are
-    in the unconstrained variable space. These might be interesting to the user, but since they come back stacked in a
-    single vector, it's not easy to work with.
-
-    This function unpacks each component of the vector into its own DataArray, with the appropriate dimensions and
-    coordinates, where possible.
-    """
-    initial_point = DictToArrayBijection.map(model.initial_point())
-
-    cursor = 0
-    unstacked_laplace_draws = {}
-    coords = model.coords | {"chain": range(chains), "draw": range(draws)}
-
-    # There are corner cases where the value_vars will not have the same dimensions as the random variable (e.g.
-    # simplex transform of a Dirichlet). In these cases, we don't try to guess what the labels should be, and just
-    # add an arviz-style default dim and label.
-    for rv, (name, shape, size, dtype) in zip(model.free_RVs, initial_point.point_map_info):
-        rv_dims = []
-        for i, dim in enumerate(
-            model.named_vars_to_dims.get(rv.name, [f"{name}_dim_{i}" for i in range(len(shape))])
-        ):
-            if coords.get(dim) and shape[i] == len(coords[dim]):
-                rv_dims.append(dim)
-            else:
-                rv_dims.append(f"{name}_dim_{i}")
-                coords[f"{name}_dim_{i}"] = np.arange(shape[i])
-
-        dims = ("chain", "draw", *rv_dims)
-
-        values = (
-            laplace_data[..., cursor : cursor + size].reshape((chains, draws, *shape)).astype(dtype)
+    fn = pm.pytensorf.compile(
+        [mu_pt, sigma_pt],
+        outputs,
+        random_seed=random_seed,
+        trust_input=True,
+        **compile_kwargs,
+    )
+    sigma = covariance if covariance is not None else standard_deviation
+    if vectorize_draws:
+        output_buffers = fn(mean, sigma)
+    else:
+        # Take one draw to find the shape of the outputs
+        output_buffers = []
+        for out_draw in fn(mean, sigma):
+            output_buffer = np.empty((draws, *out_draw.shape), dtype=out_draw.dtype)
+            output_buffer[0] = out_draw
+            output_buffers.append(output_buffer)
+        # Fill one draws at a time
+        for i in range(1, draws):
+            for out_buffer, out_draw in zip(output_buffers, fn(mean, sigma)):
+                out_buffer[i] = out_draw
+
+    model_coords, model_dims = coords_and_dims_for_inferencedata(model)
+    posterior = {
+        var_name: out_buffer[None]
+        for var_name, out_buffer in (
+            zip(var_names, output_buffers, strict=not return_unconstrained)
         )
-        unstacked_laplace_draws[name] = xr.DataArray(
-            values, dims=dims, coords={dim: list(coords[dim]) for dim in dims}
+    }
+    posterior_dataset = dict_to_dataset(posterior, coords=model_coords, dims=model_dims, library=pm)
+    unconstrained_posterior_dataset = None
+
+    if return_unconstrained:
+        unconstrained_posterior = {
+            var.name: out_buffer[None]
+            for var, out_buffer in zip(
+                model.value_vars, output_buffers[len(posterior) :], strict=True
+            )
+        }
+        # Attempt to map constrained dims to unconstrained dims
+        for var_name, var_draws in unconstrained_posterior.items():
+            if not is_transformed_name(var_name):
+                # constrained == unconstrained, dims already shared
+                continue
+            constrained_dims = model_dims.get(get_untransformed_name(var_name))
+            if constrained_dims is None or (len(constrained_dims) != (var_draws.ndim - 2)):
+                continue
+            # Reuse dims from constrained variable if they match in length with unconstrained draws
+            inferred_dims = []
+            for i, (constrained_dim, unconstrained_dim_length) in enumerate(
+                zip(constrained_dims, var_draws.shape[2:], strict=True)
+            ):
+                if model_coords.get(constrained_dim) is not None and (
+                    len(model_coords[constrained_dim]) == unconstrained_dim_length
+                ):
+                    # Assume coordinates map. This could be fooled, by e.g., having a transform that reverses values
+                    inferred_dims.append(constrained_dim)
+                else:
+                    # Size mismatch (e.g., Simplex), make no assumption about mapping
+                    inferred_dims.append(f"{var_name}_dim_{i}")
+            model_dims[var_name] = inferred_dims
+
+        unconstrained_posterior_dataset = dict_to_dataset(
+            unconstrained_posterior,
+            coords=model_coords,
+            dims=model_dims,
+            library=pm,
         )
 
-        cursor += size
-
-    unstacked_laplace_draws = xr.Dataset(unstacked_laplace_draws)
-
-    return unstacked_laplace_draws
+    return posterior_dataset, unconstrained_posterior_dataset
 
 
 def fit_laplace(
@@ -295,8 +332,9 @@ def fit_laplace(
     include_transformed: bool = True,
     freeze_model: bool = True,
     gradient_backend: GradientBackend = "pytensor",
-    chains: int = 2,
+    chains: None | int = None,
     draws: int = 500,
+    vectorize_draws: bool = True,
     optimizer_kwargs: dict | None = None,
     compile_kwargs: dict | None = None,
 ) -> az.InferenceData:
@@ -343,16 +381,14 @@ def fit_laplace(
         True.
     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 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. Totals samples will be chains * draws.
+        The number of samples to draw from the approximated posterior.
     optimizer_kwargs
         Additional keyword arguments to pass to the ``scipy.optimize`` function being used. Unless
         ``method = "basinhopping"``, ``scipy.optimize.minimize`` will be used. For ``basinhopping``,
         ``scipy.optimize.basinhopping`` will be used. See the documentation of these functions for details.
+    vectorize_draws: bool, default True
+        Whether to natively vectorize the random function or take one at a time in a python loop.
     compile_kwargs: dict, optional
         Additional keyword arguments to pass to pytensor.function.
 
@@ -385,6 +421,12 @@ def fit_laplace(
           will forward the call to 'fit_laplace'.
 
     """
+    if chains is not None:
+        raise ValueError(
+            "chains argument has been deprecated. "
+            "The behavior can be recreated by unstacking draws into multiple chains after fitting"
+        )
+
     compile_kwargs = {} if compile_kwargs is None else compile_kwargs
     optimizer_kwargs = {} if optimizer_kwargs is None else optimizer_kwargs
     model = pm.modelcontext(model) if model is None else model
@@ -410,11 +452,10 @@ def fit_laplace(
         **optimizer_kwargs,
     )
 
-    unpacked_variable_names = idata.fit["mean_vector"].coords["rows"].values.tolist()
-
     if "covariance_matrix" not in idata.fit:
         # The user didn't use `use_hess` or `use_hessp` (or an optimization method that returns an inverse Hessian), so
         # we have to go back and compute the Hessian at the MAP point now.
+        unpacked_variable_names = idata.fit["mean_vector"].coords["rows"].values.tolist()
         frozen_model = freeze_dims_and_data(model)
         initial_params = _make_initial_point(frozen_model, initvals, random_seed, jitter_rvs)
 
@@ -443,29 +484,16 @@ def fit_laplace(
             coords={"rows": unpacked_variable_names, "columns": unpacked_variable_names},
         )
 
-    with model_to_laplace_approx(model, unpacked_variable_names, chains, draws) as laplace_model:
-        new_posterior = (
-            pm.sample_posterior_predictive(
-                idata.fit.expand_dims(chain=[0], draw=[0]),
-                extend_inferencedata=False,
-                random_seed=random_seed,
-                var_names=[
-                    "laplace_approximation",
-                    *[x.name for x in laplace_model.deterministics],
-                ],
-            )
-            .posterior_predictive.squeeze(["chain", "draw"])
-            .drop_vars(["chain", "draw"])
-            .rename({"temp_chain": "chain", "temp_draw": "draw"})
-        )
-
-        if include_transformed:
-            idata.unconstrained_posterior = unstack_laplace_draws(
-                new_posterior.laplace_approximation.values, model, chains=chains, draws=draws
-            )
-
-        idata.posterior = new_posterior.drop_vars(
-            ["laplace_approximation", "unpacked_variable_names"]
-        )
-
+    # We override the posterior/unconstrained_posterior from find_MAP
+    idata.posterior, unconstrained_posterior = draws_from_laplace_approx(
+        mean=idata.fit["mean_vector"].values,
+        covariance=idata.fit["covariance_matrix"].values,
+        draws=draws,
+        return_unconstrained=include_transformed,
+        model=model,
+        vectorize_draws=vectorize_draws,
+        random_seed=random_seed,
+    )
+    if include_transformed:
+        idata.unconstrained_posterior = unconstrained_posterior
     return idata

pymc_extras/inference/pathfinder/pathfinder.py

@@ -278,12 +278,13 @@ def alpha_recover(
     z = pt.diff(g, axis=0)
     alpha_l_init = pt.ones(N)
 
-    alpha, _ = pytensor.scan(
+    alpha = pytensor.scan(
         fn=compute_alpha_l,
         outputs_info=alpha_l_init,
         sequences=[s, z],
         n_steps=Lp1 - 1,
         allow_gc=False,
+        return_updates=False,
     )
 
     # assert np.all(alpha.eval() > 0), "alpha cannot be negative"
@@ -334,11 +335,12 @@ def inverse_hessian_factors(
             return pt.set_subtensor(chi_l[j_last], diff_l)
 
         chi_init = pt.zeros((J, N))
-        chi_mat, _ = pytensor.scan(
+        chi_mat = pytensor.scan(
             fn=chi_update,
             outputs_info=chi_init,
             sequences=[diff],
             allow_gc=False,
+            return_updates=False,
         )
 
         chi_mat = pt.matrix_transpose(chi_mat)
@@ -377,14 +379,14 @@ def inverse_hessian_factors(
     eta = pt.diagonal(E, axis1=-2, axis2=-1)
 
     # beta: (L, N, 2J)
-    alpha_diag, _ = pytensor.scan(lambda a: pt.diag(a), sequences=[alpha])
+    alpha_diag = pytensor.scan(lambda a: pt.diag(a), sequences=[alpha], return_updates=False)
     beta = pt.concatenate([alpha_diag @ Z, S], axis=-1)
 
     # more performant and numerically precise to use solve than inverse: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.inv.html
 
     # E_inv: (L, J, J)
     E_inv = pt.slinalg.solve_triangular(E, Ij, check_finite=False)
-    eta_diag, _ = pytensor.scan(pt.diag, sequences=[eta])
+    eta_diag = pytensor.scan(pt.diag, sequences=[eta], return_updates=False)
 
     # block_dd: (L, J, J)
     block_dd = (
@@ -530,7 +532,9 @@ def bfgs_sample_sparse(
 
     # qr_input: (L, N, 2J)
     qr_input = inv_sqrt_alpha_diag @ beta
-    (Q, R), _ = pytensor.scan(fn=pt.linalg.qr, sequences=[qr_input], allow_gc=False)
+    Q, R = pytensor.scan(
+        fn=pt.linalg.qr, sequences=[qr_input], allow_gc=False, return_updates=False
+    )
 
     IdN = pt.eye(R.shape[1])[None, ...]
     IdN += IdN * REGULARISATION_TERM
@@ -623,10 +627,11 @@ def bfgs_sample(
 
     L, N, JJ = beta.shape
 
-    (alpha_diag, inv_sqrt_alpha_diag, sqrt_alpha_diag), _ = pytensor.scan(
+    alpha_diag, inv_sqrt_alpha_diag, sqrt_alpha_diag = pytensor.scan(
         lambda a: [pt.diag(a), pt.diag(pt.sqrt(1.0 / a)), pt.diag(pt.sqrt(a))],
         sequences=[alpha],
         allow_gc=False,
+        return_updates=False,
     )
 
     u = pt.random.normal(size=(L, num_samples, N))

pymc_extras/inference/smc/sampling.py

@@ -238,7 +238,7 @@ class SMCDiagnostics(NamedTuple):
     def update_diagnosis(i, history, info, state):
         le, lli, ancestors, weights_evolution = history
         return SMCDiagnostics(
-            le.at[i].set(state.lmbda),
+            le.at[i].set(state.tempering_param),
             lli.at[i].set(info.log_likelihood_increment),
             ancestors.at[i].set(info.ancestors),
             weights_evolution.at[i].set(state.weights),
@@ -265,7 +265,7 @@ def inference_loop(rng_key, initial_state, kernel, iterations_to_diagnose, n_par
 
     def cond(carry):
         i, state, _, _ = carry
-        return state.lmbda < 1
+        return state.tempering_param < 1
 
     def one_step(carry):
         i, state, k, previous_info = carry

pymc_extras/model/marginal/distributions.py

@@ -282,11 +282,12 @@ def finite_discrete_marginal_rv_logp(op: MarginalFiniteDiscreteRV, values, *inpu
         def logp_fn(marginalized_rv_const, *non_sequences):
             return graph_replace(joint_logp, replace={marginalized_vv: marginalized_rv_const})
 
-        joint_logps, _ = scan_map(
+        joint_logps = scan_map(
             fn=logp_fn,
             sequences=marginalized_rv_domain_tensor,
             non_sequences=[*values, *inputs],
             mode=Mode().including("local_remove_check_parameter"),
+            return_updates=False,
         )
 
     joint_logp = pt.logsumexp(joint_logps, axis=0)
@@ -350,12 +351,13 @@ def marginal_hmm_logp(op, values, *inputs, **kwargs):
 
     P_bcast_dims = (len(chain_shape) - 1) - (P.type.ndim - 2)
     log_P = pt.shape_padright(pt.log(P), P_bcast_dims)
-    log_alpha_seq, _ = scan(
+    log_alpha_seq = scan(
         step_alpha,
         non_sequences=[log_P],
         outputs_info=[log_alpha_init],
         # Scan needs the time dimension first, and we already consumed the 1st logp computing the initial value
         sequences=pt.moveaxis(batch_logp_emissions[..., 1:], -1, 0),
+        return_updates=False,
     )
     # Final logp is just the sum of the last scan state
     joint_logp = pt.logsumexp(log_alpha_seq[-1], axis=0)

pymc_extras/model/marginal/marginal_model.py

@@ -11,7 +11,7 @@ from pymc.backends.arviz import coords_and_dims_for_inferencedata, dataset_to_po
 from pymc.distributions.discrete import Bernoulli, Categorical, DiscreteUniform
 from pymc.distributions.transforms import Chain
 from pymc.logprob.transforms import IntervalTransform
-from pymc.model import Model
+from pymc.model import Model, modelcontext
 from pymc.model.fgraph import (
     ModelFreeRV,
     ModelValuedVar,
@@ -337,8 +337,9 @@ def transform_posterior_pts(model, posterior_pts):
 
 
 def recover_marginals(
-    model: Model,
     idata: InferenceData,
+    *,
+    model: Model | None = None,
     var_names: Sequence[str] | None = None,
     return_samples: bool = True,
     extend_inferencedata: bool = True,
@@ -389,6 +390,15 @@ def recover_marginals(
 
 
     """
+    # Temporary error message for helping with migration
+    # Will be removed in a future release
+    if isinstance(idata, Model):
+        raise TypeError(
+            "The order of arguments of `recover_marginals` changed. The first input must be an idata"
+        )
+
+    model = modelcontext(model)
+
     unmarginal_model = unmarginalize(model)
 
     # Find the names of the marginalized variables

pymc_extras/statespace/core/statespace.py

@@ -2500,13 +2500,14 @@ class PyMCStateSpace:
                 next_x = c + T @ x + R @ shock
                 return next_x
 
-            irf, updates = pytensor.scan(
+            irf = pytensor.scan(
                 irf_step,
                 sequences=[shock_trajectory],
                 outputs_info=[x0],
                 non_sequences=[c, T, R],
                 n_steps=n_steps,
                 strict=True,
+                return_updates=False,
             )
 
             pm.Deterministic("irf", irf, dims=[TIME_DIM, ALL_STATE_DIM])

pymc_extras/statespace/filters/distributions.py

@@ -197,10 +197,9 @@ class _LinearGaussianStateSpace(Continuous):
         n_seq = len(sequence_names)
 
         def step_fn(*args):
-            seqs, state, non_seqs = args[:n_seq], args[n_seq], args[n_seq + 1 :]
-            non_seqs, rng = non_seqs[:-1], non_seqs[-1]
+            seqs, (rng, state, *non_seqs) = args[:n_seq], args[n_seq:]
 
-            c, d, T, Z, R, H, Q = sort_args(seqs + non_seqs)
+            c, d, T, Z, R, H, Q = sort_args((*seqs, *non_seqs))
             k = T.shape[0]
             a = state[:k]
 
@@ -219,7 +218,7 @@ class _LinearGaussianStateSpace(Continuous):
 
             next_state = pt.concatenate([a_next, y_next], axis=0)
 
-            return next_state, {rng: next_rng}
+            return next_rng, next_state
 
         Z_init = Z_ if Z_ in non_sequences else Z_[0]
         H_init = H_ if H_ in non_sequences else H_[0]
@@ -229,13 +228,14 @@ class _LinearGaussianStateSpace(Continuous):
 
         init_dist_ = pt.concatenate([init_x_, init_y_], axis=0)
 
-        statespace, updates = pytensor.scan(
+        ss_rng, statespace = pytensor.scan(
             step_fn,
-            outputs_info=[init_dist_],
+            outputs_info=[rng, init_dist_],
             sequences=None if len(sequences) == 0 else sequences,
-            non_sequences=[*non_sequences, rng],
+            non_sequences=[*non_sequences],
             n_steps=steps,
             strict=True,
+            return_updates=False,
         )
 
         if append_x0:
@@ -245,7 +245,6 @@ class _LinearGaussianStateSpace(Continuous):
             statespace_ = statespace
             statespace_ = pt.specify_shape(statespace_, (steps, None))
 
-        (ss_rng,) = tuple(updates.values())
         linear_gaussian_ss_op = LinearGaussianStateSpaceRV(
             inputs=[a0_, P0_, c_, d_, T_, Z_, R_, H_, Q_, steps, rng],
             outputs=[ss_rng, statespace_],
@@ -385,10 +384,15 @@ class SequenceMvNormal(Continuous):
 
         def step(mu, cov, rng):
             new_rng, mvn = pm.MvNormal.dist(mu=mu, cov=cov, rng=rng, method=method).owner.outputs
-            return mvn, {rng: new_rng}
+            return new_rng, mvn
 
-        mvn_seq, updates = pytensor.scan(
-            step, sequences=[mus_, covs_], non_sequences=[rng], strict=True, n_steps=mus_.shape[0]
+        seq_mvn_rng, mvn_seq = pytensor.scan(
+            step,
+            sequences=[mus_, covs_],
+            outputs_info=[rng, None],
+            strict=True,
+            n_steps=mus_.shape[0],
+            return_updates=False,
         )
         mvn_seq = pt.specify_shape(mvn_seq, mus.type.shape)
 
@@ -396,8 +400,6 @@ class SequenceMvNormal(Continuous):
         if mvn_seq.ndim > 2:
             mvn_seq = pt.moveaxis(mvn_seq, 0, -2)
 
-        (seq_mvn_rng,) = tuple(updates.values())
-
         mvn_seq_op = KalmanFilterRV(
             inputs=[mus_, covs_, logp_, rng], outputs=[seq_mvn_rng, mvn_seq], ndim_supp=2
         )

pymc_extras/statespace/filters/kalman_filter.py

@@ -148,10 +148,9 @@ class BaseFilter(ABC):
         R,
         H,
         Q,
-        return_updates=False,
         missing_fill_value=None,
         cov_jitter=None,
-    ) -> list[TensorVariable] | tuple[list[TensorVariable], dict]:
+    ) -> list[TensorVariable]:
         """
         Construct the computation graph for the Kalman filter. See [1] for details.
 
@@ -211,20 +210,17 @@ class BaseFilter(ABC):
         if len(sequences) > 0:
             sequences = self.add_check_on_time_varying_shapes(data, sequences)
 
-        results, updates = pytensor.scan(
+        results = pytensor.scan(
             self.kalman_step,
             sequences=[data, *sequences],
             outputs_info=[None, a0, None, None, P0, None, None],
             non_sequences=non_sequences,
             name="forward_kalman_pass",
             strict=False,
+            return_updates=False,
         )
 
-        filter_results = self._postprocess_scan_results(results, a0, P0, n=data.type.shape[0])
-
-        if return_updates:
-            return filter_results, updates
-        return filter_results
+        return self._postprocess_scan_results(results, a0, P0, n=data.type.shape[0])
 
     def _postprocess_scan_results(self, results, a0, P0, n) -> list[TensorVariable]:
         """
@@ -652,7 +648,9 @@ class SquareRootFilter(BaseFilter):
         y_hat = Z.dot(a) + d
         v = y - y_hat
 
-        H_chol = pytensor.ifelse(pt.all(pt.eq(H, 0.0)), H, pt.linalg.cholesky(H, lower=True))
+        H_chol = pytensor.ifelse(
+            pt.all(pt.eq(H, 0.0)), H, pt.linalg.cholesky(H, lower=True, on_error="nan")
+        )
 
         # The following notation comes from https://ipnpr.jpl.nasa.gov/progress_report/42-233/42-233A.pdf
         # Construct upper-triangular block matrix A = [[chol(H), Z @ L_pred],
@@ -694,8 +692,10 @@ class SquareRootFilter(BaseFilter):
             """
             return [a, P_chol, pt.zeros(())]
 
+        degenerate = pt.eq(all_nan_flag, 1.0)
+        F_chol = pytensor.ifelse(degenerate, pt.eye(*F_chol.shape), F_chol)
         [a_filtered, P_chol_filtered, ll] = pytensor.ifelse(
-            pt.eq(all_nan_flag, 1.0),
+            degenerate,
             compute_degenerate(P_chol_filtered, F_chol, K_F_chol, v),
             compute_non_degenerate(P_chol_filtered, F_chol, K_F_chol, v),
         )
@@ -786,11 +786,12 @@ class UnivariateFilter(BaseFilter):
         H_masked = W.dot(H)
         y_masked = pt.set_subtensor(y[nan_mask], 0.0)
 
-        result, updates = pytensor.scan(
+        result = pytensor.scan(
             self._univariate_inner_filter_step,
             sequences=[y_masked, Z_masked, d, pt.diag(H_masked), nan_mask],
             outputs_info=[a, P, None, None, None],
             name="univariate_inner_scan",
+            return_updates=False,
         )
 
         a_filtered, P_filtered, obs_mu, obs_cov, ll_inner = result

pymc_extras/statespace/filters/kalman_smoother.py

@@ -76,16 +76,16 @@ class KalmanSmoother:
         self.seq_names = seq_names
         self.non_seq_names = non_seq_names
 
-        smoother_result, updates = pytensor.scan(
+        smoothed_states, smoothed_covariances = pytensor.scan(
             self.smoother_step,
             sequences=[filtered_states[:-1], filtered_covariances[:-1], *sequences],
             outputs_info=[a_last, P_last],
             non_sequences=non_sequences,
             go_backwards=True,
             name="kalman_smoother",
+            return_updates=False,
         )
 
-        smoothed_states, smoothed_covariances = smoother_result
         smoothed_states = pt.concatenate(
             [smoothed_states[::-1], pt.expand_dims(a_last, axis=(0,))], axis=0
         )

{pymc_extras-0.6.0.dist-info → pymc_extras-0.7.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pymc-extras
-Version: 0.6.0
+Version: 0.7.0
 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.
 Project-URL: Documentation, https://pymc-extras.readthedocs.io/
 Project-URL: Repository, https://github.com/pymc-devs/pymc-extras.git
@@ -235,8 +235,8 @@ Requires-Python: >=3.11
 Requires-Dist: better-optimize>=0.1.5
 Requires-Dist: preliz>=0.20.0
 Requires-Dist: pydantic>=2.0.0
-Requires-Dist: pymc>=5.26.1
-Requires-Dist: pytensor>=2.35.1
+Requires-Dist: pymc>=5.27.0
+Requires-Dist: pytensor>=2.36.3
 Requires-Dist: scikit-learn
 Provides-Extra: complete
 Requires-Dist: dask[complete]<2025.1.1; extra == 'complete'
@@ -245,7 +245,7 @@ Provides-Extra: dask-histogram
 Requires-Dist: dask[complete]<2025.1.1; extra == 'dask-histogram'
 Requires-Dist: xhistogram; extra == 'dask-histogram'
 Provides-Extra: dev
-Requires-Dist: blackjax; extra == 'dev'
+Requires-Dist: blackjax>=0.12; extra == 'dev'
 Requires-Dist: dask[all]<2025.1.1; extra == 'dev'
 Requires-Dist: pytest-mock; extra == 'dev'
 Requires-Dist: pytest>=6.0; extra == 'dev'

{pymc_extras-0.6.0.dist-info → pymc_extras-0.7.0.dist-info}/RECORD

@@ -8,7 +8,7 @@ pymc_extras/distributions/__init__.py,sha256=Cge3AP7gzD6qTJY7v2tYRtSgn-rlnIo7wQB
 pymc_extras/distributions/continuous.py,sha256=bCXOgnw2Vh_FbYOHCqB0c3ozFVay5Qwua2A211kvWNQ,11251
 pymc_extras/distributions/discrete.py,sha256=HNi-K0_hnNWTcfyBkWGh26sc71FwBgukQ_EjGAaAOjY,13036
 pymc_extras/distributions/histogram_utils.py,sha256=kkZHu1F_2qMfOEzwNP4K6QYA_xEKUk9cChImOQ2Nkjs,5847
-pymc_extras/distributions/timeseries.py,sha256=R9htJ24LCFe5zyR9v6nX7S8iFmSMoQKSIGnLWZ1849Y,12747
+pymc_extras/distributions/timeseries.py,sha256=WysWtUchfObTGmKduF47bUBqV_g1kW-uAx4_oKENgDg,12709
 pymc_extras/distributions/multivariate/__init__.py,sha256=E8OeLW9tTotCbrUjEo4um76-_WQD56PehsPzkKmhfyA,93
 pymc_extras/distributions/multivariate/r2d2m2cp.py,sha256=5SzvD41pu-EWyWlDNz4AR4Sl8MkyC-1dYwkADFh5Avg,16009
 pymc_extras/distributions/transforms/__init__.py,sha256=FUp2vyRE6_2eUcQ_FVt5Dn0-vy5I-puV-Kz13-QtLNc,104
@@ -18,25 +18,25 @@ pymc_extras/gp/latent_approx.py,sha256=cDEMM6H1BL2qyKg7BZU-ISrKn2HJe7hDaM4Y8GgQD
 pymc_extras/inference/__init__.py,sha256=hI3yqfEVzoUNlCpL1z579F9EqM-NlPTzMfHj8IKY-xE,1009
 pymc_extras/inference/fit.py,sha256=hNTqLms_mTdjfnCEVIHMcMiPZ3fkU3HEEkbt6LWWhLw,1443
 pymc_extras/inference/dadvi/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pymc_extras/inference/dadvi/dadvi.py,sha256=Ry1Of_zo0eURVKbWE0U7NDmb7HQ2I_PWBFQcyowAXbE,11799
+pymc_extras/inference/dadvi/dadvi.py,sha256=rERiyMn1ywEerWJ8rq3WNZKtKEpX2lHAdqApatZyJpQ,9698
 pymc_extras/inference/laplace_approx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/inference/laplace_approx/find_map.py,sha256=fbK0swDsSBo7pP1TBokREa2wkK1ajL_gLVVuREHH33k,13658
 pymc_extras/inference/laplace_approx/idata.py,sha256=Dxj6A8aJXn8c24vD_PZmMgIgrwEmaYDlbw5UAJq0Nyw,14172
-pymc_extras/inference/laplace_approx/laplace.py,sha256=suSxCibVry1p3VjdvG9I_9QxazKdcKKT-xmCN1qXNRA,19706
+pymc_extras/inference/laplace_approx/laplace.py,sha256=J4Ddt7Jc1nRZvxHYUz2CWSpvJCJOMG3p2ayyUf1T7tE,20377
 pymc_extras/inference/laplace_approx/scipy_interface.py,sha256=Crhix_dLA8Y_NvuUDmVQnKWAWGjufmQwDLh-bK9dz_o,10235
 pymc_extras/inference/pathfinder/__init__.py,sha256=FhAYrCWNx_dCrynEdjg2CZ9tIinvcVLBm67pNx_Y3kA,101
 pymc_extras/inference/pathfinder/idata.py,sha256=muAPc9JeI8ZmpjzSp9tSj-uNrcsoNkYb4raJqjgf5UQ,18636
 pymc_extras/inference/pathfinder/importance_sampling.py,sha256=NwxepXOFit3cA5zEebniKdlnJ1rZWg56aMlH4MEOcG4,6264
 pymc_extras/inference/pathfinder/lbfgs.py,sha256=GOoJBil5Kft_iFwGNUGKSeqzI5x_shA4KQWDwgGuQtQ,7110
-pymc_extras/inference/pathfinder/pathfinder.py,sha256=3fYq3cgchwXxTEehT2PLu32r7hP7vAMoPpOS7TYpZ-w,67193
+pymc_extras/inference/pathfinder/pathfinder.py,sha256=IdKyJvGAeRstvTprKVQ4xk1hy6KjB8h-ggbmM7kMPEw,67345
 pymc_extras/inference/smc/__init__.py,sha256=wyaT4NJl1YsSQRLiDy-i0Jq3CbJZ2BQd4nnCk-dIngY,603
-pymc_extras/inference/smc/sampling.py,sha256=AYwmKqGoV6pBtKnh9SUbBKbN7VcoFgb3MmNWV7SivMA,15365
+pymc_extras/inference/smc/sampling.py,sha256=eyRIFPf--tcPpuHPNCxGZNQZVd7MazR4l9aURNY87S0,15385
 pymc_extras/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/model/model_api.py,sha256=UHMfQXxWBujeSiUySU0fDUC5Sd_BjT8FoVz3iBxQH_4,2400
 pymc_extras/model/marginal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pymc_extras/model/marginal/distributions.py,sha256=iM1yT7_BmivgUSloQPKE2QXGPgjvLqDMY_OTBGsdAWg,15563
+pymc_extras/model/marginal/distributions.py,sha256=mf6Czm6av2nCydu6uKqjamKFPD8efWJmNlTMy4Ojrvk,15621
 pymc_extras/model/marginal/graph_analysis.py,sha256=Ft7RZC126R0TW2GuFdgb9uN-JSgDGTeffs-UuPcDHQE,15884
-pymc_extras/model/marginal/marginal_model.py,sha256=oIdikaSnefCkyMxmzAe222qGXNucxZpHYk7548fK6iA,23631
+pymc_extras/model/marginal/marginal_model.py,sha256=Wgfcq6hplACU4Kh8aKT2P_kz_yo7r0wIVTupZQtQUKw,23969
 pymc_extras/model/transforms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/model/transforms/autoreparam.py,sha256=_NltGWmNqi_X9sHCqAvWcBveLTPxVy11-wENFTcN6kk,12377
 pymc_extras/preprocessing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -45,11 +45,11 @@ pymc_extras/statespace/__init__.py,sha256=PxV8i4aa2XJarRM6aKU14_bEY1AoLu4bNXIBy_
 pymc_extras/statespace/core/__init__.py,sha256=LEhkqdMZzzcTyzYml45IM4ykWoCdbWWj2c29IpM_ey8,309
 pymc_extras/statespace/core/compile.py,sha256=GB2H7sE28OdQ6GmNIjtq1R1Oua2GPf6kWJ7IPuYJaNA,1607
 pymc_extras/statespace/core/representation.py,sha256=boY-jjlkd3KuuO2XiSuV-GwEAyEqRJ9267H72AmE3BU,18956
-pymc_extras/statespace/core/statespace.py,sha256=yu7smA5w7l1LFNjTwuKLnGarGLx4HEPJKQ9ZMDbWhDY,108161
+pymc_extras/statespace/core/statespace.py,sha256=pyXIS95MJtJFzQozzdx4tNvf5-jY0Po3Z8aqaNAD7uo,108190
 pymc_extras/statespace/filters/__init__.py,sha256=F0EtZUhArp23lj3upy6zB0mDTjLIjwGh0pKmMny0QfY,420
-pymc_extras/statespace/filters/distributions.py,sha256=-s1c5s2zm6FMc0UqKSrWnJzIF4U5bvJT_3mMNTyV_ak,11927
-pymc_extras/statespace/filters/kalman_filter.py,sha256=0Ei_ZSogSAAC9OQn4v0tvUqvRZFHh6V3Wx8013-POBQ,31542
-pymc_extras/statespace/filters/kalman_smoother.py,sha256=Q1mwKeZPEMQL3BK2IvGwLfblGGQrfJ2t60fzll1Pung,4128
+pymc_extras/statespace/filters/distributions.py,sha256=uLCs3iJObHyslOPLUFJp5G9w56AWJIofiFq2KozecXc,11881
+pymc_extras/statespace/filters/kalman_filter.py,sha256=x6J54t9cHi3tXKtCB6QW62Kyit5zjd5AnI3IFdxKtzw,31561
+pymc_extras/statespace/filters/kalman_smoother.py,sha256=JmnvXwHVzWTdmtgECTJY0FJOFZG6O9aEfRoSTWEeU2s,4111
 pymc_extras/statespace/filters/utilities.py,sha256=BBMDeWBcJWZfGc9owuMsOedVIXVDQ8Z2eMiU9vWeVr0,1494
 pymc_extras/statespace/models/DFM.py,sha256=EiZ3x4iFPGeha8bPp1tok4us8Z6UVUu1sFmKIM1i0xc,36458
 pymc_extras/statespace/models/ETS.py,sha256=LEsSKzbfm9Ol8UZQjNurcrM1CLQyozKfJtby7AzsDeI,27667
@@ -76,7 +76,7 @@ pymc_extras/utils/linear_cg.py,sha256=KkXhuimFsrKtNd_0By2ApxQQQNm5FdBtmDQJOVbLYk
 pymc_extras/utils/model_equivalence.py,sha256=9MLwSj7VwxxKupzmEkKBbwGD1X0WM2FGcGIpfb8bViw,2197
 pymc_extras/utils/prior.py,sha256=mnuFpamp04eQJuTU5NyB2PfCG5r-1McSmQGwQXSR_Lg,6670
 pymc_extras/utils/spline.py,sha256=R0u3eAcV5bRmD2YSLqDm0qnaJbEuf3V38OZ7amV7-Tc,4732
-pymc_extras-0.6.0.dist-info/METADATA,sha256=Dn4Xwbebxw8QiFA0veLrEktkowKsAx06h00p7Epec1E,18898
-pymc_extras-0.6.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-pymc_extras-0.6.0.dist-info/licenses/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
-pymc_extras-0.6.0.dist-info/RECORD,,
+pymc_extras-0.7.0.dist-info/METADATA,sha256=iTIf9JVSRSbvmGz-ASblyd_lR8Jj4eWrUsjgzx97QUw,18904
+pymc_extras-0.7.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+pymc_extras-0.7.0.dist-info/licenses/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
+pymc_extras-0.7.0.dist-info/RECORD,,