Constant folding overrides broadcastable pattern

[ferrine]

Description of your problem or feature request

Please provide a minimal, self-contained, and reproducible example.

from aesara.graph.opt_utils import optimize_graph
v = at.patternbroadcast(at.ones((1, 1)), (False, False))
assert optimize_graph(v).broadcastable == v.broadcastable

Please provide the full traceback of any errors.

AssertionError                            Traceback (most recent call last)
/tmp/ipykernel_1243842/1172206837.py in <module>
      1 from aesara.graph.opt_utils import optimize_graph
      2 v = at.patternbroadcast(at.ones((1, 1)), (False, False))
----> 3 assert optimize_graph(v).broadcastable == v.broadcastable

AssertionError: 

Please provide any additional information below.

Versions and main components

• Aesara version: '2.3.6'
• Aesara config (python -c "import aesara; print(aesara.config)") not applicable
• Python version:
• Operating system:
• How did you install Aesara: (conda/pip) pip

[ricardoV94]

@ferrine It don't think it is valid to make that constant unbroadcastable. Changing the broadcastable flag only makes sense for variables where you don't know the runtime shape yet, and which may or may not be broadcastable. However that constant is known, and is by definition broadcastable.

[ferrine]

A more real-life example highlighting the issue

from aesara.graph.opt_utils import is_same_graph, optimize_graph
import aesara.tensor as at
import aesara
v = at.matrix()
v3 = v ** 3
const = at.ones((1, 1))
v3_const = aesara.clone_replace(v3, {v: const})
assert optimize_graph(v3_const).broadcastable == v3.broadcastable

Specifying the pattern does not help as well

from aesara.graph.opt_utils import is_same_graph, optimize_graph
import aesara.tensor as at
import aesara
v = at.matrix()
v3 = v ** 3
const = at.patternbroadcast(at.ones((1, 1)), (False, False))
v3_const = aesara.clone_replace(v3, {v: const})
assert optimize_graph(v3_const).broadcastable == v3.broadcastable

[ricardoV94]

To make that rewrite valid, you would have to do the other way around, that is, change the broadcast of the matrix to (true, true), since the matrix is a more general case than the const.

[ferrine]

But I don't have an idea of the replacement before the runtime, moreover I would want to change it from time to time

[ricardoV94]

If you are depending on the broadcastable pattern to be (False, False) you have to check either before or after the replacement. A matrix with fixed shape (1, 1) is by definition broadcastable across both dimensions.

[brandonwillard]

What you are observing is the little bit of type inference/synthesis/narrowing that Aesara is now doing. Take a look at the documentation on Type for more information and/or the PR that introduced this feature.

[brandonwillard]

Also, @ferrine, I'm glad you noticed this! It's an important new feature on which we need to start building more, and the first step in that process is making people aware of it and its purpose.