Add is_left_expand_dims and is_right_expand_dims attributes to DimShuffle

Author: jessegrabowski

pytensor/tensor/elemwise.py

@@ -182,14 +182,12 @@ def __init__(self, input_broadcastable, new_order):
         self.drop = drop
 
         input_ndim = len(input_broadcastable)
-
-        self.is_left_expand_dims = "x" in new_order and (
-            input_ndim == 0 or new_order[:input_ndim] == list(range(input_ndim))
-        )
-
-        self.is_right_expand_dims = "x" in new_order and (
+        self.is_left_expand_dims = self.augment and (
             input_ndim == 0 or new_order[-input_ndim:] == list(range(input_ndim))
         )
+        self.is_right_expand_dims = self.augment and new_order[:input_ndim] == list(
+            range(input_ndim)
+        )
 
         if self.inplace:
             self.view_map = {0: [0]}