Skip to content

Fix stale output type bugs in rewrites#2204

Open
ricardoV94 wants to merge 1 commit into
pymc-devs:mainfrom
ricardoV94:stale_broadcastable_bugs
Open

Fix stale output type bugs in rewrites#2204
ricardoV94 wants to merge 1 commit into
pymc-devs:mainfrom
ricardoV94:stale_broadcastable_bugs

Conversation

@ricardoV94

@ricardoV94 ricardoV94 commented Jun 9, 2026

Copy link
Copy Markdown
Member

Follow up to #2169 (fixed by #2167).

I had claude do an exhaustive search of rewrites that may do similar fatal/wrong reasoning based on output.type that may become stale. Narrowed down to rewrites where it would return invalid results or raise errors.

It found two, which are fixed here with regression tests.

This is a symptom of a bad design. Broadcasting in always implicit from the input static types. If an input has static shape 1 along a dimension it is allowed to broadcast with other inputs by the Op. If it has None it is not. Rewrites though may replace a None dimension by a 1, and a graph that did not allow an input to broadcast now allows (see #408 for a suggested solution).

For rewrites, it means we can't do either of the following:

  1. Assume recreating the same node or equivalent one will have the same static output broadcastability.
  2. Look at the delta between output broadcastable and input broadcastable to gauge whether broadcasting is happening. Input shape of 1, output shape of None doesn't mean there is another None input that broadcasts it, it may be that all inputs are now 1 but the node is stale.

@ricardoV94

Copy link
Copy Markdown
Member Author
Agent report

@ricardoV94 ricardoV94 requested a review from jessegrabowski June 9, 2026 10:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant