pymc-devs / pymc-devs/pytensor

Avoid explicitly broadcasting RandomVariable inputs

Open
#1,562 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

graph rewriting memory opt random variables vectorization
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

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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.