Lint fixes

Author: zaxtax

pymc_experimental/model/marginal_model.py

@@ -1,10 +1,10 @@
 import warnings
 from typing import Sequence, Tuple, Union
 
-from arviz import dict_to_dataset
 import numpy as np
-import pytensor.tensor as pt
 import pymc
+import pytensor.tensor as pt
+from arviz import dict_to_dataset
 from pymc import SymbolicRandomVariable
 from pymc.backends.arviz import coords_and_dims_for_inferencedata
 from pymc.distributions.discrete import Bernoulli, Categorical, DiscreteUniform
@@ -13,7 +13,7 @@
 from pymc.logprob.basic import conditional_logp
 from pymc.logprob.transforms import IntervalTransform
 from pymc.model import Model
-from pymc.pytensorf import constant_fold, compile_pymc, inputvars
+from pymc.pytensorf import compile_pymc, constant_fold, inputvars
 from pymc.util import dataset_to_point_list, treedict
 from pytensor import Mode
 from pytensor.compile import SharedVariable
@@ -248,7 +248,9 @@ def marginalize(self, rvs_to_marginalize: Union[TensorVariable, Sequence[TensorV
         # Raise errors and warnings immediately
         self.clone()._marginalize(user_warnings=True)
 
-    def unmunmarginalize(self, idata, var_names=None, include_samples=False, extend_inferencedata=True):
+    def unmunmarginalize(
+        self, idata, var_names=None, include_samples=False, extend_inferencedata=True
+    ):
         """Computes log-likelihoods of marginalized variables conditioned on parameters
         of the model given InferenceData with posterior group
 
@@ -276,7 +278,9 @@ def unmunmarginalize(self, idata, var_names=None, include_samples=False, extend_
         posterior = idata.posterior
 
         # Remove Deterministics
-        posterior_values = posterior[[rv.name for rv in mm.free_RVs if rv not in self.marginalized_rvs]]
+        posterior_values = posterior[
+            [rv.name for rv in mm.free_RVs if rv not in self.marginalized_rvs]
+        ]
 
         sample_dims = ("chain", "draw")
         posterior_pts, stacked_dims = dataset_to_point_list(posterior_values, sample_dims)
@@ -306,10 +310,11 @@ def unmunmarginalize(self, idata, var_names=None, include_samples=False, extend_
 
             # TODO: Handle constants
             # TODO: Handle transformed variables
-            joint_logp_op = OpFromGraph([marginalized_value] + other_values, [joint_logp], inline=True)
+            joint_logp_op = OpFromGraph(
+                [marginalized_value] + other_values, [joint_logp], inline=True
+            )
             joint_logps = [
-                joint_logp_op(rv_domain_tensor[i], *other_values)
-                for i in range(len(rv_domain))
+                joint_logp_op(rv_domain_tensor[i], *other_values) for i in range(len(rv_domain))
             ]
 
             rv_loglike_fn = None
@@ -333,20 +338,17 @@ def unmunmarginalize(self, idata, var_names=None, include_samples=False, extend_
                 logps, samples = zip(*logvs)
                 logps = np.array(logps)
                 rv_dict[rv.name] = np.reshape(
-                    samples,
-                    tuple(len(coord) for coord in stacked_dims.values())
+                    samples, tuple(len(coord) for coord in stacked_dims.values())
                 )
                 rv_dims_dict[rv.name] = sample_dims
                 rv_dict["lp_" + rv.name] = np.reshape(
-                    logps,
-                    tuple(len(coord) for coord in stacked_dims.values()) + logps.shape[1:]
+                    logps, tuple(len(coord) for coord in stacked_dims.values()) + logps.shape[1:]
                 )
                 rv_dims_dict["lp_" + rv.name] = sample_dims + ("lp_" + rv.name + "_dims",)
             else:
                 logps = np.array(logvs)
                 rv_dict["lp_" + rv.name] = np.reshape(
-                    logps,
-                    tuple(len(coord) for coord in stacked_dims.values()) + logps.shape[1:]
+                    logps, tuple(len(coord) for coord in stacked_dims.values()) + logps.shape[1:]
                 )
                 rv_dims_dict["lp_" + rv.name] = sample_dims + ("lp_" + rv.name + "_dims",)
 
@@ -361,7 +363,7 @@ def unmunmarginalize(self, idata, var_names=None, include_samples=False, extend_
         )
 
         if extend_inferencedata:
-            rv_dict = {k:(rv_dims_dict[k], v) for (k,v) in rv_dict.items()}
+            rv_dict = {k: (rv_dims_dict[k], v) for (k, v) in rv_dict.items()}
             idata = idata.posterior.assign(**rv_dict)
             return idata
         else:

pymc_experimental/tests/model/test_marginal_model.py

@@ -52,7 +52,12 @@ def test_marginalized_bernoulli_logp():
 
     idx = pm.Bernoulli.dist(0.7, name="idx")
     y = pm.Normal.dist(mu=mu[idx], sigma=1.0, name="y")
-    marginal_rv_node = FiniteDiscreteMarginalRV([mu], [idx, y], ndim_supp=None, n_updates=0,)(
+    marginal_rv_node = FiniteDiscreteMarginalRV(
+        [mu],
+        [idx, y],
+        ndim_supp=None,
+        n_updates=0,
+    )(
         mu
     )[0].owner
 
@@ -250,6 +255,7 @@ def test_marginalized_change_point_model_sampling(disaster_model):
         rtol=1e-2,
     )
 
+
 @pytest.mark.slow
 @pytest.mark.filterwarnings("error")
 def test_unmarginalized_basic(disaster_model):
@@ -269,6 +275,7 @@ def test_unmarginalized_basic(disaster_model):
     assert idata.switchpoint.shape == idata.early_mean.shape
     assert idata.lp_switchpoint.shape == idata.switchpoint.shape + (len(years),)
 
+
 @pytest.mark.filterwarnings("error")
 def test_not_supported_marginalized():
     """Marginalized graphs with non-Elemwise Operations are not supported as they