Should shape graphs assume inputs are valid?

[ricardoV94]

According to the documentation: infer_shape implementations do not generally assert the inputs are valid: https://pytensor.readthedocs.io/en/latest/tutorial/shape_info.html#problems-with-shape-inference

This was the case in the times of Theano, and they mention Elemwise specifically.
In Aesara, the Elemwise.infer_shape was eventually modified to validate shape information when we moved towards dynamic broadcasting (see brandonwillard/aesara@064e72f)

Although the two issues need not be related, the implementation logic changes depending which view we want to take.

This is how the optimized compiled graph for an addition shape looks like:

import pytensor
import pytensor.tensor as pt

x = pt.vector("x")
y = pt.vector("y")
out = (x + y).shape
fn = pytensor.function([x, y], out)
pytensor.dprint(fn)

MakeVector{dtype='int64'} [id A] 8
 |TensorFromScalar [id B] 7
   |Assert{msg=Could not broadcast dimensions} [id C] 6
     |Composite [id D] 5
     | |ScalarFromTensor [id E] 3
     | | |Shape_i{0} [id F] 2
     | |   |x [id G]
     | |ScalarFromTensor [id H] 1
     |   |Shape_i{0} [id I] 0
     |     |y [id J]
     |Composite [id K] 4
       |ScalarFromTensor [id E] 3
       |ScalarFromTensor [id H] 1

Inner graphs:

Composite [id D]
 >Abs [id L]
 > |maximum [id M]
 >   |Switch [id N]
 >   | |EQ [id O]
 >   | | |<int64> [id P]
 >   | | |ScalarConstant{1} [id Q]
 >   | |neg [id R]
 >   | | |ScalarConstant{1} [id S]
 >   | |<int64> [id P]
 >   |Switch [id T]
 >     |EQ [id U]
 >     | |<int64> [id V]
 >     | |ScalarConstant{1} [id Q]
 >     |neg [id R]
 >     |<int64> [id V]
Composite [id K]
 >AND [id W]
 > |OR [id X]
 > | |EQ [id Y]
 > | | |Switch [id Z]
 > | | | |EQ [id BA]
 > | | | | |<int64> [id BB]
 > | | | | |ScalarConstant{1} [id Q]
 > | | | |neg [id BC]
 > | | | | |ScalarConstant{1} [id S]
 > | | | |<int64> [id BB]
 > | | |neg [id BC]
 > | |EQ [id BD]
 > |   |Switch [id Z]
 > |   |Composite [id BE]
 > |     |<int64> [id BB]
 > |     |<int64> [id BF]
 > |OR [id BG]
 >   |EQ [id BH]
 >   | |Switch [id BI]
 >   | | |EQ [id BJ]
 >   | | | |<int64> [id BF]
 >   | | | |ScalarConstant{1} [id Q]
 >   | | |neg [id BC]
 >   | | |<int64> [id BF]
 >   | |neg [id BC]
 >   |EQ [id BK]
 >     |Switch [id BI]
 >     |Composite [id BE]
Composite [id BE]
 >Abs [id L]

This issue cropped up in the implementation of AdvancedSubtensor.infer_shape, which originally did not support the boolean index case because of this? aesara-devs/aesara#39

In that case the argument against boolean index in particular didn't seem relevant, because advanced integer indexing can also be invalid:

import pytensor.tensor as pt
x = pt.vector("x")
x[[0, 20, 0]].shape.eval({x: [0, 1, 2]})  # [3], but there is no valid entry at position 20

But it touches on the very same question!

Originally posted by @ricardoV94 in #329 (comment)

[ricardoV94]

We would add a config flag assume_shapes_correct to toggle between simple and defensive graphs. This could also be used in the rewrites. I would probably leave the default to True