Fix pivot_to_permutation

Author: jessegrabowski

pytensor/tensor/slinalg.py

@@ -585,41 +585,30 @@ def lu(
 
 
 class PivotToPermutations(Op):
-    __props__ = ("inverse", "inplace")
+    __props__ = ("inverse",)
 
-    def __init__(self, inverse=True, inplace=False):
+    def __init__(self, inverse=True):
         self.inverse = inverse
-        self.inplace = inplace
-        self.destroy_map = {}
-        if self.inplace:
-            self.destroy_map = {0: [0]}
 
     def make_node(self, pivots):
         pivots = as_tensor_variable(pivots)
         if pivots.ndim != 1:
             raise ValueError("PivotToPermutations only works on 1-D inputs")
-        permutations = pivots.type()
 
+        permutations = pivots.type.clone(dtype="int64")()
         return Apply(self, [pivots], [permutations])
 
-    def inplace_on_inputs(self, allowed_inplace_inputs: list[int]) -> "Op":
-        if 0 in allowed_inplace_inputs:
-            new_props = self._props_dict()  # type: ignore
-            new_props["inplace"] = True
-            return type(self)(**new_props)
-        else:
-            return self
-
     def perform(self, node, inputs, outputs):
-        [p] = inputs
-        p_inv = np.arange(len(p)).astype(p.dtype)
-        for i in range(len(p)):
-            p_inv[i], p_inv[p[i]] = p_inv[p[i]], p_inv[i]
+        [pivots] = inputs
+        p_inv = np.arange(len(pivots), dtype=pivots.dtype)
+
+        for i in range(len(pivots)):
+            p_inv[i], p_inv[pivots[i]] = p_inv[pivots[i]], p_inv[i]
 
         if self.inverse:
             outputs[0][0] = p_inv
-
-        outputs[0][0] = np.argsort(p_inv)
+        else:
+            outputs[0][0] = np.argsort(p_inv)
 
 
 def pivot_to_permutation(p: TensorLike, inverse=False) -> Variable:
@@ -629,14 +618,14 @@ def pivot_to_permutation(p: TensorLike, inverse=False) -> Variable:
 
 class LUFactor(Op):
     __props__ = ("overwrite_a", "check_finite", "permutation_indices")
+    gufunc_signature = "(m,m)->(m,m),(m)"
 
     def __init__(
         self, *, overwrite_a=False, check_finite=True, permutation_indices=False
     ):
         self.overwrite_a = overwrite_a
         self.check_finite = check_finite
         self.permutation_indices = permutation_indices
-        self.gufunc_signature = "(m,m)->(m,m),(m)"
 
         if self.overwrite_a:
             self.destroy_map = {1: [0]}

tests/tensor/test_slinalg.py

@@ -18,7 +18,6 @@
     Solve,
     SolveBase,
     SolveTriangular,
-    _pivot_to_permutation,
     block_diag,
     cho_solve,
     cholesky,
@@ -27,6 +26,7 @@
     lu,
     lu_factor,
     lu_solve,
+    pivot_to_permutation,
     solve,
     solve_continuous_lyapunov,
     solve_discrete_are,
@@ -667,14 +667,19 @@ def f_pt(A):
     utt.verify_grad(f_pt, [A_value], rng=rng)
 
 
-def test_pivot_to_permutation():
+@pytest.mark.parametrize("inverse", [True, False], ids=["inverse", "no_inverse"])
+def test_pivot_to_permutation(inverse):
     rng = np.random.default_rng(utt.fetch_seed())
     A_val = rng.normal(size=(5, 5))
     _, pivots = scipy.linalg.lu_factor(A_val)
     perm_idx, *_ = scipy.linalg.lu(A_val, p_indices=True)
 
-    permutations = pt.argsort(_pivot_to_permutation(pivots)).eval()
-    np.testing.assert_array_equal(permutations, perm_idx)
+    if not inverse:
+        perm_idx_pt = pivot_to_permutation(pivots, inverse=False).eval()
+        np.testing.assert_array_equal(perm_idx_pt, perm_idx)
+    else:
+        p_inv_pt = pivot_to_permutation(pivots, inverse=True).eval()
+        np.testing.assert_array_equal(p_inv_pt, np.argsort(perm_idx))
 
 
 class TestLUSolve(utt.InferShapeTester):
@@ -686,8 +691,8 @@ def factor_and_solve(A, b, sum=False, **lu_kwargs):
             return x
         return x.sum()
 
-    @pytest.mark.parametrize("b_shape", [(5,), (5, 5)])
-    @pytest.mark.parametrize("trans", [True, False])
+    @pytest.mark.parametrize("b_shape", [(5,), (5, 5)], ids=["b_vec", "b_matrix"])
+    @pytest.mark.parametrize("trans", [True, False], ids=["x_T", "x"])
     def test_lu_solve(self, b_shape: tuple[int], trans):
         rng = np.random.default_rng(utt.fetch_seed())
         A = pt.tensor("A", shape=(5, 5))