pymc-extras 0.6.0__tar.gz → 0.8.0__tar.gz

{pymc_extras-0.6.0 → pymc_extras-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pymc-extras
-Version: 0.6.0
+Version: 0.8.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.1
+Requires-Dist: pytensor>=2.37.0
 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 → pymc_extras-0.8.0}/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.6.0'
-__version_tuple__ = version_tuple = (0, 6, 0)
+__version__ = version = '0.8.0'
+__version_tuple__ = version_tuple = (0, 8, 0)
 
 __commit_id__ = commit_id = None

{pymc_extras-0.6.0 → pymc_extras-0.8.0}/conda-envs/environment-test.yml

@@ -3,8 +3,8 @@ channels:
 - conda-forge
 - nodefaults
 dependencies:
-  - pymc>=5.26.1
-  - pytensor>=2.35.1
+  - pymc>=5.27.1
+  - pytensor>=2.37.0
   - scikit-learn
   - better-optimize>=0.1.5
   - dask<2025.1.1

{pymc_extras-0.6.0 → pymc_extras-0.8.0}/docs/statespace/models/structural.rst

@@ -6,9 +6,9 @@ Structural Components
 .. autosummary::
    :toctree: generated
 
-   LevelTrendComponent
-   AutoregressiveComponent
+   LevelTrend
+   Autoregressive
    TimeSeasonality
    FrequencySeasonality
    MeasurementError
-   CycleComponent
+   Cycle

{pymc_extras-0.6.0 → pymc_extras-0.8.0}/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-0.6.0 → pymc_extras-0.8.0}/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-0.6.0 → pymc_extras-0.8.0}/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,
@@ -137,7 +137,7 @@ def get_conditional_gaussian_approximation(
     hess = pytensor.graph.replace.graph_replace(hess, {x: x0})
 
     # Full log(p(x | y, params)) using the Laplace approximation (up to a constant)
-    _, logdetQ = pt.nlinalg.slogdet(Q)
+    _, logdetQ = pt.linalg.slogdet(Q)
     conditional_gaussian_approx = (
         -0.5 * x.T @ (-hess + Q) @ x + x.T @ (Q @ mu + jac - hess @ x0) + 0.5 * logdetQ
     )
@@ -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, keep_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