pymc-devs / pymc-devs/pytensor

Convolve1d rewrites

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

Nobody has claimed this yet.

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

Description

Description

Follow up to #1318

Some of these may apply to higher-dimensional convolutions.

Sum of convolution as product of sum of inputs

If we have convolve1d(x, y, mode="full").sum(), we can rewrite it as x.sum() * y.sum()

import numpy as np

rng = np.random.default_rng(37)
x = rng.normal(size=(5))
y = rng.normal(size=(9))
np.testing.assert_allclose(np.convolve(x, y).sum(), x.sum() * y.sum())

convolve constant kernels

If we have two sequential applications of convolve, with constant inputs we can convolve the constant inputs first, reducing the number of runtime convolutions

import numpy as np

rng = np.random.default_rng(37)
x = rng.normal(size=(5))
y1 = rng.normal(size=(9))
y2 = rng.normal(size=(3,))

r1 = np.convolve(np.convolve(x, y1), y2)
r2 = np.convolve(x, np.convolve(y1, y2))
np.testing.assert_allclose(r1, r2)

Merge convolutions with flipped inputs

Convolutions give the same output regardless of order of inputs, so if we see two with the same inputs but in different order we can merge them. In general, we may want to have an Op property that tells us when the output is invariant to the order of inputs to apply such merge automatically. This Applies to Add, Mul, ...

import numpy as np

rng = np.random.default_rng(37)
x = rng.normal(size=(5))
y = rng.normal(size=(9))

r1 = np.convolve(x, y)
r2 = np.convolve(y, x)
np.testing.assert_allclose(r1, r2)

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 follow-up #1318 and trace the existing convolve1d rewrite path. Compare the three NumPy examples and determine how each proposed optimization should be represented and validated, including input-order invariance. Done means the applicable rewrites are implemented with coverage for the stated equivalences.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
compilers, performance
Issue type
Feature
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.