add check for k = 0

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -552,9 +552,12 @@ def _find_solve_with_eye(node):
     # Check whether second input to solve is Eye
     solve_inputs = node.inputs
     potential_eye_input = solve_inputs[1].owner
+
     if not (potential_eye_input and isinstance(potential_eye_input.op, Eye)):
         return False
 
+    if getattr(potential_eye_input.inputs[-1], "data", -1).item() != 0:
+        return False
     return True
 
 

tests/tensor/rewriting/test_linalg.py

@@ -564,5 +564,5 @@ def get_pt_function(x, op_name):
     valid_inverses = (MatrixInverse, MatrixPinv, Solve, SolveTriangular)
 
     assert not any(isinstance(node.op, valid_inverses) for node in nodes)
-    x_testing = np.random.rand(10, 10)
+    x_testing = np.random.rand(10, 10).astype(config.floatX)
     np.testing.assert_allclose(f_rewritten(x_testing), x_testing)