Fix failing JAX Split test

Author: ricardoV94

pytensor/link/jax/dispatch/tensor_basic.py

@@ -119,20 +119,22 @@ def jax_funcify_Split(op: Split, node, **kwargs):
     def split(x, axis, splits):
         if constant_axis is not None:
             axis = constant_axis
+            if len(splits) != op.len_splits:
+                raise ValueError("Length of splits is not equal to n_splits")
+
         if constant_splits is not None:
             splits = constant_splits
             cumsum_splits = np.cumsum(splits[:-1])
+            if (splits < 0).any():
+                raise ValueError("Split sizes cannot be negative")
         else:
             cumsum_splits = jnp.cumsum(splits[:-1])
 
-        if len(splits) != op.len_splits:
-            raise ValueError("Length of splits is not equal to n_splits")
-        if np.sum(splits) != x.shape[axis]:
-            raise ValueError(
-                f"Split sizes do not sum up to input length along axis: {x.shape[axis]}"
-            )
-        if np.any(splits < 0):
-            raise ValueError("Split sizes cannot be negative")
+        if constant_axis is not None and constant_splits is not None:
+            if splits.sum() != x.shape[axis]:
+                raise ValueError(
+                    f"Split sizes do not sum up to input length along axis: {x.shape[axis]}"
+                )
 
         return jnp.split(x, cumsum_splits, axis=axis)
 

tests/link/jax/test_tensor_basic.py

@@ -182,8 +182,7 @@ def test_jax_split_not_supported(self):
             UserWarning, match="Split node does not have constant split positions."
         ):
             fn = pytensor.function([a], a_splits, mode="JAX")
-        # It raises an informative ConcretizationTypeError, but there's an AttributeError that surpasses it
-        with pytest.raises(AttributeError):
+        with pytest.raises(errors.ConcretizationTypeError):
             fn(np.zeros((6, 4), dtype=pytensor.config.floatX))
 
         split_axis = iscalar("split_axis")
@@ -192,7 +191,7 @@ def test_jax_split_not_supported(self):
             fn = pytensor.function([a, split_axis], a_splits, mode="JAX")
         # Same as above, an AttributeError surpasses the `TracerIntegerConversionError`
         # Both errors are included for backwards compatibility
-        with pytest.raises((AttributeError, errors.TracerIntegerConversionError)):
+        with pytest.raises((AttributeError, errors.ConcretizationTypeError)):
             fn(np.zeros((6, 6), dtype=pytensor.config.floatX), 0)