Additional fixes requested by ruff

Author: jessegrabowski

pymc_experimental/__init__.py

@@ -13,6 +13,10 @@
 #   limitations under the License.
 import logging
 
+from pymc_experimental import distributions, gp, statespace, utils
+from pymc_experimental.inference.fit import fit
+from pymc_experimental.model.marginal_model import MarginalModel
+from pymc_experimental.model.model_api import as_model
 from pymc_experimental.version import __version__
 
 _log = logging.getLogger("pmx")
@@ -23,7 +27,14 @@
         handler = logging.StreamHandler()
         _log.addHandler(handler)
 
-from pymc_experimental import distributions, gp, statespace, utils
-from pymc_experimental.inference.fit import fit
-from pymc_experimental.model.marginal_model import MarginalModel
-from pymc_experimental.model.model_api import as_model
+
+__all__ = [
+    "distributions",
+    "gp",
+    "statespace",
+    "utils",
+    "fit",
+    "MarginalModel",
+    "as_model",
+    "__version__",
+]

pymc_experimental/distributions/multivariate/__init__.py

@@ -1 +1,3 @@
 from pymc_experimental.distributions.multivariate.r2d2m2cp import R2D2M2CP
+
+__all__ = ["R2D2M2CP"]

pymc_experimental/gp/__init__.py

@@ -14,3 +14,5 @@
 
 
 from pymc_experimental.gp.latent_approx import KarhunenLoeveExpansion, ProjectedProcess
+
+__all__ = ["KarhunenLoeveExpansion", "ProjectedProcess"]

pymc_experimental/inference/__init__.py

@@ -14,3 +14,5 @@
 
 
 from pymc_experimental.inference.fit import fit
+
+__all__ = ["fit"]

pymc_experimental/inference/fit.py

@@ -11,6 +11,7 @@
 #   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.
+from importlib.util import find_spec
 
 
 def fit(method, **kwargs):
@@ -30,10 +31,8 @@ def fit(method, **kwargs):
     arviz.InferenceData
     """
     if method == "pathfinder":
-        try:
-            import blackjax
-        except ImportError as exc:
-            raise RuntimeError("Need BlackJAX to use `pathfinder`") from exc
+        if find_spec("blackjax") is None:
+            raise RuntimeError("Need BlackJAX to use `pathfinder`")
 
         from pymc_experimental.inference.pathfinder import fit_pathfinder
 

pymc_experimental/model/marginal_model.py

@@ -21,6 +21,7 @@
 from pytensor import Mode, scan
 from pytensor.compile import SharedVariable
 from pytensor.graph import Constant, FunctionGraph, ancestors, clone_replace
+from pytensor.graph.basic import graph_inputs
 from pytensor.graph.replace import graph_replace, vectorize_graph
 from pytensor.scan import map as scan_map
 from pytensor.tensor import TensorType, TensorVariable
@@ -638,9 +639,6 @@ def is_elemwise_subgraph(rv_to_marginalize, other_input_rvs, output_rvs):
     return True
 
 
-from pytensor.graph.basic import graph_inputs
-
-
 def collect_shared_vars(outputs, blockers):
     return [
         inp for inp in graph_inputs(outputs, blockers=blockers) if isinstance(inp, SharedVariable)

pymc_experimental/model/transforms/autoreparam.py

@@ -419,6 +419,9 @@ def vip_reparametrize(
         lambda_names.append(lam.name)
     toposort_replace(fmodel, replacements, reverse=True)
     reparam_model = model_from_fgraph(fmodel)
-    model_lambdas = {n: reparam_model[l] for l, n in zip(lambda_names, var_names)}
+    model_lambdas = {
+        var_name: reparam_model[lambda_name]
+        for lambda_name, var_name in zip(lambda_names, var_names)
+    }
     vip = VIP(model_lambdas)
     return reparam_model, vip

pymc_experimental/statespace/core/representation.py

@@ -1,7 +1,5 @@
 import copy
 
-from typing import Union
-
 import numpy as np
 import pytensor
 import pytensor.tensor as pt
@@ -12,7 +10,7 @@
 )
 
 floatX = pytensor.config.floatX
-KeyLike = Union[tuple[str | int, ...], str]
+KeyLike = tuple[str | int, ...] | str
 
 
 class PytensorRepresentation:

pymc_experimental/statespace/filters/distributions.py

@@ -96,8 +96,6 @@ def update(self, node: Node):
 
 
 class _LinearGaussianStateSpace(Continuous):
-    rv_op = LinearGaussianStateSpaceRV
-
     def __new__(
         cls,
         name,
@@ -360,8 +358,6 @@ def update(self, node: Node):
 
 
 class SequenceMvNormal(Continuous):
-    rv_op = KalmanFilterRV
-
     def __new__(cls, *args, **kwargs):
         return super().__new__(cls, *args, **kwargs)
 

pymc_experimental/statespace/filters/kalman_filter.py

@@ -351,11 +351,14 @@ def handle_missing_values(
         self, y, Z, H
     ) -> tuple[TensorVariable, TensorVariable, TensorVariable, float]:
         """
-        This function handles missing values in the observation data `y` and adjusts the design matrix `Z` and the
-        observation noise covariance matrix `H` accordingly. Missing values are replaced with zeros to prevent
-        propagating NaNs through the computation. The function also returns a binary flag tensor `all_nan_flag`,
-        indicating if all values in the observation data are missing. This flag is used for numerical adjustments in
-        the update method.
+        Handle missing values in the observation data `y`
+
+        Adjusts the design matrix `Z` and the observation noise covariance matrix `H` by removing rows and/or columns
+        associated with the data that is not observed at this iteration. Missing values are replaced with zeros to prevent
+        propagating NaNs through the computation.
+
+        Return a binary flag tensor `all_nan_flag`,indicating if all values in the observation data are missing. This
+        flag is used for numerical adjustments in the update method.
 
         Parameters
         ----------
@@ -660,7 +663,7 @@ def update(self, a, P, y, c, d, Z, H, all_nan_flag):
 
 
 class CholeskyFilter(BaseFilter):
-    """ "
+    """
     Kalman filter with Cholesky factorization
 
     Kalman filter implementation using a Cholesky factorization plus pt.solve_triangular to (attempt) to speed up
@@ -712,7 +715,7 @@ class SingleTimeseriesFilter(BaseFilter):
 
     # TODO: This class should eventually be made irrelevant by pytensor re-writes.
     def check_params(self, data, a0, P0, c, d, T, Z, R, H, Q):
-        """ "
+        """
         Wrap the data in an `Assert` `Op` to ensure there is only one observed state.
         """
         data = assert_data_is_1d(data, pt.eq(data.shape[1], 1))