Provide static shape in output of Split

Author: ricardoV94

pytensor/tensor/basic.py

@@ -2201,8 +2201,28 @@ def make_node(self, x, axis, splits):
             raise TypeError("`axis` parameter must be an integer scalar")
 
         inputs = [x, axis, splits]
-        out_type = TensorType(dtype=x.dtype, shape=(None,) * x.type.ndim)
-        outputs = [out_type() for i in range(self.len_splits)]
+
+        x_dtype = x.type.dtype
+        if isinstance(axis, Constant):
+            # In this case we can preserve more static shape info
+            static_axis = axis.data.item()
+            outputs = []
+            x_static_shape = list(x.type.shape)
+            for i in range(self.len_splits):
+                try:
+                    static_split_size = int(get_scalar_constant_value(splits[i]))
+                except NotScalarConstantError:
+                    static_split_size = None
+                except IndexError:
+                    raise ValueError("Number of splits is larger than splits size")
+                static_out_shape = x_static_shape.copy()
+                static_out_shape[static_axis] = static_split_size
+                outputs.append(tensor(shape=tuple(static_out_shape), dtype=x_dtype))
+        else:
+            outputs = [
+                tensor(shape=(None,) * x.type.ndim, dtype=x_dtype)
+                for i in range(self.len_splits)
+            ]
 
         return Apply(self, inputs, outputs)
 

tests/link/jax/test_tensor_basic.py

@@ -150,12 +150,14 @@ def test_runtime_errors(self):
         ):
             fn(np.zeros((6, 4), dtype=pytensor.config.floatX))
 
-        a_splits = ptb.split(a, splits_size=[2, 4], n_splits=3, axis=0)
-        fn = pytensor.function([a], a_splits, mode="JAX")
+        # This check is triggered at compile time if splits_size has incompatible static length
+        splits_size = vector("splits_size", shape=(None,), dtype=int)
+        a_splits = ptb.split(a, splits_size=splits_size, n_splits=3, axis=0)
+        fn = pytensor.function([a, splits_size], a_splits, mode="JAX")
         with pytest.raises(
             ValueError, match="Length of splits is not equal to n_splits"
         ):
-            fn(np.zeros((6, 4), dtype=pytensor.config.floatX))
+            fn(np.zeros((6, 4), dtype=pytensor.config.floatX), [2, 2])
 
         a_splits = ptb.split(a, splits_size=[2, 4], n_splits=2, axis=0)
         fn = pytensor.function([a], a_splits, mode="JAX")