Remove array and potrf copies

Author: Fyrebright

pytensor/tensor/slinalg.py

@@ -68,35 +68,38 @@ def perform(self, node, inputs, outputs):
         [x] = inputs
         [out] = outputs
 
+        (potrf,) = scipy_linalg.get_lapack_funcs(("potrf",), (x,))
+
         # Quick return for square empty array
         if x.size == 0:
-            eye = np.eye(1, dtype=x.dtype)
-            (potrf,) = scipy_linalg.get_lapack_funcs(("potrf",), (eye,))
-            c, _ = potrf(eye, lower=False, overwrite_a=False, clean=True)
-            out[0] = np.empty_like(x, dtype=c.dtype)
+            out[0] = np.empty_like(x, dtype=potrf.dtype)
             return
 
-        x1 = np.asarray_chkfinite(x) if self.check_finite else x
+        if self.check_finite and not np.isfinite(x).all():
+            if self.on_error == "nan":
+                out[0] = np.full(x.shape, np.nan, dtype=node.outputs[0].type.dtype)
+                return
+            else:
+                raise ValueError("array must not contain infs or NaNs")
 
         # Squareness check
-        if x1.shape[0] != x1.shape[1]:
+        if x.shape[0] != x.shape[1]:
             raise ValueError(
-                "Input array is expected to be square but has "
-                f"the shape: {x1.shape}."
+                "Input array is expected to be square but has " f"the shape: {x.shape}."
             )
 
         # Scipy cholesky only makes use of overwrite_a when it is F_CONTIGUOUS
         # If we have a `C_CONTIGUOUS` array we transpose to benefit from it
-        if self.overwrite_a and x.flags["C_CONTIGUOUS"]:
-            x1 = x1.T
+        c_contiguous_input = self.overwrite_a and x.flags["C_CONTIGUOUS"]
+        if c_contiguous_input:
+            x = x.T
             lower = not self.lower
             overwrite_a = True
         else:
             lower = self.lower
             overwrite_a = self.overwrite_a
 
-        (potrf,) = scipy_linalg.get_lapack_funcs(("potrf",), (x1,))
-        c, info = potrf(x1, lower=lower, overwrite_a=overwrite_a, clean=True)
+        c, info = potrf(x, lower=lower, overwrite_a=overwrite_a, clean=True)
 
         if info != 0:
             if self.on_error == "nan":
@@ -112,7 +115,7 @@ def perform(self, node, inputs, outputs):
                 )
         else:
             # Transpose result if input was transposed
-            out[0] = c.T if (self.overwrite_a and x.flags["C_CONTIGUOUS"]) else c
+            out[0] = c.T if c_contiguous_input else c
 
     def L_op(self, inputs, outputs, gradients):
         """