minor changes

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -591,7 +591,6 @@ def rewrite_inv_for_diag_eye_mul(fgraph, node):
     list of Variable, optional
         List of optimized variables, or None if no optimization was performed
     """
-    # List of useful operations : Inv, Pinv
     valid_inverses = (MatrixInverse, MatrixPinv)
     core_op = node.op.core_op
     if not (isinstance(core_op, valid_inverses)):
@@ -609,15 +608,15 @@ def rewrite_inv_for_diag_eye_mul(fgraph, node):
     if len(non_eye_inputs) != 1:
         return None
 
-    useful_eye, useful_non_eye = eye_input[0], non_eye_inputs[0]
+    eye_input, non_eye_input = eye_input[0], non_eye_inputs[0]
 
-    # For a matrix, we can first get the diagonal and then only use those
-    if useful_non_eye.type.broadcastable[-2:] == (False, False):
+    # For a matrix, we have to first extract the diagonal (non-zero values) and then only use those
+    if non_eye_input.type.broadcastable[-2:] == (False, False):
         # For Matrix
-        return [useful_eye * 1 / useful_non_eye.diagonal(axis1=-1, axis2=-2)]
+        return [eye_input / non_eye_input.diagonal(axis1=-1, axis2=-2)]
     else:
         # For Scalar/Vector
-        return [useful_eye * 1 / useful_non_eye]
+        return [eye_input / non_eye_input]
 
 
 def rewrite_inv_for_diag_ptdiag(fgraph, node):
@@ -656,7 +655,7 @@ def rewrite_inv_for_orthonormal(fgraph, node):
     ):
         return None
 
-    # To make sure input is orthonormal, we have to check that its not S (output order is U, S, Vh, so S is index 1)
+    # To make sure input is orthonormal, we have to check that its not S (output order of SVD is U, S, Vh, so S is index 1) (S matrix consists of singular values and it is not orthonormal)
     if input_to_inv == input_to_inv.owner.outputs[1]:
         return None
 

tests/tensor/rewriting/test_linalg.py

@@ -40,6 +40,9 @@
 from tests.test_rop import break_op
 
 
+ATOL = RTOL = 1e-3 if config.floatX == "float32" else 1e-8
+
+
 def test_rop_lop():
     mx = matrix("mx")
     mv = matrix("mv")
@@ -589,8 +592,8 @@ def test_inv_diag_from_eye_mul(shape):
     assert_allclose(
         inverse_matrix,
         rewritten_inverse,
-        atol=1e-3 if config.floatX == "float32" else 1e-8,
-        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+        atol=ATOL,
+        rtol=RTOL,
     )
 
 
@@ -599,7 +602,7 @@ def test_inv_orthonormal():
     u, s, vh = pt.linalg.svd(x)
     # Calculating inverse using pt.linalg.inv
     u_inv = pt.linalg.inv(u)
-    print(u_inv.dprint())
+
     # REWRITE TEST
     f_rewritten = function([x], u_inv, mode="FAST_RUN")
     nodes = f_rewritten.maker.fgraph.apply_nodes
@@ -616,6 +619,6 @@ def test_inv_orthonormal():
     assert_allclose(
         inverse_matrix,
         rewritten_inverse,
-        atol=1e-3 if config.floatX == "float32" else 1e-8,
-        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+        atol=ATOL,
+        rtol=RTOL,
     )