Linter fixes.

Author: twiecki

pymc/distributions/multivariate.py

@@ -73,7 +73,12 @@
     rv_size_is_none,
     to_tuple,
 )
-from pymc.distributions.transforms import Interval, ZeroSumTransform, _default_transform
+from pymc.distributions.transforms import (
+    CholeskyCorr,
+    Interval,
+    ZeroSumTransform,
+    _default_transform,
+)
 from pymc.logprob.abstract import _logprob
 from pymc.logprob.rewriting import (
     specialization_ir_rewrites_db,

pymc/distributions/transforms.py

@@ -16,7 +16,6 @@
 
 import numpy as np
 import pytensor.tensor as pt
-import pytensor
 
 from pytensor.graph import Op
 from pytensor.npy_2_compat import normalize_axis_tuple
@@ -214,13 +213,10 @@ def forward(self, x, *inputs):
         chol = pt.zeros((self.n, self.n), dtype=x.dtype)
 
         # Assign the unconstrained values to the lower triangular part
-        chol = pt.set_subtensor(
-            chol[self.tril_r_idxs, self.tril_c_idxs],
-            x
-        )
+        chol = pt.set_subtensor(chol[self.tril_r_idxs, self.tril_c_idxs], x)
 
         # Normalize each row to have unit L2 norm
-        row_norms = pt.sqrt(pt.sum(chol ** 2, axis=1, keepdims=True))
+        row_norms = pt.sqrt(pt.sum(chol**2, axis=1, keepdims=True))
         chol = chol / row_norms
 
         return chol[self.tril_r_idxs, self.tril_c_idxs]
@@ -241,17 +237,14 @@ def backward(self, y, *inputs):
         """
         # Reconstruct the full Cholesky matrix
         chol = pt.zeros((self.n, self.n), dtype=y.dtype)
-        chol = pt.set_subtensor(
-            chol[self.triu_r_idxs, self.triu_c_idxs],
-            y
-        )
+        chol = pt.set_subtensor(chol[self.triu_r_idxs, self.triu_c_idxs], y)
         chol = chol + pt.transpose(chol) + pt.eye(self.n, dtype=y.dtype)
 
         # Perform Cholesky decomposition
         chol = pt.linalg.cholesky(chol)
 
         # Extract the unconstrained parameters by normalizing
-        row_norms = pt.sqrt(pt.sum(chol ** 2, axis=1))
+        row_norms = pt.sqrt(pt.sum(chol**2, axis=1))
         unconstrained = chol / row_norms[:, None]
 
         return unconstrained[self.tril_r_idxs, self.tril_c_idxs]
@@ -262,7 +255,7 @@ def log_jac_det(self, y, *inputs):
 
         The Jacobian determinant for normalization is the product of row norms.
         """
-        row_norms = pt.sqrt(pt.sum(y ** 2, axis=1))
+        row_norms = pt.sqrt(pt.sum(y**2, axis=1))
         return -pt.sum(pt.log(row_norms), axis=-1)
 
 

tests/distributions/test_transform.py

@@ -23,8 +23,8 @@
 
 import pymc as pm
 import pymc.distributions.transforms as tr
-from pymc.distributions.transforms import CholeskyCorr
 
+from pymc.distributions.transforms import CholeskyCorr
 from pymc.logprob.basic import transformed_conditional_logp
 from pymc.logprob.transforms import Transform
 from pymc.pytensorf import floatX, jacobian
@@ -676,7 +676,9 @@ def test_lkjcorr_transform_round_trip():
     with pm.Model() as model:
         rho = pm.LKJCorr("rho", n=3, eta=2)
 
-    trace = pm.sample(100, tune=100, chains=1, cores=1, progressbar=False, return_inferencedata=False)
+    trace = pm.sample(
+        100, tune=100, chains=1, cores=1, progressbar=False, return_inferencedata=False
+    )
 
     # Extract the sampled correlation matrices
     rho_samples = trace["rho"]
@@ -755,7 +757,7 @@ def test_lkjcorr_invalid_n():
 
     with pytest.raises(TypeError):
         # 'n' must be an integer
-        CholeskyCorr(n='three')
+        CholeskyCorr(n="three")
 
 
 def test_lkjcorr_positive_definite():
@@ -765,7 +767,9 @@ def test_lkjcorr_positive_definite():
     with pm.Model() as model:
         rho = pm.LKJCorr("rho", n=4, eta=2)
 
-    trace = pm.sample(100, tune=100, chains=1, cores=1, progressbar=False, return_inferencedata=False)
+    trace = pm.sample(
+        100, tune=100, chains=1, cores=1, progressbar=False, return_inferencedata=False
+    )
 
     # Extract the sampled correlation matrices
     rho_samples = trace["rho"]
@@ -800,4 +804,4 @@ def test_lkjcorr_round_trip_various_sizes():
         reconstructed = transform.backward(y).eval()
 
         # Assert that the original and reconstructed unconstrained parameters are close
-        assert_allclose(x, reconstructed, atol=1e-6)
+        assert_allclose(x, reconstructed, atol=1e-6)