Linter fixes.

Author: twiecki

pymc/distributions/multivariate.py

@@ -69,7 +69,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.math import kron_diag, kron_dot
 from pymc.pytensorf import normalize_rng_param

pymc/distributions/transforms.py

@@ -17,7 +17,6 @@
 
 import numpy as np
 import pytensor.tensor as pt
-import pytensor
 
 
 # ignore mypy error because it somehow considers that
@@ -213,13 +212,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]
@@ -240,17 +236,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]
@@ -261,7 +254,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
@@ -684,7 +684,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"]
@@ -763,7 +765,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():
@@ -773,7 +775,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"]
@@ -808,4 +812,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)