Fix AdvancedSubtensor static shape with newaxis

Author: ricardoV94

pytensor/tensor/subtensor.py

@@ -2629,9 +2629,13 @@ def make_node(self, x, *indices):
         advanced_indices = []
         adv_group_axis = None
         last_adv_group_axis = None
-        expanded_x_shape = tuple(
-            np.insert(np.array(x.type.shape, dtype=object), 1, new_axes)
-        )
+        if new_axes:
+            expanded_x_shape_list = list(x.type.shape)
+            for new_axis in new_axes:
+                expanded_x_shape_list.insert(new_axis, 1)
+            expanded_x_shape = tuple(expanded_x_shape_list)
+        else:
+            expanded_x_shape = x.type.shape
         for i, (idx, dim_length) in enumerate(
             zip_longest(explicit_indices, expanded_x_shape, fillvalue=NoneSliceConst)
         ):

tests/tensor/test_subtensor.py

@@ -1856,6 +1856,7 @@ def test_static_shape(self):
 
         assert x[idx1].type.shape == (10, None)
         assert x[:, idx1].type.shape == (None, 10)
+        assert x[None, :, idx1].type.shape == (1, None, 10)
         assert x[idx2, :5].type.shape == (3, None, None)
         assert specify_shape(x, (None, 7))[idx2, :5].type.shape == (3, None, 5)
         assert specify_shape(x, (None, 3))[idx2, :5].type.shape == (3, None, 3)