Implement batched convolve1d #1318

ricardoV94 · 2025-03-24T13:49:35Z

Spinoff from #548. It was obvious the old fast Ops were not really fast, just not as slow as the AbstractConv Ops. I suggest we start removing those, by more sane implementations.

This PR implements a simple convolve1d (that can be batched by blockwise). This is actually not offered by either scipy.signal nor numpy.convolve. The scipy impl has no concept of batch dims, so if you pass two 2d inputs you get a convolve2d. The numpy one doesn't support batch dims, period.

Because the 1d case implemented here matches with the one you would get from calling scipy.convolve I added a small wrapper that allows you to go that route. Consider it WIP.

Also note that same behaves differently between numpy and scipy. In numpy it returns something with the shape of the largest of the two inputs, in scipy with the shape of the first input. The one in this PR behaves like scipy. This may be a good argument to go down the jax route and implement a pytensor.numpy and a pytensor.scipy, because such differences are likely to hit users, specially if they don't know which one we are mimicking.

This Op can be used to replace batched_convolution in pymc-marketing with much better graph compile time and runtime (specially as the size of the convolution grows).

📚 Documentation preview 📚: https://pytensor--1318.org.readthedocs.build/en/1318/

ricardoV94 · 2025-03-24T14:39:23Z

One day I'll murder floatX="float32", I swear

pytensor/tensor/signal/conv.py

codecov · 2025-03-24T15:35:36Z

Codecov Report

Attention: Patch coverage is 96.73913% with 3 lines in your changes missing coverage. Please review.

Project coverage is 82.03%. Comparing base (8a7356c) to head (39ab436).
Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/signal/conv.py	97.05%	1 Missing and 1 partial ⚠️
pytensor/link/numba/dispatch/signal/conv.py	90.90%	1 Missing ⚠️

❌ Your patch status has failed because the patch coverage (96.73%) 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    #1318      +/-   ##
==========================================
+ Coverage   81.98%   82.03%   +0.04%     
==========================================
  Files         188      193       +5     
  Lines       48489    48622     +133     
  Branches     8673     8684      +11     
==========================================
+ Hits        39756    39886     +130     
- Misses       6582     6584       +2     
- Partials     2151     2152       +1

Files with missing lines	Coverage Δ
pytensor/link/jax/dispatch/__init__.py	`100.00% <100.00%> (ø)`
pytensor/link/jax/dispatch/signal/__init__.py	`100.00% <100.00%> (ø)`
pytensor/link/jax/dispatch/signal/conv.py	`100.00% <100.00%> (ø)`
pytensor/link/numba/dispatch/__init__.py	`100.00% <100.00%> (ø)`
pytensor/link/numba/dispatch/signal/__init__.py	`100.00% <100.00%> (ø)`
pytensor/link/numba/dispatch/signal/conv.py	`90.90% <90.90%> (ø)`
pytensor/tensor/signal/conv.py	`97.05% <97.05%> (ø)`

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

jessegrabowski · 2025-03-24T15:35:24Z

pytensor/tensor/signal/conv.py

+    return cast(TensorVariable, Blockwise(Conv1d(mode=mode))(in1, in2))
+
+
+def convolve(


I'm not in love with the name and docstring. The docstring says it works on two n-dim arrays, but then raises if they're not 1d? But the name doesn't imply anything about the shapes?

Is this something that is meant to be a generic front-end to more functionality in the future?

This is supposed to match the scipy one but right now only supports 1d. It raises a NotImplementedError, not a ValueError. The idea is when we add an arbitrary order conv this function will dispatch to those

Want me to remove this until we actually support all the cases the scipy one does?

I think that's better. As it stands, it doesn't add anything (except confusion)

jessegrabowski · 2025-03-24T15:37:43Z

pytensor/tensor/signal/conv.py

+        in1_batch_shape = tuple(in1.shape)[:-1]
+        zeros_left = in2.shape[0] // 2
+        zeros_right = (in2.shape[0] - 1) // 2
+        in1 = join(


pad wasn't useful here?

I don't want pad until we figure out inline for it. I want PyTensor to optimize across the boundary, specially when gradients get involved

pytensor/tensor/signal/conv.py

jessegrabowski · 2025-03-24T15:40:33Z

tests/link/numba/signal/test_conv.py

+    rng = np.random.default_rng()
+    test_x = rng.normal(size=(3, 5))
+    test_y = rng.normal(size=(7, 11))
+    # Object mode is not supported for numba


Because we didn't implement it. This is only used for coverage and wouldn't cover the relevant code since it's all low level numba

Specifically it's the numba dispatch of Blockwise that doesn't have an obj implementation.. I'll tweak the comment

Clarified comment

jessegrabowski · 2025-03-24T15:40:46Z

tests/link/jax/signal/test_conv.py

+    out = convolve1d(x[None], y[:, None], mode=mode)
+
+    rng = np.random.default_rng()
+    test_x = rng.normal(size=(3, 5))


Suggested change

test_x = rng.normal(size=(3, 5))

test_x = rng.normal(size=(3, 5)).astype(floatX)

?

Nope no stupid float32 in jax/numba tests and also I explicitly used dmatrix

jessegrabowski · 2025-03-24T15:41:22Z

tests/tensor/signal/test_conv.py

+    y = matrix("kernel")
+    out = convolve1d(x, y)
+
+    # Convolution is unchanged by order


I don't understand the comment

Doesn't matter if x, y, or y,x. I'm using that property to confirm batching is working as expected

Clarified comment

zaxtax · 2025-03-25T14:06:17Z

It would also be good to have like checkboxes on an issue of what we still need to support

…

On Tue, 25 Mar 2025, 14:11 Ricardo Vieira, ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In pytensor/tensor/signal/conv.py <#1318 (comment)>: > + # We implement "same" as "valid" with padded `in1`. + in1_batch_shape = tuple(in1.shape)[:-1] + zeros_left = in2.shape[0] // 2 + zeros_right = (in2.shape[0] - 1) // 2 + in1 = join( + -1, + zeros((*in1_batch_shape, zeros_left), dtype=in2.dtype), + in1, + zeros((*in1_batch_shape, zeros_right), dtype=in2.dtype), + ) + mode = "valid" + + return cast(TensorVariable, Blockwise(Conv1d(mode=mode))(in1, in2)) + + +def convolve( Want me to remove this until we actually support all the cases the scipy one does? — Reply to this email directly, view it on GitHub <#1318 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUIXPA46ZTJND65OJJL2WFBXLAVCNFSM6AAAAABZU66RNSVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZDOMJTGY4TANJRGY> . You are receiving this because your review was requested.Message ID: ***@***.***>

jessegrabowski

Remove convolve for now and I think this is good to merge.

If you don't want to lose track of the code, just removing it from signal.init.__all__ is probably good enough, and add a one-time-per-runtime warning that the function is a work in progress and users should prefer convolve1d for now

ricardoV94 · 2025-03-27T10:48:34Z

I'll remove it

Co-authored-by: Rob Zinkov <zaxtax@users.noreply.github.com> Co-authored-by: Ricardo Vieira <28983449+ricardov94@users.noreply.github.com>

ricardoV94 added SciPy compatibility convolution enhancement labels Mar 24, 2025

ricardoV94 mentioned this pull request Mar 24, 2025

Convolve1d rewrites #1319

Open

ricardoV94 force-pushed the implement_convolve branch from 0fa631a to fcd70c1 Compare March 24, 2025 14:12

ricardoV94 mentioned this pull request Mar 24, 2025

Replace batched_convolution by PyTensor native implementation pymc-labs/pymc-marketing#1583

Draft

ricardoV94 force-pushed the implement_convolve branch from fcd70c1 to 8e96c41 Compare March 24, 2025 14:41

ricardoV94 requested review from jessegrabowski and zaxtax March 24, 2025 14:41

ricardoV94 commented Mar 24, 2025

View reviewed changes

pytensor/tensor/signal/conv.py Show resolved Hide resolved

ricardoV94 force-pushed the implement_convolve branch from 8e96c41 to 7c23e43 Compare March 24, 2025 14:55

jessegrabowski reviewed Mar 24, 2025

View reviewed changes

ricardoV94 mentioned this pull request Mar 25, 2025

Add correlate helper function #1321

Open

jessegrabowski requested changes Mar 27, 2025

View reviewed changes

ricardoV94 force-pushed the implement_convolve branch from 7c23e43 to 8b55a71 Compare March 27, 2025 11:10

ricardoV94 requested a review from jessegrabowski March 27, 2025 11:12

Implement batched convolve1d

39ab436

Co-authored-by: Rob Zinkov <zaxtax@users.noreply.github.com> Co-authored-by: Ricardo Vieira <28983449+ricardov94@users.noreply.github.com>

ricardoV94 force-pushed the implement_convolve branch from 8b55a71 to 39ab436 Compare March 27, 2025 11:55

jessegrabowski approved these changes Mar 27, 2025

View reviewed changes

jessegrabowski merged commit 0b56ed9 into pymc-devs:main Mar 27, 2025
72 of 73 checks passed

This was referenced Apr 1, 2025

BUG: convolve1d mode="same" indexing error #1336

Closed

Fix indexing in convolve1d with mode="same" #1337

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement batched convolve1d #1318

Implement batched convolve1d #1318

ricardoV94 commented Mar 24, 2025 •

edited

Loading

ricardoV94 commented Mar 24, 2025

codecov bot commented Mar 24, 2025 •

edited

Loading

jessegrabowski Mar 24, 2025

ricardoV94 Mar 24, 2025

ricardoV94 Mar 25, 2025

jessegrabowski Mar 27, 2025

jessegrabowski Mar 24, 2025

ricardoV94 Mar 24, 2025

jessegrabowski Mar 24, 2025

ricardoV94 Mar 24, 2025

ricardoV94 Mar 24, 2025

ricardoV94 Mar 25, 2025

jessegrabowski Mar 24, 2025

ricardoV94 Mar 24, 2025

jessegrabowski Mar 24, 2025

ricardoV94 Mar 24, 2025

ricardoV94 Mar 25, 2025

zaxtax commented Mar 25, 2025 via email

jessegrabowski left a comment

ricardoV94 commented Mar 27, 2025

		return cast(TensorVariable, Blockwise(Conv1d(mode=mode))(in1, in2))


		def convolve(

	test_x = rng.normal(size=(3, 5))
	test_x = rng.normal(size=(3, 5)).astype(floatX)

Implement batched convolve1d #1318

Implement batched convolve1d #1318

Conversation

ricardoV94 commented Mar 24, 2025 • edited Loading

ricardoV94 commented Mar 24, 2025

codecov bot commented Mar 24, 2025 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zaxtax commented Mar 25, 2025 via email

jessegrabowski left a comment

Choose a reason for hiding this comment

ricardoV94 commented Mar 27, 2025

ricardoV94 commented Mar 24, 2025 •

edited

Loading

codecov bot commented Mar 24, 2025 •

edited

Loading