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

{pymc_extras-0.2.4.dist-info → pymc_extras-0.2.5.dist-info}/RECORD

@@ -1,26 +1,28 @@
-pymc_extras/__init__.py,sha256=IFIEZdPX_Ugq57Bu7jlyrJLpKng-P0FBAAAzl2pFXLE,1266
+pymc_extras/__init__.py,sha256=lYGf9TcwUHROIElkX7Epnb7-IppcmiSEYuxdtRzqS3s,1195
 pymc_extras/linearmodel.py,sha256=6eitl15Ec15mSZu7zoHZ7Wwy4U1DPwqfAgwEt6ILeIc,3920
 pymc_extras/model_builder.py,sha256=sAw77fxdiy046BvDPjocuMlbJ0Efj-CDAGtmcwYmoG0,26361
 pymc_extras/printing.py,sha256=G8mj9dRd6i0PcsbcEWZm56ek6V8mmil78RI4MUhywBs,6506
 pymc_extras/version.py,sha256=VxPGCBzhtSegu-Jp5cjzn0n4DGU0wuPUh-KyZKB6uPM,240
-pymc_extras/version.txt,sha256=FyW_plJMUmXnwXHPBlaEF9OblH__ScJC8DhZR5yCM0s,6
-pymc_extras/distributions/__init__.py,sha256=gTX7tvX8NcgP7V72URV7GeqF1aAEjGVbuW8LMxhXceY,1295
+pymc_extras/version.txt,sha256=6Vn3UOktu3YUriislvCjcnLK7YHYu7dYeRr3v7thBqA,6
+pymc_extras/distributions/__init__.py,sha256=fDbrBt9mxEVp2CDPwnyCW3oiutzZ0PduB8EUH3fUrjI,1377
 pymc_extras/distributions/continuous.py,sha256=z-nvQgGncYISdRY8cWsa-56V0bQGq70jYwU-i8VZ0Uk,11253
 pymc_extras/distributions/discrete.py,sha256=vrARNuiQAEXrs7yQgImV1PO8AV1uyEC_LBhr6F9IcOg,13032
 pymc_extras/distributions/histogram_utils.py,sha256=5RTvlGCUrp2qzshrchmPyWxjhs6RIYL62SMikjDM1TU,5814
 pymc_extras/distributions/timeseries.py,sha256=M5MZ-nik_tgkaoZ1hdUGEZ9g04DQyVLwszVJqSKwNcY,12719
 pymc_extras/distributions/multivariate/__init__.py,sha256=E8OeLW9tTotCbrUjEo4um76-_WQD56PehsPzkKmhfyA,93
 pymc_extras/distributions/multivariate/r2d2m2cp.py,sha256=bUj9bB-hQi6CpaJfvJjgNPi727uTbvAdxl9fm1zNBqY,16005
+pymc_extras/distributions/transforms/__init__.py,sha256=FUp2vyRE6_2eUcQ_FVt5Dn0-vy5I-puV-Kz13-QtLNc,104
+pymc_extras/distributions/transforms/partial_order.py,sha256=oEZlc9WgnGR46uFEjLzKEUxlhzIo2vrUUbBE3vYrsfQ,8404
 pymc_extras/gp/__init__.py,sha256=sFHw2y3lEl5tG_FDQHZUonQ_k0DF1JRf0Rp8dpHmge0,745
 pymc_extras/gp/latent_approx.py,sha256=cDEMM6H1BL2qyKg7BZU-ISrKn2HJe7hDaM4Y8GgQDf4,6682
-pymc_extras/inference/__init__.py,sha256=5cXpaQQnW0mJJ3x8wSxmYu63l--Xab5D_gMtjA6Q3uU,666
+pymc_extras/inference/__init__.py,sha256=UH6S0bGfQKKyTSuqf7yezdy9PeE2bDU8U1v4eIRv4ZI,887
 pymc_extras/inference/find_map.py,sha256=vl5l0ei48PnX-uTuHVTr-9QpCEHc8xog-KK6sOnJ8LU,16513
-pymc_extras/inference/fit.py,sha256=S9R48dh74s6K0MC9Iys4NAwVjP6rVRfx6SF-kPiR70E,1165
-pymc_extras/inference/laplace.py,sha256=uOZGp8ssQuhvCHV_Y_v3icsr4rhcYgr_qlr9dS7pcSM,21761
+pymc_extras/inference/fit.py,sha256=oe20RAajImZ-VD9Ucbzri8Bof4Y2KHNhNRG19v9O3lI,1336
+pymc_extras/inference/laplace.py,sha256=cqarAdbFaOH74AkPUF4c7c_Hswa5mqmhgHpsgrkebHY,21860
 pymc_extras/inference/pathfinder/__init__.py,sha256=FhAYrCWNx_dCrynEdjg2CZ9tIinvcVLBm67pNx_Y3kA,101
 pymc_extras/inference/pathfinder/importance_sampling.py,sha256=NwxepXOFit3cA5zEebniKdlnJ1rZWg56aMlH4MEOcG4,6264
-pymc_extras/inference/pathfinder/lbfgs.py,sha256=P0UIOVtspdLzDU6alK-y91qzVAzXjYAXPuGmZ1nRqMo,5715
-pymc_extras/inference/pathfinder/pathfinder.py,sha256=baw8NUN4hdylM0o4JpCh32xxig-fNFLjh_W9qsvvmM0,64495
+pymc_extras/inference/pathfinder/lbfgs.py,sha256=GOoJBil5Kft_iFwGNUGKSeqzI5x_shA4KQWDwgGuQtQ,7110
+pymc_extras/inference/pathfinder/pathfinder.py,sha256=GW04HQurj_3Nlo1C6_K2tEIeigo8x0buV3FqDLA88PQ,64439
 pymc_extras/inference/smc/__init__.py,sha256=wyaT4NJl1YsSQRLiDy-i0Jq3CbJZ2BQd4nnCk-dIngY,603
 pymc_extras/inference/smc/sampling.py,sha256=AYwmKqGoV6pBtKnh9SUbBKbN7VcoFgb3MmNWV7SivMA,15365
 pymc_extras/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -52,21 +54,22 @@ pymc_extras/statespace/models/utilities.py,sha256=G9GuHKsghmIYOlfkPtvxBWF-FZY5-5
 pymc_extras/statespace/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/statespace/utils/constants.py,sha256=Kf6j75ABaDQeRODxKQ76wTUQV4F5sTjn1KBcZgCQx20,2403
 pymc_extras/statespace/utils/coord_tools.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pymc_extras/statespace/utils/data_tools.py,sha256=caanvrxDu9g-dEKff2bbmaVTs6-71kkSoYIiiSUXhw4,5985
+pymc_extras/statespace/utils/data_tools.py,sha256=01sz6XDtLYK9I5xghxYpD-PuDzGXv9D-wFGfTV6FGEw,6566
 pymc_extras/utils/__init__.py,sha256=yxI9cJ7fCtVQS0GFw0y6mDGZIQZiK53vm3UNKqIuGSk,758
 pymc_extras/utils/linear_cg.py,sha256=KkXhuimFsrKtNd_0By2ApxQQQNm5FdBtmDQJOVbLYkA,10056
 pymc_extras/utils/model_equivalence.py,sha256=8QIftID2HDxD659i0RXHazQ-l2Q5YegCRLcDqb2p9Pc,2187
 pymc_extras/utils/pivoted_cholesky.py,sha256=QtnjP0pAl9b77fLAu-semwT4_9dcoiqx3dz1xKGBjMk,1871
 pymc_extras/utils/prior.py,sha256=QlWVr7uKIK9VncBw7Fz3YgaASKGDfqpORZHc-vz_9gQ,6841
 pymc_extras/utils/spline.py,sha256=qGq0gcoMG5dpdazKFzG0RXkkCWP8ADPPXN-653-oFn4,4820
+pymc_extras-0.2.5.dist-info/licenses/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
 tests/__init__.py,sha256=-ree9OWVCyTeXLR944OWjrQX2os15HXrRNkhJ7QdRjc,603
 tests/test_blackjax_smc.py,sha256=jcNgcMBxaKyPg9UvHnWQtwoL79LXlSpZfALe3RGEZnQ,7233
 tests/test_find_map.py,sha256=B8ThnXNyfTQeem24QaLoTitFrsxKoq2VQINUdOwzna0,3379
 tests/test_histogram_approximation.py,sha256=w-xb2Rr0Qft6sm6F3BTmXXnpuqyefC1SUL6YxzqA5X4,4674
-tests/test_laplace.py,sha256=u4o-0y4v1emaTMYr_rOyL_EKY_bQIz0DUXFuwuDbfNg,9314
+tests/test_laplace.py,sha256=fArHjwMR7x98K-gZLvrvb3AwNZ7_fo7E0A4SJyt4EGU,9843
 tests/test_linearmodel.py,sha256=iB8ApNqIX9_nUHoo-Tm51xuPdrva5t4VLLut6qXB5Ao,6906
 tests/test_model_builder.py,sha256=QiINEihBR9rx8xM4Nqlg4urZKoyo58aTKDtxl9SJF1s,11249
-tests/test_pathfinder.py,sha256=-ekjetUSWnNRe7YausDvD00Cqh0zpBW3xn5z1hJ37MI,6027
+tests/test_pathfinder.py,sha256=vlMI1p2Ja5X4QIaSV4h6U41I303rEppfO0JqE3xe1Rs,10023
 tests/test_pivoted_cholesky.py,sha256=PuMdMSCzO4KdQWpUF4SEBeuH_qsINCIH8TYtmmJ1NKo,692
 tests/test_printing.py,sha256=HnvwwjrjBuxXFAJdyU0K_lvKGLgh4nzHAnhsIUpenbY,5211
 tests/test_prior_from_trace.py,sha256=HOzR3l98pl7TEJquo_kSugED4wBTgHo4-8lgnpmacs8,5516
@@ -77,6 +80,7 @@ tests/distributions/test_continuous.py,sha256=1-bu-IP6RgLUJnuPYpOD8ZS1ahYbKtsJ9o
 tests/distributions/test_discrete.py,sha256=CjjaUpppsvQ6zLzV15ZsbwNOKrDmEdz4VWcleoCXUi0,7776
 tests/distributions/test_discrete_markov_chain.py,sha256=8RCHZXSB8IWjniuKaGGlM_iTWGmdrcOqginxmrAeEJg,9212
 tests/distributions/test_multivariate.py,sha256=LBvBuoT_3rzi8rR38b8L441Y-9Ff0cIXeRBKiEn6kjs,10452
+tests/distributions/test_transform.py,sha256=QM9sSQ5eSbuT2pM76nUMWqb-tQa7DGZbT9uwFDqIRUk,2672
 tests/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/model/test_model_api.py,sha256=FJvMTmexovRELZOUcUyk-6Vwk9qSiH7hIFoiArgl5mk,1040
 tests/model/marginal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -87,7 +91,7 @@ tests/statespace/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 tests/statespace/test_ETS.py,sha256=IPg3uQ7xEGqDMEHu993vtUTV7r-uNAxmw23sr5MVGfQ,15582
 tests/statespace/test_SARIMAX.py,sha256=1BYNOm9aSHnpn-qbpe3YsQVH8m-mXcp_gvKgWhWn1W4,12948
 tests/statespace/test_VARMAX.py,sha256=rJnea9_WEGo9I0iv2eaSbwwFQv0tlIjpN7KAE0eQewU,6336
-tests/statespace/test_coord_assignment.py,sha256=2GBm46-0eI4QNh4bvp3D7az58stcA5Zo6VgOo_JkCig,3821
+tests/statespace/test_coord_assignment.py,sha256=2Mo5196ibkBTwscE7kqQoUsgQphdaagVkOccDi7D4RI,5980
 tests/statespace/test_distributions.py,sha256=WQ_ROyd-PL3cimXTyEtyVaMEVtS7Hue2Z0lN7UnGDyo,9122
 tests/statespace/test_kalman_filter.py,sha256=s2n62FzXl9elU_uqaMNaEaexUfq3SXe3_YvQ2lM6hiQ,11600
 tests/statespace/test_representation.py,sha256=1KAJY4ZaVhb1WdAJLx2UYSXuVYsMNWX98gEDF7P0B4s,6210
@@ -98,8 +102,7 @@ tests/statespace/utilities/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJ
 tests/statespace/utilities/shared_fixtures.py,sha256=SNw8Bvj1Yw11TxAW6n20Bq0B8oaYtVTiFFEVNH_wnp4,164
 tests/statespace/utilities/statsmodel_local_level.py,sha256=SQAzaYaSDwiVhUQ1iWjt4MgfAd54RuzVtnslIs3xdS8,1225
 tests/statespace/utilities/test_helpers.py,sha256=oH24a6Q45NFFFI3Kx9mhKbxsCvo9ErCorKFoTjDB3-4,9159
-pymc_extras-0.2.4.dist-info/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
-pymc_extras-0.2.4.dist-info/METADATA,sha256=ozmK251JzJsJLI9yx8NFhhVCgOy5nfcfSfE5IfTP3ok,5154
-pymc_extras-0.2.4.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-pymc_extras-0.2.4.dist-info/top_level.txt,sha256=D6RkgBiXiZCel0nvsYg_zYEoT_VuwocyIY98EMaulj0,18
-pymc_extras-0.2.4.dist-info/RECORD,,
+pymc_extras-0.2.5.dist-info/METADATA,sha256=02v5liTQQ55sV8xeFl5EjFpwbOKSYKG6g5lE_4htpBo,5227
+pymc_extras-0.2.5.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+pymc_extras-0.2.5.dist-info/top_level.txt,sha256=D6RkgBiXiZCel0nvsYg_zYEoT_VuwocyIY98EMaulj0,18
+pymc_extras-0.2.5.dist-info/RECORD,,

{pymc_extras-0.2.4.dist-info → pymc_extras-0.2.5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (76.0.0)
+Generator: setuptools (78.1.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

tests/distributions/test_transform.py

@@ -0,0 +1,77 @@
+#   Copyright 2025 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
+
+from pymc_extras.distributions.transforms import PartialOrder
+
+
+class TestPartialOrder:
+    adj_mats = np.array(
+        [
+            # 0 < {1, 2} < 3
+            [[0, 1, 1, 0], [0, 0, 0, 1], [0, 0, 0, 1], [0, 0, 0, 0]],
+            # 1 < 0 < 3 < 2
+            [[0, 0, 0, 1], [1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 1, 0]],
+        ]
+    )
+
+    valid_values = np.array([[0, 2, 1, 3], [1, 0, 3, 2]], dtype=float)
+
+    # Test that forward and backward are inverses of eachother
+    # And that it works when extra dimensions are added in data
+    def test_forward_backward_dimensionality(self):
+        po = PartialOrder(self.adj_mats)
+        po0 = PartialOrder(self.adj_mats[0])
+        vv = self.valid_values
+        vv0 = self.valid_values[0]
+
+        testsets = [
+            (vv, po),
+            (po.initvals(), po),
+            (vv0, po0),
+            (po0.initvals(), po0),
+            (np.tile(vv0, (2, 1)), po0),
+            (np.tile(vv0, (2, 3, 2, 1)), po0),
+            (np.tile(vv, (2, 3, 2, 1, 1)), po),
+        ]
+
+        for vv, po in testsets:
+            fw = po.forward(vv)
+            bw = po.backward(fw)
+            np.testing.assert_allclose(bw.eval(), vv)
+
+    def test_sample_model(self):
+        po = PartialOrder(self.adj_mats)
+        with pm.Model() as model:
+            x = pm.Normal(
+                "x",
+                size=(3, 2, 4),
+                transform=po,
+                initval=po.initvals(shape=(3, 2, 4), lower=-1, upper=1),
+            )
+            idata = pm.sample()
+
+        # Check that the order constraints are satisfied
+        # Move chain, draw and "3" dimensions to the back
+        xvs = idata.posterior.x.values.transpose(3, 4, 0, 1, 2)
+        x0 = xvs[0]  # 0 < {1, 2} < 3
+        assert (
+            (x0[0] < x0[1]).all()
+            and (x0[0] < x0[2]).all()
+            and (x0[1] < x0[3]).all()
+            and (x0[2] < x0[3]).all()
+        )
+        x1 = xvs[1]  # 1 < 0 < 3 < 2
+        assert (x1[1] < x1[0]).all() and (x1[0] < x1[3]).all() and (x1[3] < x1[2]).all()

tests/statespace/test_coord_assignment.py

@@ -8,6 +8,7 @@ import pytensor.tensor as pt
 import pytest
 
 from pymc_extras.statespace.models import structural
+from pymc_extras.statespace.models.structural import LevelTrendComponent
 from pymc_extras.statespace.utils.constants import (
     FILTER_OUTPUT_DIMS,
     FILTER_OUTPUT_NAMES,
@@ -114,3 +115,67 @@ def test_data_index_is_coord(f, warning, create_model):
     with warning:
         pymc_model = create_model(f)
     assert TIME_DIM in pymc_model.coords
+
+
+def make_model(index):
+    n = len(index)
+    a = pd.DataFrame(index=index, columns=["A", "B", "C", "D"], data=np.arange(n * 4).reshape(n, 4))
+
+    mod = LevelTrendComponent(order=2, innovations_order=[0, 1])
+    ss_mod = mod.build(name="a", verbose=False)
+
+    initial_trend_dims, sigma_trend_dims, P0_dims = ss_mod.param_dims.values()
+    coords = ss_mod.coords
+
+    with pm.Model(coords=coords) as model:
+        P0_diag = pm.Gamma("P0_diag", alpha=5, beta=5)
+        P0 = pm.Deterministic("P0", pt.eye(ss_mod.k_states) * P0_diag, dims=P0_dims)
+
+        initial_trend = pm.Normal("initial_trend", dims=initial_trend_dims)
+        sigma_trend = pm.Gamma("sigma_trend", alpha=2, beta=50, dims=sigma_trend_dims)
+
+        with pytest.warns(UserWarning, match="No time index found on the supplied data"):
+            ss_mod.build_statespace_graph(
+                a["A"],
+                mode="JAX",
+            )
+    return model
+
+
+def test_integer_index():
+    index = np.arange(8).astype(int)
+    model = make_model(index)
+    assert TIME_DIM in model.coords
+    np.testing.assert_allclose(model.coords[TIME_DIM], index)
+
+
+def test_float_index_raises():
+    index = np.linspace(0, 1, 8)
+
+    with pytest.raises(IndexError, match="Provided index is not an integer index"):
+        make_model(index)
+
+
+def test_non_strictly_monotone_index_raises():
+    # Decreases
+    index = [0, 1, 2, 1, 2, 3]
+    with pytest.raises(IndexError, match="Provided index is not monotonic increasing"):
+        make_model(index)
+
+    # Has gaps
+    index = [0, 1, 2, 3, 5, 6]
+    with pytest.raises(IndexError, match="Provided index is not monotonic increasing"):
+        make_model(index)
+
+    # Has duplicates
+    index = [0, 1, 1, 2, 3, 4]
+    with pytest.raises(IndexError, match="Provided index is not monotonic increasing"):
+        make_model(index)
+
+
+def test_multiindex_raises():
+    index = pd.MultiIndex.from_tuples([(0, 0), (1, 1), (2, 2), (3, 3)])
+    with pytest.raises(
+        NotImplementedError, match="MultiIndex panel data is not currently supported"
+    ):
+        make_model(index)

tests/test_laplace.py

@@ -263,3 +263,19 @@ def test_fit_laplace(fit_in_unconstrained_space, mode, gradient_backend: Gradien
         else:
             assert idata.fit.rows.values.tolist() == ["mu", "sigma"]
             np.testing.assert_allclose(idata.fit.mean_vector.values, np.array([3.0, 1.5]), atol=0.1)
+
+
+def test_laplace_scalar():
+    # Example model from Statistical Rethinking
+    data = np.array([0, 0, 0, 1, 1, 1, 1, 1, 1])
+
+    with pm.Model():
+        p = pm.Uniform("p", 0, 1)
+        w = pm.Binomial("w", n=len(data), p=p, observed=data.sum())
+
+        idata_laplace = pmx.fit_laplace(progressbar=False)
+
+    assert idata_laplace.fit.mean_vector.shape == (1,)
+    assert idata_laplace.fit.covariance_matrix.shape == (1, 1)
+
+    np.testing.assert_allclose(idata_laplace.fit.mean_vector.values.item(), data.mean(), atol=0.1)

tests/test_pathfinder.py

@@ -12,14 +12,14 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+import re
 import sys
 
 import numpy as np
 import pymc as pm
+import pytensor.tensor as pt
 import pytest
 
-pytestmark = pytest.mark.filterwarnings("ignore:compile_pymc was renamed to compile:FutureWarning")
-
 import pymc_extras as pmx
 
 
@@ -52,7 +52,90 @@ def reference_idata():
     return idata
 
 
+def unstable_lbfgs_update_mask_model() -> pm.Model:
+    # data and model from: https://github.com/pymc-devs/pymc-extras/issues/445
+    # this scenario made LBFGS struggle leading to a lot of rejected iterations, (result.nit being moderate, but only history.count <= 1).
+    # this scenario is used to test that the LBFGS history manager is rejecting iterations as expected and PF can run to completion.
+
+    # fmt: off
+    inp = np.array([0, 0, 1, 0, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 2, 0, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 2, 0, 1, 0, 0, 0, 0, 1, 1, 1, 2, 0, 1, 2, 1, 0, 1, 0, 1, 0, 1, 0])
+
+    res = np.array([[0,0,1,0,0],[0,0,1,0,0],[0,0,1,0,0],[0,1,0,0,0],[0,0,1,0,0],[0,1,0,0,0],[0,0,0,1,0],[0,0,0,1,0],[1,0,0,0,0],[0,1,0,0,0],[0,0,1,0,0],[1,0,0,0,0],[0,0,1,0,0],[0,1,0,0,0],[0,0,0,1,0],[0,0,1,0,0],[0,0,1,0,0],[0,0,0,1,0],[0,0,0,1,0],[0,1,0,0,0],[0,0,0,1,0],[0,0,1,0,0],[0,1,0,0,0],[1,0,0,0,0],[1,0,0,0,0],[0,0,1,0,0],[0,0,1,0,0],[0,0,1,0,0],[0,0,1,0,0],[0,1,0,0,0],[0,1,0,0,0],[0,0,0,1,0],[0,0,1,0,0],[0,1,0,0,0],[0,0,1,0,0],[0,0,0,1,0],[0,0,1,0,0],[1,0,0,0,0],[1,0,0,0,0],[0,0,1,0,0],[0,0,0,1,0],[0,0,0,1,0],[1,0,0,0,0],[1,0,0,0,0],[0,1,0,0,0],[1,0,0,0,0],[0,0,1,0,0],[0,0,1,0,0],[1,0,0,0,0],[0,0,0,1,0]])
+    # fmt: on
+
+    n_ordered = res.shape[1]
+    coords = {
+        "obs": np.arange(len(inp)),
+        "inp": np.arange(max(inp) + 1),
+        "outp": np.arange(res.shape[1]),
+    }
+    with pm.Model(coords=coords) as mdl:
+        mu = pm.Normal("intercept", sigma=3.5)[None]
+
+        offset = pm.Normal(
+            "offset", dims=("inp"), transform=pm.distributions.transforms.ZeroSumTransform([0])
+        )
+
+        scale = 3.5 * pm.HalfStudentT("scale", nu=5)
+        mu += (scale * offset)[inp]
+
+        phi_delta = pm.Dirichlet("phi_diffs", [1.0] * (n_ordered - 1))
+        phi = pt.concatenate([[0], pt.cumsum(phi_delta)])
+        s_mu = pm.Normal(
+            "stereotype_intercept",
+            size=n_ordered,
+            transform=pm.distributions.transforms.ZeroSumTransform([-1]),
+        )
+        fprobs = pm.math.softmax(s_mu[None, :] + phi[None, :] * mu[:, None], axis=-1)
+
+        pm.Multinomial("y_res", p=fprobs, n=np.ones(len(inp)), observed=res, dims=("obs", "outp"))
+
+    return mdl
+
+
+@pytest.mark.parametrize("jitter", [12.0, 500.0, 1000.0])
+def test_unstable_lbfgs_update_mask(capsys, jitter):
+    model = unstable_lbfgs_update_mask_model()
+
+    if jitter < 1000:
+        with model:
+            idata = pmx.fit(
+                method="pathfinder",
+                jitter=jitter,
+                random_seed=4,
+            )
+        out, err = capsys.readouterr()
+        status_pattern = [
+            r"INIT_FAILED_LOW_UPDATE_PCT\s+\d+",
+            r"LOW_UPDATE_PCT\s+\d+",
+            r"LBFGS_FAILED\s+\d+",
+            r"SUCCESS\s+\d+",
+        ]
+        for pattern in status_pattern:
+            assert re.search(pattern, out) is not None
+
+    else:
+        with pytest.raises(ValueError, match="All paths failed"):
+            with model:
+                idata = pmx.fit(
+                    method="pathfinder",
+                    jitter=1000,
+                    random_seed=2,
+                    num_paths=4,
+                )
+            out, err = capsys.readouterr()
+
+            status_pattern = [
+                r"INIT_FAILED_LOW_UPDATE_PCT\s+2",
+                r"LOW_UPDATE_PCT\s+2",
+                r"LBFGS_FAILED\s+4",
+            ]
+            for pattern in status_pattern:
+                assert re.search(pattern, out) is not None
+
+
 @pytest.mark.parametrize("inference_backend", ["pymc", "blackjax"])
+@pytest.mark.filterwarnings("ignore:JAXopt is no longer maintained.:DeprecationWarning")
 def test_pathfinder(inference_backend, reference_idata):
     if inference_backend == "blackjax" and sys.platform == "win32":
         pytest.skip("JAX not supported on windows")
@@ -149,12 +232,11 @@ def test_bfgs_sample():
     # get factors
     x_full = pt.as_tensor(x_data, dtype="float64")
     g_full = pt.as_tensor(g_data, dtype="float64")
-    epsilon = 1e-11
 
     x = x_full[1:]
     g = g_full[1:]
-    alpha, S, Z, update_mask = alpha_recover(x_full, g_full, epsilon)
-    beta, gamma = inverse_hessian_factors(alpha, S, Z, update_mask, J)
+    alpha, s, z = alpha_recover(x_full, g_full)
+    beta, gamma = inverse_hessian_factors(alpha, s, z, J)
 
     # sample
     phi, logq = bfgs_sample(
@@ -169,8 +251,8 @@ def test_bfgs_sample():
     # check shapes
     assert beta.eval().shape == (L, N, 2 * J)
     assert gamma.eval().shape == (L, 2 * J, 2 * J)
-    assert phi.eval().shape == (L, num_samples, N)
-    assert logq.eval().shape == (L, num_samples)
+    assert all(phi.shape.eval() == (L, num_samples, N))
+    assert all(logq.shape.eval() == (L, num_samples))
 
 
 @pytest.mark.parametrize("importance_sampling", ["psis", "psir", "identity", None])
@@ -201,3 +283,15 @@ def test_pathfinder_importance_sampling(importance_sampling):
         assert idata.posterior["mu"].shape == (1, num_draws)
         assert idata.posterior["tau"].shape == (1, num_draws)
         assert idata.posterior["theta"].shape == (1, num_draws, 8)
+
+
+def test_pathfinder_initvals():
+    # Run a model with an ordered transform that will fail unless initvals are in place
+    with pm.Model() as mdl:
+        pm.Normal("ordered", size=10, transform=pm.distributions.transforms.ordered)
+        idata = pmx.fit_pathfinder(initvals={"ordered": np.linspace(0, 1, 10)})
+
+    # Check that the samples are ordered to make sure transform was applied
+    assert np.all(
+        idata.posterior["ordered"][..., 1:].values > idata.posterior["ordered"][..., :-1].values
+    )