pymc-devs / pymc-devs/pytensor

Encode broadcasting explicitly, not implicitly from "volatile" static shape of inputs

Open
#408 1 comment 0 reactions 0 assignees View on GitHub

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:

  1. Not force static shape to be unknown
  2. 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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.