Implement specialized PivotToPermutation

Author: ricardoV94

pytensor/link/numba/dispatch/slinalg.py

@@ -18,12 +18,14 @@
     int_ptr_to_val,
     val_to_int_ptr,
 )
-from pytensor.link.numba.dispatch.basic import numba_funcify
+from pytensor.link.numba.dispatch.basic import numba_funcify, numba_njit
 from pytensor.tensor.slinalg import (
     LU,
     BlockDiagonal,
     Cholesky,
     CholeskySolve,
+    LUFactor,
+    PivotToPermutations,
     Solve,
     SolveTriangular,
 )
@@ -36,6 +38,18 @@
 )
 
 
+@numba_funcify.register(PivotToPermutations)
+def pivot_to_permutation(op, node, **kwargs):
+    @numba_basic.numba_njit
+    def pivot_to_permutation(p):
+        p_inv = np.arange(len(p))
+        for i in range(len(p)):
+            p_inv[i], p_inv[p[i]] = p_inv[p[i]], p_inv[i]
+        return p_inv
+
+    return pivot_to_permutation
+
+
 @numba_basic.numba_njit(inline="always")
 def _copy_to_fortran_order_even_if_1d(x):
     # Numba's _copy_to_fortran_order doesn't do anything for vectors
@@ -818,6 +832,20 @@ def lu(a):
     return lu
 
 
+@numba_funcify.register(LUFactor)
+def numba_funcify_LUFactor(op, node, **kwargs):
+    overwrite_a = op.overwrite_a
+
+    @numba_njit
+    def lu_factor(A):
+        N = A.shape[1]
+        LU, IPIV, INFO = _getrf(A, overwrite_a=overwrite_a)
+        _solve_check(N, INFO)
+        return LU, IPIV - 1
+
+    return lu_factor
+
+
 def _getrs(
     LU: np.ndarray, B: np.ndarray, IPIV: np.ndarray, trans: int, overwrite_b: bool
 ) -> tuple[np.ndarray, int]:

pytensor/tensor/slinalg.py

@@ -21,7 +21,7 @@
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.nlinalg import kron, matrix_dot
 from pytensor.tensor.shape import reshape
-from pytensor.tensor.type import matrix, tensor, vector
+from pytensor.tensor.type import ivector, matrix, tensor, vector
 from pytensor.tensor.variable import TensorVariable
 
 
@@ -583,31 +583,21 @@ def lu(
     )
 
 
-def _pivot_to_permutation(pivots):
-    """
-    Converts a sequence of row exchanges to a permutation matrix that represents the same row exchanges. This
-    represents the inverse permutation, which can be used to reconstruct the original matrix from its LU factorization.
-    To get the actual permutation, the inverse permutation must be argsorted.
-    """
+class PivotToPermutations(Op):
+    itypes = [ivector]
+    otypes = [ivector]
+
+    __props__ = ()
+
+    def perform(self, node, inputs, outputs):
+        [p] = inputs
+        p_inv = np.arange(len(p))
+        for i in range(len(p)):
+            p_inv[i], p_inv[p[i]] = p_inv[p[i]], p_inv[i]
+        outputs[0][0] = p_inv
 
-    def step(i, permutation, swaps):
-        j = swaps[i]
-        x = permutation[i]
-        y = permutation[j]
-
-        permutation = permutation[i].set(y)
-        return permutation[j].set(x)
-
-    pivots = as_tensor_variable(pivots)
-    n = pivots.shape[0]
-    p_inv, _ = pytensor.scan(
-        step,
-        sequences=[pt.arange(n.copy())],
-        outputs_info=[pt.arange(n.copy())],
-        non_sequences=[pivots],
-    )
 
-    return p_inv[-1]
+_pivot_to_permutation = PivotToPermutations()
 
 
 class LUFactor(Op):
@@ -810,13 +800,7 @@ def lu_solve(
     )
     x = x[pt.argsort(inv_permutation)] if trans else x
 
-    return LUSolve(
-        inputs=[LU, pivots, b],
-        outputs=[x],
-        trans=trans,
-        b_ndim=b_ndim,
-        check_finite=check_finite,
-    )(LU, pivots, b)
+    return x
 
 
 class SolveTriangular(SolveBase):