Split: Return disconnected gradient for split sizes

Author: ricardoV94

pytensor/tensor/basic.py

@@ -2254,18 +2254,19 @@ def infer_shape(self, fgraph, node, in_shapes):
             out_shapes.append(temp)
         return out_shapes
 
+    def connection_pattern(self, node):
+        n_out = len(node.outputs)
+        return [
+            [True] * n_out,
+            [True] * n_out,
+            [False] * n_out,
+        ]
+
     def L_op(self, inputs, outputs, g_outputs):
         """Join the gradients along the axis that was used to split x."""
-        _x, axis, n = inputs
+        _x, axis, _n = inputs
 
-        # If all the output gradients are disconnected, then so are the inputs
-        if builtins.all(isinstance(g.type, DisconnectedType) for g in g_outputs):
-            return [
-                DisconnectedType()(),
-                grad_undefined(self, 1, axis),
-                grad_undefined(self, 2, n),
-            ]
-        # Else, we have to make them zeros before joining them
+        # We have to convert disconnected outputs to zeros before joining them
         new_g_outputs = []
         for o, g in zip(outputs, g_outputs, strict=True):
             if isinstance(g.type, DisconnectedType):
@@ -2276,7 +2277,7 @@ def L_op(self, inputs, outputs, g_outputs):
         return [
             join(axis, *new_g_outputs),
             grad_undefined(self, 1, axis),
-            grad_undefined(self, 2, n),
+            DisconnectedType()(),
         ]
 
     def R_op(self, inputs, eval_points):

tests/tensor/test_reshape.py

@@ -6,6 +6,7 @@
 from pytensor import config, function
 from pytensor import tensor as pt
 from pytensor.graph import rewrite_graph, vectorize_graph
+from pytensor.graph.op import io_connection_pattern
 from pytensor.tensor.reshape import (
     _analyze_axes_list,
     join_dims,
@@ -289,3 +290,12 @@ def test_pack_unpack_round_trip(self, axes):
 
         for input_val, output_val in zip(input_dict.values(), output_vals, strict=True):
             np.testing.assert_allclose(input_val, output_val)
+
+
+def test_unpack_connection():
+    x = pt.vector("x")
+    d0 = pt.scalar("d0", dtype=int)
+    d1 = pt.scalar("d1", dtype=int)
+    x0, x1 = pt.unpack(x, axes=None, packed_shapes=[d0, d1])
+    out = x0.sum() + x1.sum()
+    assert io_connection_pattern([x, d0, d1], [out]) == [[True], [False], [False]]