Add rewrite to remove useless Blockwise AdvancedIncSubtensor

Author: ricardoV94

pytensor/tensor/rewriting/subtensor.py

@@ -29,6 +29,7 @@
     register_infer_shape,
     switch,
 )
+from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.math import Dot, add
@@ -1865,3 +1866,37 @@ def local_uint_constant_indices(fgraph, node):
     copy_stack_trace(node.outputs, new_outs)
 
     return new_outs
+
+
+@register_canonicalize
+@register_stabilize
+@register_specialize
+@node_rewriter([Blockwise])
+def useless_blockwise_advanced_inc_subtensor(fgraph, node):
+    """Rewrite blockwise advanced inc_subtensor whithout batched indexes as an inc_subtensor with prepended empty slices."""
+    if not isinstance(node.op.core_op, AdvancedIncSubtensor):
+        return None
+
+    op: Blockwise = node.op
+
+    x, y, *idxs = node.inputs
+
+    batch_ndim = node.outputs[0].type.ndim - len(op.outputs_sig[0])
+
+    new_idxs = []
+    for idx in idxs:
+        if all(idx.type.broadcastable[:batch_ndim]):
+            new_idxs.append(idx.squeeze(tuple(range(batch_ndim))))
+        else:
+            # Rewrite does not apply
+            return None
+
+    # Need to broadcast x to final shape
+    if any(x.type.broadcastable[:batch_ndim]):
+        x = alloc(x, *tuple(y.shape)[:batch_ndim], *tuple(x.shape)[batch_ndim:])
+
+    new_idxs = [slice(None)] * batch_ndim + new_idxs
+    symbolic_idxs = x[tuple(new_idxs)].owner.inputs[1:]
+    new_out = op.core_op.make_node(x, y, *symbolic_idxs).outputs
+    copy_stack_trace(node.outputs, new_out)
+    return new_out