Faster convolve1d in numba backend #1378
Conversation
ricardoV94 force-pushed the faster_conv1d_numba branch 2 times, most recently from 66fa69a to 02823cc Compare 
Codecov Report❌ Patch coverage is
Additional details and impacted files
@@ Coverage Diff @@ ## main #1378 +/- ## ========================================== - Coverage 82.07% 82.01% -0.06% ========================================== Files 206 207 +1 Lines 49174 49250 +76 Branches 8720 8734 +14 ========================================== + Hits 40359 40394 +35 - Misses 6656 6692 +36 - Partials 2159 2164 +5
🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull Request Overview
This PR reimplements the core logic of convolve1d in the numba backend for a 6× speedup in benchmarks with small inputs, while also optimizing the gradient computation for valid convolutions when the smaller input’s shape is known statically. In addition, the PR renames Conv1d to Convolve1d for improved consistency in function naming and updates various test and dispatch files to reflect these changes.
- Renames Conv1d to Convolve1d across modules.
- Adds new tests for gradient optimization and benchmarks for numba convolve1d.
- Updates rewriting and dispatch code to support the new implementation.
Reviewed Changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| tests/tensor/signal/test_conv.py | Updated to import Convolve1d and added a test for gradient rewrite optimization. |
| tests/link/numba/signal/test_conv.py | Adjusted tests to optionally swap inputs, and added a benchmark test. |
| pytensor/tensor/signal/conv.py | Renamed Conv1d to Convolve1d and updated internal variable naming for clarity. |
| pytensor/tensor/rewriting/conv.py | Added a rewrite rule to optimize valid convolution gradients for static shapes. |
| pytensor/tensor/rewriting/init.py | Imported the new conv rewriting module. |
| pytensor/link/numba/dispatch/signal/conv.py | Updated to register Convolve1d and implemented specialized numba functions. |
| pytensor/link/jax/dispatch/signal/conv.py | Updated to register Convolve1d. |
jessegrabowski left a comment
jessegrabowski left a comment There was a problem hiding this comment.
lgtm, left ignorable suggestions
| | ||
| if ( | ||
| start == len_y - 1 | ||
| # equivalent to stop = conv.shape[-1] - len_y - 1 |
There was a problem hiding this comment.
Why not use that form then? I don't understand this comment
There was a problem hiding this comment.
Because I already extracted len_x, and I can use that directly
ricardoV94 force-pushed the faster_conv1d_numba branch from 02823cc to f1102ba Compare 
ricardoV94 force-pushed the faster_conv1d_numba branch from f1102ba to e2c8464 Compare These show up in the gradient of Convolve1D
ricardoV94 force-pushed the faster_conv1d_numba branch from e2c8464 to f0ef8fb Compare 2 participants
Reimplementing the core logic in the numba overload of
convolve/correlategives a speedup of 6x in the benchmarked test with relatively small inputs. I guess the overloads don't optimize/propagate constant checks as well? It's a bit surprising but the results are crystal clear.Also added a rewrite to optimize the gradient of valid convolutions wrt to the smallest inputs, in which case we don't need a full convolve. This is done at the rewrite level because static shape may not be known at the time of grad.
Finally, renamed
Conv1dtoConvolve1dwhich is more in line with the user-facing function