Description
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)