Various fixes

Author: fonnesbeck

pymc_extras/distributions/histogram_utils.py

@@ -15,7 +15,6 @@
 
 import numpy as np
 import pymc as pm
-import pymc_extras as pmx
 
 from numpy.typing import ArrayLike
 

pymc_extras/model/marginal/distributions.py

@@ -92,7 +92,6 @@ def reduce_batch_dependent_logps(
        as well as transpose the remaining axis of dep1 logp before adding the two element-wise.
 
     """
-    from pymc_extras.model.marginal.graph_analysis import get_support_axes
 
     reduced_logps = []
     for dependent_op, dependent_logp, dependent_dims_connection in zip(

pymc_extras/model/model_api.py

@@ -1,7 +1,6 @@
 from functools import wraps
 
 from pymc import Model
-import pymc_extras as pmx
 
 
 def as_model(*model_args, **model_kwargs):

pymc_extras/statespace/models/ETS.py

@@ -10,10 +10,6 @@
 from pymc_extras.statespace.core.statespace import PyMCStateSpace, floatX
 from pymc_extras.statespace.models.utilities import make_default_coords
 from pymc_extras.statespace.utils.constants import (
-    JITTER_DEFAULT,
-    LONG_MATRIX_NAMES,
-    MISSING_FILL,
-    SHORT_NAME_TO_LONG,
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
     ETS_SEASONAL_DIM,

pymc_extras/statespace/models/VARMAX.py

@@ -18,12 +18,7 @@
     OBS_STATE_DIM,
     SHOCK_AUX_DIM,
     SHOCK_DIM,
-    JITTER_DEFAULT,
-    LONG_MATRIX_NAMES,
-    MISSING_FILL,
-    SHORT_NAME_TO_LONG,
 )
-import pymc_extras.statespace as pmss
 
 floatX = pytensor.config.floatX
 

pymc_extras/statespace/models/utilities.py

@@ -11,9 +11,6 @@
     SHOCK_AUX_DIM,
     SHOCK_DIM,
     VECTOR_VALUED,
-    JITTER_DEFAULT,
-    MISSING_FILL,
-    SHORT_NAME_TO_LONG,
 )
 
 

tests/model/marginal/test_graph_analysis.py

@@ -5,8 +5,6 @@
 from pytensor.tensor.type_other import NoneTypeT
 
 from pymc_extras.model.marginal.graph_analysis import (
-    find_conditional_dependent_rvs,
-    find_conditional_input_rvs,
     is_conditional_dependent,
     subgraph_batch_dim_connection,
 )

tests/test_blackjax_smc.py

@@ -20,28 +20,17 @@
 from numpy import dtype
 from xarray.core.utils import Frozen
 
-jax = pytest.importorskip("jax")
-pytest.importorskip("blackjax")
-
 from pymc_extras.inference.smc.sampling import (
-    SMC_KERNEL_DEFAULTS,
-    SMC_KERNEL_NAMES,
-    SMC_KERNEL_PARAMS,
-    SMC_KERNEL_PARAMS_DEFAULTS,
-    SMC_KERNEL_PARAMS_NAMES,
-    SMC_KERNEL_PARAMS_TYPES,
-    SMC_KERNEL_TYPES,
-    SMC_KERNEL_VALID_PARAMS,
-    SMC_KERNEL_VALID_PARAMS_TYPES,
-    SMC_KERNELS,
-    SMC_KERNELS_PARAMS,
-    SMC_KERNELS_PARAMS_DEFAULTS,
-    SMC_KERNELS_PARAMS_NAMES,
-    SMC_KERNELS_PARAMS_TYPES,
-    SMC_KERNELS_VALID_PARAMS,
-    SMC_KERNELS_VALID_PARAMS_TYPES,
+    arviz_from_particles,
+    blackjax_particles_from_pymc_population,
+    get_jaxified_loglikelihood,
+    get_jaxified_logprior,
+    sample_smc_blackjax,
 )
 
+jax = pytest.importorskip("jax")
+pytest.importorskip("blackjax")
+
 
 def two_gaussians_model():
     n = 4

tests/test_find_map.py

@@ -1,12 +1,13 @@
+from typing import Literal
+
 import numpy as np
 import pymc as pm
 import pytensor.tensor as pt
 import pytest
 
+from pymc_extras.find_map import scipy_optimize_funcs_from_loss
 from pymc_extras.inference.find_map import (
     find_MAP,
-    find_MAP_scipy,
-    find_MAP_pytensor,
 )
 
 pytest.importorskip("jax")
@@ -18,6 +19,10 @@ def rng():
     return np.random.default_rng(seed)
 
 
+# Define GradientBackend type alias
+GradientBackend = Literal["jax", "pytensor"]
+
+
 @pytest.mark.parametrize("gradient_backend", ["jax", "pytensor"], ids=str)
 def test_jax_functions_from_graph(gradient_backend: GradientBackend):
     x = pt.tensor("x", shape=(2,))

tests/test_laplace.py

@@ -143,9 +143,7 @@ def test_fit_laplace_coords(rng, transform_samples, mode):
             transform_samples=transform_samples,
         )
 
-        idata = sample_laplace(
-            mu=mu, H_inv=H_inv, model=model, transform_samples=transform_samples
-        )
+        idata = sample_laplace(mu=mu, H_inv=H_inv, model=model, transform_samples=transform_samples)
 
     np.testing.assert_allclose(np.mean(idata.posterior.mu, axis=1), np.full((2, 3), 3), atol=0.5)
     np.testing.assert_allclose(