Forbid runtime broadcasting in Elemwise #372

ricardoV94 · 2023-07-04T15:09:53Z

Related to #100
Related to #149
Related to #371

pytensor/tensor/elemwise.py

aseyboldt · 2023-07-07T16:45:26Z

pytensor/tensor/elemwise.py


-        out_shape = pytensor.tensor.broadcast_shape(*i_shapes, arrays_are_shapes=True)
+    def infer_shape(self, fgraph, node, i_shapes) -> List[Tuple[TensorVariable, ...]]:


I think this could just use this function: https://github.com/pymc-devs/pytensor/blob/main/pytensor/tensor/extra_ops.py#L1465

The make_node method doesn't seem to properly take into account the broadcastable flag either though, maybe that needs an update as well?

I didn't want to introduce checks or comparison between shapes, which that function does. This allows it to return a more optimized graph like Theano used to by assuming no invalid shapes were provided

The question then is whether we want to refactor that helper to do the same when arrays_are_shapes=False?

I think the make_node is correct insofar as it uses static shape and it's not possible to have broadcastable=False and shape=1

That one still requires some thinking and would be tackled in a separate PR.

I didn't want to introduce checks or comparison between shapes, which that function does. This allows it to return a more optimized graph like Theano used to by assuming no invalid shapes were provided

So we allow undefined behavior in the shapes and in rewrites? I'm not sure I see that much downside with having that check here...

But at least I think we shouldn't have this logic in both places. Maybe the function should have a flag if it should return shape with or without checks?

I am thinking we should add a config.assume_shapes_correct flag (default to True) to toggle that behavior in both shape_inference and rewrites that can return simplified cases.

Actually that helper works differently in that it expects either shapes or arrays, but here we are combining information from both shapes and arrays so it would require some refactoring. We don't want to simply pass node.inputs since infer_shape wants us to return a graph from ishapes.

I don't know if that is the right place to implement this logic since it is a user facing function. WDYT?

Okay I reverted to using the helper. Things are a bit weird in shape compilation because it will just use the static type shape of the node if that's available. Because the Elemwise make_node assumes valid shapes, the check introduced by infer_shape is only triggered when all dims are None.

Not much we can do about that then I think without a major rewrite of the shape handling...

pytensor/tensor/elemwise_cgen.py

aseyboldt · 2023-07-07T16:58:24Z

tests/tensor/test_elemwise.py

-        """
+    @staticmethod
+    def check_runtime_broadcast_error(mode):
+        """Check we emmit a clear error when runtime broadcasting would occur according to Numpy rules."""


I think I'd feel better if those tests were a bit more complete, ie inputs with different lengths etc...

Do you mean different runtime shapes (3 vs 5)? I am sure there are old tests for that already.

There are tests for invalid static shapes.

This one test was added when we specifically allowed runtime broadcasting in Aesara. The other thing I considered doing was to just remove it.

I'll confirm other tests for invalid shapes exist and maybe combine with this if they are not too convoluted.

Added a test for incompatible non-broadcast shapes. Let me know if you meant something else

tests/tensor/test_elemwise.py

tests/link/numba/test_elemwise.py

ricardoV94 · 2023-07-10T13:17:38Z

pytensor/link/numba/dispatch/elemwise_codegen.py

@@ -35,15 +35,20 @@ def compute_itershape(
                with builder.if_then(
                    builder.icmp_unsigned("!=", length, shape[i]), likely=False
                ):
-                    with builder.if_else(builder.icmp_unsigned("==", length, one)) as (
+                    with builder.if_else(
+                        builder.or_(


Weird the changes cause a SegmentationFault on the BroadcastTo numba test, but only on python 3.11? I couldn't replicate locally on 3.8 either. https://github.com/pymc-devs/pytensor/actions/runs/5507826611/jobs/10039563156?pr=372

Did I do something obviously wrong @aseyboldt?

I don't see anything wrong, I can try locally with py311 and if I can reproduce I can try to look at it in a debugger (with no debugging symbols, but well...)

If you can quickly try to reproduce that's already helpful (even if you don't dig down)

No luck so far, for me the tests run just fine...

It reliably segfaults here. I'll remove the numba changes for now and put the new test as an xfail

Does it segfault during the test_BroadcastTo test?

Yes... tests/link/numba/test_extra_ops.py::test_BroadcastTo[x0-shape0].

https://github.com/pymc-devs/pytensor/actions/runs/5507826611/jobs/10039563156?pr=372#step:6:281

But I don't see how it could be a problem in those tests. There is nothing else in the compiled graph other than the BroadcastTo

codecov-commenter · 2023-07-11T08:40:44Z

Codecov Report

Merging #372 (d044271) into main (5c87d74) will decrease coverage by 0.01%.
The diff coverage is 100.00%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #372      +/-   ##
==========================================
- Coverage   80.40%   80.40%   -0.01%     
==========================================
  Files         156      156              
  Lines       45401    45390      -11     
  Branches    11106    11103       -3     
==========================================
- Hits        36505    36496       -9     
  Misses       6689     6689              
+ Partials     2207     2205       -2

Impacted Files	Coverage Δ
pytensor/tensor/extra_ops.py	`89.00% <ø> (ø)`
pytensor/link/jax/dispatch/elemwise.py	`81.69% <100.00%> (+1.09%)`	⬆️
pytensor/tensor/elemwise.py	`88.07% <100.00%> (+0.01%)`	⬆️
pytensor/tensor/elemwise_cgen.py	`95.34% <100.00%> (-0.40%)`	⬇️

... and 2 files with indirect coverage changes

aseyboldt

Looks good :-)

ricardoV94 added major shape problem Op implementation shape inference labels Jul 4, 2023

ricardoV94 force-pushed the revert_dynamic_broadcast_elemwise branch 2 times, most recently from d1a0ff7 to 3c1d876 Compare July 4, 2023 15:18

ricardoV94 added jax backend compatibility C-backend labels Jul 4, 2023

ricardoV94 mentioned this pull request Jul 4, 2023

Static broadcast #149

Closed

6 tasks

ricardoV94 force-pushed the revert_dynamic_broadcast_elemwise branch 3 times, most recently from 3beaaff to 60b6d6f Compare July 5, 2023 13:52

ricardoV94 marked this pull request as ready for review July 5, 2023 14:07

ricardoV94 requested review from aseyboldt, michaelosthege and ferrine July 5, 2023 14:29

ricardoV94 mentioned this pull request Jul 7, 2023

HalfNormal in JAX failing due to implicit downcasting of constant 0d TensorVariable to float #373

Closed

aseyboldt reviewed Jul 7, 2023

View reviewed changes

ricardoV94 commented Jul 8, 2023

View reviewed changes

tests/tensor/test_elemwise.py Outdated Show resolved Hide resolved

ricardoV94 force-pushed the revert_dynamic_broadcast_elemwise branch from 60b6d6f to a21ae05 Compare July 10, 2023 08:05

ricardoV94 commented Jul 10, 2023

View reviewed changes

tests/link/numba/test_elemwise.py Outdated Show resolved Hide resolved

ricardoV94 force-pushed the revert_dynamic_broadcast_elemwise branch from a21ae05 to b2c2743 Compare July 10, 2023 08:19

ricardoV94 added the numba label Jul 10, 2023

ricardoV94 force-pushed the revert_dynamic_broadcast_elemwise branch 2 times, most recently from a26e46b to f3ad19a Compare July 10, 2023 12:04

ricardoV94 commented Jul 10, 2023

View reviewed changes

ricardoV94 force-pushed the revert_dynamic_broadcast_elemwise branch from eb98809 to 1ab333d Compare July 11, 2023 08:12

ricardoV94 force-pushed the revert_dynamic_broadcast_elemwise branch from 1ab333d to 28b3b46 Compare July 11, 2023 10:00

ricardoV94 added 2 commits July 11, 2023 12:37

Forbid runtime broadcasting in Elemwise

3806650

Revert numba runtime broadcast check

d044271

ricardoV94 force-pushed the revert_dynamic_broadcast_elemwise branch from 28b3b46 to d044271 Compare July 11, 2023 10:37

ricardoV94 requested a review from aseyboldt July 11, 2023 11:52

aseyboldt approved these changes Jul 12, 2023

View reviewed changes

ricardoV94 merged commit 981be2a into pymc-devs:main Jul 12, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Forbid runtime broadcasting in Elemwise #372

Forbid runtime broadcasting in Elemwise #372

ricardoV94 commented Jul 4, 2023 •

edited

Loading

aseyboldt Jul 7, 2023

ricardoV94 Jul 8, 2023 •

edited

Loading

ricardoV94 Jul 8, 2023

aseyboldt Jul 10, 2023

ricardoV94 Jul 10, 2023

ricardoV94 Jul 11, 2023 •

edited

Loading

ricardoV94 Jul 11, 2023 •

edited

Loading

aseyboldt Jul 12, 2023

aseyboldt Jul 7, 2023

ricardoV94 Jul 8, 2023 •

edited

Loading

ricardoV94 Jul 10, 2023 •

edited

Loading

ricardoV94 Jul 10, 2023

aseyboldt Jul 10, 2023

ricardoV94 Jul 10, 2023

aseyboldt Jul 11, 2023

ricardoV94 Jul 11, 2023

aseyboldt Jul 11, 2023

ricardoV94 Jul 11, 2023

ricardoV94 Jul 11, 2023 •

edited

Loading

codecov-commenter commented Jul 11, 2023 •

edited

Loading

aseyboldt left a comment


		out_shape = pytensor.tensor.broadcast_shape(*i_shapes, arrays_are_shapes=True)
		def infer_shape(self, fgraph, node, i_shapes) -> List[Tuple[TensorVariable, ...]]:

Forbid runtime broadcasting in Elemwise #372

Forbid runtime broadcasting in Elemwise #372

Conversation

ricardoV94 commented Jul 4, 2023 • edited Loading

Choose a reason for hiding this comment

ricardoV94 Jul 8, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Jul 11, 2023 • edited Loading

Choose a reason for hiding this comment

ricardoV94 Jul 11, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Jul 8, 2023 • edited Loading

Choose a reason for hiding this comment

ricardoV94 Jul 10, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Jul 11, 2023 • edited Loading

Choose a reason for hiding this comment

codecov-commenter commented Jul 11, 2023 • edited Loading

Codecov Report

aseyboldt left a comment

Choose a reason for hiding this comment

ricardoV94 commented Jul 4, 2023 •

edited

Loading

ricardoV94 Jul 8, 2023 •

edited

Loading

ricardoV94 Jul 11, 2023 •

edited

Loading

ricardoV94 Jul 11, 2023 •

edited

Loading

ricardoV94 Jul 8, 2023 •

edited

Loading

ricardoV94 Jul 10, 2023 •

edited

Loading

ricardoV94 Jul 11, 2023 •

edited

Loading

codecov-commenter commented Jul 11, 2023 •

edited

Loading