pymc-extras 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

tests/statespace/utilities/shared_fixtures.py

@@ -1,9 +0,0 @@
-import numpy as np
-import pytest
-
-TEST_SEED = sum(map(ord, "statespace"))
-
-
-@pytest.fixture(scope="session")
-def rng():
-    return np.random.default_rng(TEST_SEED)

tests/statespace/utilities/statsmodel_local_level.py

@@ -1,42 +0,0 @@
-import numpy as np
-import statsmodels.api as sm
-
-
-class LocalLinearTrend(sm.tsa.statespace.MLEModel):
-    def __init__(self, endog, **kwargs):
-        # Model order
-        k_states = k_posdef = 2
-
-        # Initialize the statespace
-        super().__init__(endog, k_states=k_states, k_posdef=k_posdef, **kwargs)
-
-        # Initialize the matrices
-        self.ssm["design"] = np.array([1, 0])
-        self.ssm["transition"] = np.array([[1, 1], [0, 1]])
-        self.ssm["selection"] = np.eye(k_states)
-
-        # Cache some indices
-        self._state_cov_idx = ("state_cov", *np.diag_indices(k_posdef))
-
-    @property
-    def param_names(self):
-        return ["sigma2.measurement", "sigma2.level", "sigma2.trend"]
-
-    @property
-    def start_params(self):
-        return [np.std(self.endog)] * 3
-
-    def transform_params(self, unconstrained):
-        return unconstrained**2
-
-    def untransform_params(self, constrained):
-        return constrained**0.5
-
-    def update(self, params, *args, **kwargs):
-        params = super().update(params, *args, **kwargs)
-
-        # Observation covariance
-        self.ssm["obs_cov", 0, 0] = params[0]
-
-        # State covariance
-        self.ssm[self._state_cov_idx] = params[1:]

tests/statespace/utilities/test_helpers.py

@@ -1,310 +0,0 @@
-import numpy as np
-import pandas as pd
-import pytensor
-import pytensor.tensor as pt
-import statsmodels.api as sm
-
-from numpy.testing import assert_allclose
-from pymc import modelcontext
-
-from pymc_extras.statespace.filters.kalman_smoother import KalmanSmoother
-from pymc_extras.statespace.utils.constants import (
-    MATRIX_NAMES,
-    SHORT_NAME_TO_LONG,
-)
-from tests.statespace.utilities.statsmodel_local_level import LocalLinearTrend
-
-floatX = pytensor.config.floatX
-
-
-def load_nile_test_data():
-    from importlib.metadata import version
-
-    nile = pd.read_csv("tests/statespace/test_data/nile.csv", dtype={"x": floatX})
-    major, minor, rev = map(int, version("pandas").split("."))
-    if major >= 2 and minor >= 2 and rev >= 0:
-        freq_str = "YS-JAN"
-    else:
-        freq_str = "AS-JAN"
-    nile.index = pd.date_range(start="1871-01-01", end="1970-01-01", freq=freq_str)
-    nile.rename(columns={"x": "height"}, inplace=True)
-    nile = (nile - nile.mean()) / nile.std()
-    nile = nile.astype(floatX)
-
-    return nile
-
-
-def initialize_filter(kfilter, mode=None, p=None, m=None, r=None, n=None):
-    ksmoother = KalmanSmoother()
-    data = pt.tensor(name="data", dtype=floatX, shape=(n, p))
-    a0 = pt.tensor(name="x0", dtype=floatX, shape=(m,))
-    P0 = pt.tensor(name="P0", dtype=floatX, shape=(m, m))
-    c = pt.tensor(name="c", dtype=floatX, shape=(m,))
-    d = pt.tensor(name="d", dtype=floatX, shape=(p,))
-    Q = pt.tensor(name="Q", dtype=floatX, shape=(r, r))
-    H = pt.tensor(name="H", dtype=floatX, shape=(p, p))
-    T = pt.tensor(name="T", dtype=floatX, shape=(m, m))
-    R = pt.tensor(name="R", dtype=floatX, shape=(m, r))
-    Z = pt.tensor(name="Z", dtype=floatX, shape=(p, m))
-
-    inputs = [data, a0, P0, c, d, T, Z, R, H, Q]
-
-    (
-        filtered_states,
-        predicted_states,
-        observed_states,
-        filtered_covs,
-        predicted_covs,
-        observed_covs,
-        ll_obs,
-    ) = kfilter.build_graph(*inputs, mode=mode)
-
-    smoothed_states, smoothed_covs = ksmoother.build_graph(T, R, Q, filtered_states, filtered_covs)
-
-    outputs = [
-        filtered_states,
-        predicted_states,
-        smoothed_states,
-        filtered_covs,
-        predicted_covs,
-        smoothed_covs,
-        ll_obs.sum(),
-        ll_obs,
-    ]
-
-    return inputs, outputs
-
-
-def add_missing_data(data, n_missing, rng):
-    n = data.shape[0]
-    missing_idx = rng.choice(n, n_missing, replace=False)
-    data[missing_idx] = np.nan
-
-    return data
-
-
-def make_test_inputs(p, m, r, n, rng, missing_data=None, H_is_zero=False):
-    data = np.arange(n * p, dtype=floatX).reshape(-1, p)
-    if missing_data is not None:
-        data = add_missing_data(data, missing_data, rng)
-
-    a0 = np.zeros(m, dtype=floatX)
-    P0 = np.eye(m, dtype=floatX)
-    c = np.zeros(m, dtype=floatX)
-    d = np.zeros(p, dtype=floatX)
-    Q = np.eye(r, dtype=floatX)
-    H = np.zeros((p, p), dtype=floatX) if H_is_zero else np.eye(p, dtype=floatX)
-    T = np.eye(m, k=-1, dtype=floatX)
-    T[0, :] = 1 / m
-    R = np.eye(m, dtype=floatX)[:, :r]
-    Z = np.eye(m, dtype=floatX)[:p, :]
-
-    data, a0, P0, c, d, T, Z, R, H, Q = map(
-        np.ascontiguousarray, [data, a0, P0, c, d, T, Z, R, H, Q]
-    )
-
-    return data, a0, P0, c, d, T, Z, R, H, Q
-
-
-def get_expected_shape(name, p, m, r, n):
-    if name == "log_likelihood":
-        return ()
-    elif name == "ll_obs":
-        return (n,)
-    filter_type, variable = name.split("_")
-    if variable == "states":
-        return n, m
-    if variable == "covs":
-        return n, m, m
-
-
-def get_sm_state_from_output_name(res, name):
-    if name == "log_likelihood":
-        return res.llf
-    elif name == "ll_obs":
-        return res.llf_obs
-
-    filter_type, variable = name.split("_")
-    sm_states = getattr(res, "states")
-
-    if variable == "states":
-        return getattr(sm_states, filter_type)
-    if variable == "covs":
-        m = res.filter_results.k_states
-        # remove the "s" from "covs"
-        return getattr(sm_states, name[:-1]).reshape(-1, m, m)
-
-
-def nile_test_test_helper(rng, n_missing=0):
-    a0 = np.zeros(2, dtype=floatX)
-    P0 = np.eye(2, dtype=floatX) * 1e6
-    c = np.zeros(2, dtype=floatX)
-    d = np.zeros(1, dtype=floatX)
-    Q = np.eye(2, dtype=floatX) * np.array([0.5, 0.01], dtype=floatX)
-    H = np.eye(1, dtype=floatX) * 0.8
-    T = np.array([[1.0, 1.0], [0.0, 1.0]], dtype=floatX)
-    R = np.eye(2, dtype=floatX)
-    Z = np.array([[1.0, 0.0]], dtype=floatX)
-
-    data = load_nile_test_data().values
-    if n_missing > 0:
-        data = add_missing_data(data, n_missing, rng)
-
-    sm_model = LocalLinearTrend(
-        endog=data,
-        initialization="known",
-        initial_state_cov=P0,
-        initial_state=a0.ravel(),
-    )
-
-    res = sm_model.fit_constrained(
-        constraints={
-            "sigma2.measurement": 0.8,
-            "sigma2.level": 0.5,
-            "sigma2.trend": 0.01,
-        }
-    )
-
-    inputs = [data, a0, P0, c, d, T, Z, R, H, Q]
-
-    return res, inputs
-
-
-def fast_eval(var):
-    return pytensor.function([], var, mode="FAST_COMPILE")()
-
-
-def delete_rvs_from_model(rv_names: list[str]) -> None:
-    """Remove all model mappings referring to rv
-
-    This can be used to "delete" an RV from a model
-    """
-    mod = modelcontext(None)
-    all_rvs = mod.basic_RVs + mod.deterministics
-    all_rv_names = [x.name for x in all_rvs]
-
-    for name in rv_names:
-        assert name in all_rv_names, f"{name} is not part of the Model: {all_rv_names}"
-
-        rv_idx = all_rv_names.index(name)
-        rv = all_rvs[rv_idx]
-
-        mod.named_vars.pop(name)
-        if name in mod.named_vars_to_dims:
-            mod.named_vars_to_dims.pop(name)
-
-        if rv in mod.deterministics:
-            mod.deterministics.remove(rv)
-            continue
-
-        value = mod.rvs_to_values.pop(rv)
-        mod.values_to_rvs.pop(value)
-        mod.rvs_to_transforms.pop(rv)
-        if rv in mod.free_RVs:
-            mod.free_RVs.remove(rv)
-            mod.rvs_to_initial_values.pop(rv)
-        else:
-            mod.observed_RVs.remove(rv)
-
-
-def unpack_statespace(ssm):
-    return [ssm[SHORT_NAME_TO_LONG[x]] for x in MATRIX_NAMES]
-
-
-def unpack_symbolic_matrices_with_params(mod, param_dict, data_dict=None, mode="FAST_COMPILE"):
-    inputs = list(mod._name_to_variable.values())
-    if data_dict is not None:
-        inputs += list(mod._name_to_data.values())
-    else:
-        data_dict = {}
-
-    f_matrices = pytensor.function(
-        inputs,
-        unpack_statespace(mod.ssm),
-        on_unused_input="raise",
-        mode=mode,
-    )
-
-    x0, P0, c, d, T, Z, R, H, Q = f_matrices(**param_dict, **data_dict)
-
-    return x0, P0, c, d, T, Z, R, H, Q
-
-
-def simulate_from_numpy_model(mod, rng, param_dict, data_dict=None, steps=100):
-    """
-    Helper function to visualize the components outside of a PyMC model context
-    """
-    x0, P0, c, d, T, Z, R, H, Q = unpack_symbolic_matrices_with_params(mod, param_dict, data_dict)
-    k_states = mod.k_states
-    k_posdef = mod.k_posdef
-
-    x = np.zeros((steps, k_states))
-    y = np.zeros(steps)
-
-    x[0] = x0
-    y[0] = (Z @ x0).squeeze() if Z.ndim == 2 else (Z[0] @ x0).squeeze()
-
-    if not np.allclose(H, 0):
-        y[0] += rng.multivariate_normal(mean=np.zeros(1), cov=H).squeeze()
-
-    for t in range(1, steps):
-        if k_posdef > 0:
-            shock = rng.multivariate_normal(mean=np.zeros(k_posdef), cov=Q)
-            innov = R @ shock
-        else:
-            innov = 0
-
-        if not np.allclose(H, 0):
-            error = rng.multivariate_normal(mean=np.zeros(1), cov=H)
-        else:
-            error = 0
-
-        x[t] = c + T @ x[t - 1] + innov
-        if Z.ndim == 2:
-            y[t] = (d + Z @ x[t] + error).squeeze()
-        else:
-            y[t] = (d + Z[t] @ x[t] + error).squeeze()
-
-    return x, y
-
-
-def assert_pattern_repeats(y, T, atol, rtol):
-    val = np.diff(y.reshape(-1, T), axis=0)
-    if floatX.endswith("64"):
-        # Round this before going into the test, otherwise it behaves poorly (atol = inf)
-        n_digits = len(str(1 / atol))
-        val = np.round(val, n_digits)
-
-    assert_allclose(
-        val,
-        0,
-        err_msg="seasonal pattern does not repeat",
-        atol=atol,
-        rtol=rtol,
-    )
-
-
-def make_stationary_params(data, p, d, q, P, D, Q, S):
-    sm_sarimax = sm.tsa.SARIMAX(data, order=(p, d, q), seasonal_order=(P, D, Q, S))
-    res = sm_sarimax.fit(disp=False)
-
-    param_dict = dict(ar_params=[], ma_params=[], seasonal_ar_params=[], seasonal_ma_params=[])
-
-    for name, param in zip(res.param_names, res.params):
-        if name.startswith("ar.S"):
-            param_dict["seasonal_ar_params"].append(param)
-        elif name.startswith("ma.S"):
-            param_dict["seasonal_ma_params"].append(param)
-        elif name.startswith("ar."):
-            param_dict["ar_params"].append(param)
-        elif name.startswith("ma."):
-            param_dict["ma_params"].append(param)
-        else:
-            param_dict["sigma_state"] = param
-
-    param_dict = {
-        k: np.array(v, dtype=floatX)
-        for k, v in param_dict.items()
-        if isinstance(v, float) or len(v) > 0
-    }
-    return param_dict

tests/test_blackjax_smc.py

@@ -1,222 +0,0 @@
-#   Copyright 2023 The PyMC Developers
-#
-#   Licensed under the Apache License, Version 2.0 (the "License");
-#   you may not use this file except in compliance with the License.
-#   You may obtain a copy of the License at
-#
-#       http://www.apache.org/licenses/LICENSE-2.0
-#
-#   Unless required by applicable law or agreed to in writing, software
-#   distributed under the License is distributed on an "AS IS" BASIS,
-#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#   See the License for the specific language governing permissions and
-#   limitations under the License.
-import numpy as np
-import pymc as pm
-import pytensor.tensor as pt
-import pytest
-import scipy
-
-from numpy import dtype
-from xarray.core.utils import Frozen
-
-jax = pytest.importorskip("jax")
-pytest.importorskip("blackjax")
-
-from pymc_extras.inference.smc.sampling import (
-    arviz_from_particles,
-    blackjax_particles_from_pymc_population,
-    get_jaxified_loglikelihood,
-    get_jaxified_logprior,
-    sample_smc_blackjax,
-)
-
-
-def two_gaussians_model():
-    n = 4
-    mu1 = np.ones(n) * 0.5
-    mu2 = -mu1
-
-    stdev = 0.1
-    sigma = np.power(stdev, 2) * np.eye(n)
-    isigma = np.linalg.inv(sigma)
-    dsigma = np.linalg.det(sigma)
-
-    w1 = stdev
-    w2 = 1 - stdev
-
-    def two_gaussians(x):
-        """
-        Mixture of gaussians likelihood
-        """
-        log_like1 = (
-            -0.5 * n * pt.log(2 * np.pi)
-            - 0.5 * pt.log(dsigma)
-            - 0.5 * (x - mu1).T.dot(isigma).dot(x - mu1)
-        )
-        log_like2 = (
-            -0.5 * n * pt.log(2 * np.pi)
-            - 0.5 * pt.log(dsigma)
-            - 0.5 * (x - mu2).T.dot(isigma).dot(x - mu2)
-        )
-        return pt.log(w1 * pt.exp(log_like1) + w2 * pt.exp(log_like2))
-
-    with pm.Model() as m:
-        X = pm.Uniform("X", lower=-2, upper=2.0, shape=n)
-        llk = pm.Potential("muh", two_gaussians(X))
-
-    return m, mu1
-
-
-def fast_model():
-    with pm.Model() as m:
-        x = pm.Normal("x", 0, 1)
-        y = pm.Normal("y", x, 1, observed=0)
-    return m
-
-
-@pytest.mark.parametrize(
-    "kernel, check_for_integration_steps, inner_kernel_params",
-    [
-        ("HMC", True, {"step_size": 0.1, "integration_steps": 11}),
-        ("NUTS", False, {"step_size": 0.1}),
-    ],
-)
-def test_sample_smc_blackjax(kernel, check_for_integration_steps, inner_kernel_params):
-    """
-    When running the two gaussians model
-    with BlackJax SMC, we sample them correctly,
-    the shape of a posterior variable is (1, particles, dimension)
-    and the inference_data has the right attributes.
-
-    """
-    model, muref = two_gaussians_model()
-    iterations_to_diagnose = 2
-    n_particles = 1000
-    with model:
-        inference_data = sample_smc_blackjax(
-            n_particles=n_particles,
-            kernel=kernel,
-            inner_kernel_params=inner_kernel_params,
-            iterations_to_diagnose=iterations_to_diagnose,
-        )
-
-    x = inference_data.posterior["X"]
-
-    assert x.to_numpy().shape == (1, n_particles, 4)
-    mu1d = np.abs(x).mean(axis=0).mean(axis=0)
-    np.testing.assert_allclose(muref, mu1d, rtol=0.0, atol=0.03)
-
-    for attribute, value in [
-        ("particles", n_particles),
-        ("step_size", 0.1),
-        ("num_mcmc_steps", 10),
-        ("iterations_to_diagnose", iterations_to_diagnose),
-        ("sampler", f"Blackjax SMC with {kernel} kernel"),
-    ]:
-        assert inference_data.posterior.attrs[attribute] == value
-
-    for diagnostic in ["lambda_evolution", "log_likelihood_increments"]:
-        assert inference_data.posterior.attrs[diagnostic].shape == (iterations_to_diagnose,)
-
-    for diagnostic in ["ancestors_evolution", "weights_evolution"]:
-        assert inference_data.posterior.attrs[diagnostic].shape == (
-            iterations_to_diagnose,
-            n_particles,
-        )
-
-    for attribute in ["running_time_seconds", "iterations"]:
-        assert attribute in inference_data.posterior.attrs
-
-    if check_for_integration_steps:
-        assert inference_data.posterior.attrs["integration_steps"] == 11
-
-
-def test_blackjax_particles_from_pymc_population_univariate():
-    model = fast_model()
-    population = {"x": np.array([2, 3, 4])}
-    blackjax_particles = blackjax_particles_from_pymc_population(model, population)
-    jax.tree.map(np.testing.assert_allclose, blackjax_particles, [np.array([[2], [3], [4]])])
-
-
-def test_blackjax_particles_from_pymc_population_multivariate():
-    with pm.Model() as model:
-        x = pm.Normal("x", 0, 1)
-        z = pm.Normal("z", 0, 1)
-        y = pm.Normal("y", x + z, 1, observed=0)
-
-    population = {"x": np.array([0.34614613, 1.09163261, -0.44526825]), "z": np.array([1, 2, 3])}
-    blackjax_particles = blackjax_particles_from_pymc_population(model, population)
-    jax.tree.map(
-        np.testing.assert_allclose,
-        blackjax_particles,
-        [np.array([[0.34614613], [1.09163261], [-0.44526825]]), np.array([[1], [2], [3]])],
-    )
-
-
-def simple_multivariable_model():
-    """
-    A simple model that has a multivariate variable,
-    a has more than one variable (multivariable)
-    """
-    with pm.Model() as model:
-        x = pm.Normal("x", 0, 1, shape=2)
-        z = pm.Normal("z", 0, 1)
-        y = pm.Normal("y", z, 1, observed=0)
-    return model
-
-
-def test_blackjax_particles_from_pymc_population_multivariable():
-    model = simple_multivariable_model()
-    population = {"x": np.array([[2, 3], [5, 6], [7, 9]]), "z": np.array([11, 12, 13])}
-    blackjax_particles = blackjax_particles_from_pymc_population(model, population)
-
-    jax.tree.map(
-        np.testing.assert_allclose,
-        blackjax_particles,
-        [np.array([[2, 3], [5, 6], [7, 9]]), np.array([[11], [12], [13]])],
-    )
-
-
-def test_arviz_from_particles():
-    model = simple_multivariable_model()
-    particles = [np.array([[2, 3], [5, 6], [7, 9]]), np.array([[11], [12], [13]])]
-    with model:
-        inference_data = arviz_from_particles(model, particles)
-
-    assert inference_data.posterior.sizes == Frozen({"chain": 1, "draw": 3, "x_dim_0": 2})
-    assert inference_data.posterior.data_vars.dtypes == Frozen(
-        {"x": dtype("float64"), "z": dtype("float64")}
-    )
-
-
-def test_get_jaxified_logprior():
-    """
-    Given a model with a Normal prior
-    for a RV, the jaxified logprior
-    indeed calculates that number,
-    and can be jax.vmap'ed
-    """
-    logprior = get_jaxified_logprior(fast_model())
-    for point in [-0.5, 0.0, 0.5]:
-        jax.tree.map(
-            np.testing.assert_allclose,
-            jax.vmap(logprior)([np.array([point])]),
-            np.log(scipy.stats.norm(0, 1).pdf(point)),
-        )
-
-
-def test_get_jaxified_loglikelihood():
-    """
-    Given a model with a Normal Likelihood, a single observation
-    0 and std=1, the only free parameter of that function is the mean.
-    When computing the logliklikelihood
-    Then the function can be jax.vmap'ed, and the calculation matches the likelihood.
-    """
-    loglikelihood = get_jaxified_loglikelihood(fast_model())
-    for point in [-0.5, 0.0, 0.5]:
-        jax.tree.map(
-            np.testing.assert_allclose,
-            jax.vmap(loglikelihood)([np.array([point])]),
-            np.log(scipy.stats.norm(point, 1).pdf(0)),
-        )

tests/test_find_map.py

@@ -1,103 +0,0 @@
-import numpy as np
-import pymc as pm
-import pytensor.tensor as pt
-import pytest
-
-from pymc_extras.inference.find_map import (
-    GradientBackend,
-    find_MAP,
-    scipy_optimize_funcs_from_loss,
-)
-
-pytest.importorskip("jax")
-
-
-@pytest.fixture(scope="session")
-def rng():
-    seed = sum(map(ord, "test_fit_map"))
-    return np.random.default_rng(seed)
-
-
-@pytest.mark.parametrize("gradient_backend", ["jax", "pytensor"], ids=str)
-def test_jax_functions_from_graph(gradient_backend: GradientBackend):
-    x = pt.tensor("x", shape=(2,))
-
-    def compute_z(x):
-        z1 = x[0] ** 2 + 2
-        z2 = x[0] * x[1] + 3
-        return z1, z2
-
-    z = pt.stack(compute_z(x))
-    f_loss, f_hess, f_hessp = scipy_optimize_funcs_from_loss(
-        loss=z.sum(),
-        inputs=[x],
-        initial_point_dict={"x": np.array([1.0, 2.0])},
-        use_grad=True,
-        use_hess=True,
-        use_hessp=True,
-        gradient_backend=gradient_backend,
-        compile_kwargs=dict(mode="JAX"),
-    )
-
-    x_val = np.array([1.0, 2.0])
-    expected_z = sum(compute_z(x_val))
-
-    z_jax, grad_val = f_loss(x_val)
-    np.testing.assert_allclose(z_jax, expected_z)
-    np.testing.assert_allclose(grad_val.squeeze(), np.array([2 * x_val[0] + x_val[1], x_val[0]]))
-
-    hess_val = np.array(f_hess(x_val))
-    np.testing.assert_allclose(hess_val.squeeze(), np.array([[2, 1], [1, 0]]))
-
-    hessp_val = np.array(f_hessp(x_val, np.array([1.0, 0.0])))
-    np.testing.assert_allclose(hessp_val.squeeze(), np.array([2, 1]))
-
-
-@pytest.mark.parametrize(
-    "method, use_grad, use_hess, use_hessp",
-    [
-        ("nelder-mead", False, False, False),
-        ("powell", False, False, False),
-        ("CG", True, False, False),
-        ("BFGS", True, False, False),
-        ("L-BFGS-B", True, False, False),
-        ("TNC", True, False, False),
-        ("SLSQP", True, False, False),
-        ("dogleg", True, True, False),
-        ("Newton-CG", True, True, False),
-        ("Newton-CG", True, False, True),
-        ("trust-ncg", True, True, False),
-        ("trust-ncg", True, False, True),
-        ("trust-exact", True, True, False),
-        ("trust-krylov", True, True, False),
-        ("trust-krylov", True, False, True),
-        ("trust-constr", True, True, False),
-    ],
-)
-@pytest.mark.parametrize("gradient_backend", ["jax", "pytensor"], ids=str)
-def test_JAX_map(method, use_grad, use_hess, use_hessp, gradient_backend: GradientBackend, rng):
-    extra_kwargs = {}
-    if method == "dogleg":
-        # HACK -- dogleg requires that the hessian of the objective function is PSD, so we have to pick a point
-        # where this is true
-        extra_kwargs = {"initvals": {"mu": 2, "sigma_log__": 1}}
-
-    with pm.Model() as m:
-        mu = pm.Normal("mu")
-        sigma = pm.Exponential("sigma", 1)
-        pm.Normal("y_hat", mu=mu, sigma=sigma, observed=rng.normal(loc=3, scale=1.5, size=100))
-
-        optimized_point = find_MAP(
-            method=method,
-            **extra_kwargs,
-            use_grad=use_grad,
-            use_hess=use_hess,
-            use_hessp=use_hessp,
-            progressbar=False,
-            gradient_backend=gradient_backend,
-            compile_kwargs={"mode": "JAX"},
-        )
-    mu_hat, log_sigma_hat = optimized_point["mu"], optimized_point["sigma_log__"]
-
-    assert np.isclose(mu_hat, 3, atol=0.5)
-    assert np.isclose(np.exp(log_sigma_hat), 1.5, atol=0.5)