pymc-devs / pymc-devs/pytensor

Rewrite scalar dot as multiplication

Open
#1,205 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

beginner friendly graph rewriting
Dominant language
Python
Stars
644
Forks
209
Avg merge
2d 14h
Merged PRs (30d)
16

Description

Description

In https://github.com/pymc-devs/pytensor/pull/1178 we rewrite batched dots that are just multiplication away, but left core dots the same due to use of BLAS operations for those (whether they are worth it or not is a question on its own). But there is one case that is definitely not worth it: scalar multiplication.

The following graph should definitely be simplified:

import pytensor
import pytensor.tensor as pt

x = pt.tensor("x", shape=(1, 1))
y = pt.tensor("y", shape=(1, 1))
out = x @ y
pytensor.function([x, y], out).dprint()
CGer{non-destructive} [id A] 2
 ├─ [[0.]] [id B]
 ├─ 1.0 [id C]
 ├─ DropDims{axis=1} [id D] 1
 │  └─ x [id E]
 └─ DropDims{axis=0} [id F] 0
    └─ y [id G]

Or without BLAS stuff

pytensor.function([x, y], out, mode="FAST_COMPILE").dprint()
Dot22 [id A] 0
 ├─ x [id B]
 └─ y [id C]

Those should just be mul because that can be fused with other Elemwise operations (and calling BLAS for it is the silliest thing ever)

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

Reproduce the issue with the provided pytensor.tensor examples and inspect the optimizer path used by pytensor.function, including FAST_COMPILE and dprint output. The change is done when scalar matrix dots are represented as mul rather than Dot22 or CGer, while preserving the shown computation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.