Add is_zero_offset helper to Eye

Author: jessegrabowski

pytensor/tensor/basic.py

@@ -1334,6 +1334,25 @@ def infer_shape(self, fgraph, node, in_shapes):
     def grad(self, inp, grads):
         return [grad_undefined(self, i, inp[i]) for i in range(3)]
 
+    @staticmethod
+    def is_offset_zero(node) -> bool:
+        """
+        Test if an Eye Op has a diagonal offset of zero
+
+        Parameters
+        ----------
+        node
+            Eye node to test
+
+        Returns
+        -------
+        is_offset_zero: bool
+            True if the offset is zero (``k = 0``).
+        """
+
+        offset = node.inputs[-1]
+        return isinstance(offset, Constant) and offset.data.item() == 0
+
 
 def eye(n, m=None, k=0, dtype=None):
     """Return a 2-D array with ones on the diagonal and zeros elsewhere.

pytensor/tensor/rewriting/linalg.py

@@ -3,7 +3,7 @@
 from typing import cast
 
 from pytensor import Variable
-from pytensor.graph import Apply, Constant, FunctionGraph
+from pytensor.graph import Apply, FunctionGraph
 from pytensor.graph.rewriting.basic import (
     copy_stack_trace,
     node_rewriter,
@@ -438,15 +438,15 @@ def _find_diag_from_eye_mul(potential_mul_input):
         # Check if 1's are being put on the main diagonal only (k = 0)
         # and if the identity matrix is degenerate (column or row matrix)
         if not (
-            isinstance(inner_eye.owner.inputs[-1], Constant)
-            and inner_eye.owner.inputs[-1].data == 0
+            Eye.is_offset_zero(inner_eye.owner)
             and inner_eye.broadcastable[-1:] != (False, False)
         ):
             return None
 
-    elif getattr(
-        eye_input.owner.inputs[-1], "data", -1
-    ).item() != 0 or eye_input.broadcastable[-2:] != (False, False):
+    elif not (
+        Eye.is_offset_zero(eye_input.owner)
+        and eye_input.broadcastable[-1:] != (False, False)
+    ):
         return None
 
     # Get all non Eye inputs (scalars/matrices/vectors)