Faster convolve1d in numba backend

[ricardoV94]

Reimplementing the core logic in the numba overload of convolve/correlate gives 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 Conv1d to Convolve1d which is more in line with the user-facing function

[codecov]

Codecov Report

Attention: Patch coverage is 50.53763% with 46 lines in your changes missing coverage. Please review.

Project coverage is 82.01%. Comparing base (e98cbbc) to head (f0ef8fb).
Report is 16 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/numba/dispatch/signal/conv.py	26.08%	33 Missing and 1 partial ⚠️
pytensor/tensor/rewriting/conv.py	70.00%	6 Missing and 6 partials ⚠️

❌ Your patch status has failed because the patch coverage (50.53%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

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     

Files with missing lines	Coverage Δ
pytensor/link/jax/dispatch/signal/conv.py	100.00% <100.00%> (ø)
pytensor/tensor/signal/conv.py	97.05% <100.00%> (ø)
pytensor/tensor/rewriting/conv.py	70.00% <70.00%> (ø)
pytensor/link/numba/dispatch/signal/conv.py	32.00% <26.08%> (-58.91%)	⬇️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

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]

lgtm, left ignorable suggestions

[jessegrabowski]

Why not use that form then? I don't understand this comment

[ricardoV94]

Because I already extracted len_x, and I can use that directly