pymc-devs / pymc-devs/pytensor
Encode broadcasting explicitly, not implicitly from "volatile" static shape of inputs
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Latest proposal
Make the broadcasting behavior, not a function of the static type of the Variables but of the Ops that do the implicit broadcasting. So Elemwise may have a broadcast_pattern with the indexes of the inputs that will be broadcasted along each dimension. That way if a graph replacement or a rewrite provides something with more defined static shape (1,) instead of the original (None,) the Op will not change it's behavior, and we no longer need this distinction at the type level.
The first time an Op is created, we would read it from the input static types, so everything works with backwards compatibility. This means however, that there would no longer be a single instance of non-unary Elemwises like pt.add, they would need to be created by the helper functions (or like we did for RVs with the dtype argument, with the call method).
This makes sense to me as the same core operation with and without broadcasting are semantically quite different
Old proposal
Less in favor because broacasting is not really a property of the input, but the Op.
Allow TensorType(shape=(1,), broadcastable=(False,))
This requires re-introducing the broadcastable flags as independent from shape. It seems needed to:
- Not force static shape to be unknown
- Not change the meaning of the graph accidentally due to shape inference / rewrites
Affected Ops (anything that performs broadcasting of existing dims):
- Elemwise
- Alloc
- GEMM Ops
- Unbroadcast
Will require re-introducing Rebroadcast which could toggle broadcastable flags directly independently from static shape gains from SpecifyShape. Probably better named SpecifyBroadcastable.
Probably Elemwise outputs will have to be unbroadcastable as long as at least on input is also unbroadcastable along the same dimension.
x = pt.vector(shape=(1,), broadcastable=(False,))
y = x + x
assert y.type.broadcastable == (False,)
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reading the TensorType and Elemwise behavior described in the proposal, then trace how Alloc, GEMM Ops, and Unbroadcast handle broadcasting. The work is done when broadcasting behavior is explicit at the operation level, remains stable through graph rewrites, and preserves the stated backward-compatibility expectations.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100