pymc-devs / pymc-devs/pytensor
Avoid explicitly broadcasting RandomVariable inputs
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
RVs broadcast batch inputs by default. We are not avoiding materializing broadcast (which in PyTensor is always dense), similar to how we're failing to do it with #1561
import pytensor
import pytensor.tensor as pt
mu = pt.vector("x")
sigma = pt.vector("sigma")
mu_b, sigma_b = pt.broadcast_arrays(mu[:, None], sigma[None, :])
out = pt.random.normal(mu_b, sigma_b)
fn = pytensor.function([mu, sigma], [out])
fn.dprint()
# normal_rv{"(),()->()"}.1 [id A] 4
# ├─ RNG(<Generator(PCG64) at 0x7F9655BAC120>) [id B]
# ├─ NoneConst{None} [id C]
# ├─ Second [id D] 3
# │ ├─ ExpandDims{axis=0} [id E] 0
# │ │ └─ sigma [id F]
# │ └─ ExpandDims{axis=1} [id G] 1
# │ └─ x [id H]
# └─ Second [id I] 2
# ├─ ExpandDims{axis=1} [id G] 1
# │ └─ ···
# └─ ExpandDims{axis=0} [id E] 0
# └─ ···
The Second operation will iterate over each entry of sigma to copy mu, and over each entry of mu to copy sigma. We should write those as Alloc, but that's currently not implemented: https://github.com/pymc-devs/pytensor/blob/6770f46ed575f8f2d5367146d678d00c1c0b1c0b/pytensor/tensor/rewriting/basic.py#L402-L405
More importantly we shouldn't do this at all for RVs, as they automatically broadcast! This is equivalent to
pt.random.normal(mu[:, None], sigma[None, :])
This redundant broadcast is also true when there's a size argument:
import pytensor
import pytensor.tensor as pt
mu = pt.vector("x")
out = pt.random.exponential(pt.broadcast_to(mu, (5, 3)), size=(5, 3))
fn = pytensor.function([mu], [out])
fn.dprint()
# exponential_rv{"()->()"}.1 [id A] 1
# ├─ RNG(<Generator(PCG64) at 0x7F9655BAF220>) [id B]
# ├─ [5 3] [id C]
# └─ Alloc [id D] 0
# ├─ x [id E]
# ├─ 5 [id F]
# └─ 3 [id G]
An alloc on batch dimensions is never necessary if we have size.
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 with the random-variable examples in the issue and inspect the graph produced by pytensor.function and dprint(). Read pytensor/tensor/rewriting/basic.py around the linked rewrite code, then trace the relevant RV broadcasting behavior. Done means the redundant broadcast and batch-dimension Alloc operations no longer appear in the shown graphs.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100