Add xtensor hashing and restore is_full_slice for pymc-extras compatibility

Author: jaanerik

pytensor/tensor/rewriting/subtensor.py

@@ -162,6 +162,11 @@ def transform_take(a, indices, axis):
     return transform_take(a, indices.flatten(), axis).reshape(shape, ndim=ndim)
 
 
+def is_full_slice(x):
+    # Replace this function in pymc-extras and pymc with x==slice(None)
+    return x == slice(None)
+
+
 def get_advsubtensor_axis(indices):
     """Determine the axis at which an array index is applied.
 

pytensor/tensor/subtensor.py

@@ -2870,7 +2870,6 @@ def get_slice_val(comp):
                 start_val = get_slice_val(entry.start)
                 stop_val = get_slice_val(entry.stop)
                 step_val = get_slice_val(entry.step)
-
                 full_indices.append(slice(start_val, stop_val, step_val))
             else:
                 assert isinstance(entry, int)

pytensor/xtensor/indexing.py

@@ -326,6 +326,19 @@ def __init__(self, mode: Literal["set", "inc"], idx_list):
         self.mode = mode
         self.idx_list = idx_list
 
+    def __hash__(self):
+        """Hash using mode and idx_list. Slices are not hashable in Python < 3.12."""
+        return hash((type(self), self.mode, self._hashable_idx_list()))
+
+    def _hashable_idx_list(self):
+        """Return a hashable version of idx_list (slices converted to tuples)."""
+        return tuple(
+            (slice, entry.start, entry.stop, entry.step)
+            if isinstance(entry, slice)
+            else entry
+            for entry in self.idx_list
+        )
+
     def make_node(self, x, y, x_view, *index_inputs):
         try:
             y = as_xtensor(y)