Preserve static shape after `AdvancedSubtensor` that only shuffles elements

[jessegrabowski]

Description

I'm thinking about graphs like this:

import pytensor.tensor as pt

x = pt.tensor('x', shape=(3,)) 
x[[2, 0, 1]].type.shape # (None, )

For each index input, it should be possible to check in the make_node of AdvancedSubtensor and check if:

1. x has known static shape at the dimension being indexed, and if so;
2. check if index is a constant, and if so;
3. check if index it is only permuting element (len(index) == x.type.shape and np.sort(index) == np.arange(index.shape[0]))

Then we know the shape of x will not change.

[ricardoV94]

The contents of integer indices don't matter for static shape just their length

[ricardoV94]

Complexity cones from when you have multiple advanced indices plus mixed basic.

Only for slices or boolean indices do you need to know the contents to know the static shape.

[jessegrabowski]

Ah good point. So we can infer the static shape in more situations.

Anyway I'm not asking for anything exhaustive here, just that we preserve/figure out the shape info in the easy cases.

[ricardoV94]

Sure it's just you need to know the complexity to know what the simple cases are because indices interact with each other

[jessegrabowski]

I'm even happy to restrict to single indexing as a first try