rewrite deals with pt.diag as well now

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -2,6 +2,7 @@
 from collections.abc import Callable
 from typing import cast
 
+import pytensor.tensor as pt
 from pytensor import Variable
 from pytensor.graph import Apply, FunctionGraph
 from pytensor.graph.rewriting.basic import (
@@ -605,7 +606,7 @@ def rewrite_inv_diag_to_diag_reciprocal(fgraph, node):
     ):
         inv_input = inputs.owner.inputs[0]
         if inv_input.type.ndim == 1:
-            inv_val = diagonal(1 / inv_input)
+            inv_val = pt.diag(1 / inv_input)
             return [inv_val]
 
     # Check if the input is an elemwise multiply with identity matrix -- this also results in a diagonal matrix
@@ -619,21 +620,14 @@ def rewrite_inv_diag_to_diag_reciprocal(fgraph, node):
     if len(non_eye_inputs) != 1:
         return None
 
-    eye_input, non_eye_input = eye_input[0], non_eye_inputs[0]
+    non_eye_input = non_eye_inputs[0]
 
     # 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 [eye_input / non_eye_input.diagonal(axis1=-1, axis2=-2)]
-    elif non_eye_input.type.broadcastable[-2:] == (True, True):
-        # For Scalar
-        inv = eye_input / non_eye_input
-        print(eye_input.type)
-        print(non_eye_input.type)
-        print(inv.type)
-        return [eye_input / non_eye_input]
     else:
-        # For Vector
+        # For Vector or Scalar
         return [eye_input / non_eye_input]
 
 

tests/tensor/rewriting/test_linalg.py

@@ -43,9 +43,6 @@
 ATOL = RTOL = 1e-3 if config.floatX == "float32" else 1e-8
 
 
-ATOL = RTOL = 1e-3 if config.floatX == "float32" else 1e-8
-
-
 def test_rop_lop():
     mx = matrix("mx")
     mv = matrix("mv")
@@ -616,7 +613,7 @@ def test_inv_diag_from_diag():
     x_test = np.random.rand(10)
     x_test_matrix = np.eye(10) * x_test
     inverse_matrix = np.linalg.inv(x_test_matrix)
-    rewritten_inverse = f_rewritten(x_test_matrix)
+    rewritten_inverse = f_rewritten(x_test)
 
     assert_allclose(
         inverse_matrix,