New batched_convolution can be slower for small datasets

[ricardoV94]

I noticed the changes in #1583 cause a slowdown in models with small datasets such as the MMM workflow example when using Numba / C backends (I didn't measure JAX as the JAX samplers will take gradients via JAX).

The speed regression in the numba backend (both in PyMC / nutpie) should be fixed by pymc-devs/pytensor#1378. I confirmed the model is faster than before for the 7 channels / ~188 days of data used the intro example. I also used freeze_data_and_dims which is the default in the non-PyMC samplers. The gradient of the convolution wrt to the smallest input is much faster that way.

Fixing the speed regression on the C backend requires more effort, as neither Blockwise nor Convolve1d have a native C implementation right now.

The old batched_convolution was basically an unrolled convolution loop (with a native C-implementation in the C-backend). It's not surprising this can sometimes be faster for small convolutions (at the cost of higher compile time), but the worst case scenario can be much worse (as showed in the timinigs in the original PR)

[daniel-s-tccc]

Just some more context for this story. On a project, we reverted to an old pymc-marketing because we were finding the new convolution to be quite slow in the optimizer. I did some profiling. Our typical dataset is 30 channels x 4 years so I generate random numbers about that size and asked for 1000 calls to a simple adstock-only pytensor graph.

The tldr is that which convolution is faster will depend a little bit on how many grad evaluations you are doing vs function evaluations and also the choice of backend. More function evals means the old is a bit faster. More grad evals means the new is MUCH faster.

Profiling work

pytensor.__version__ # 2.17.0

C backend

version of adstock	grad or func	time	comments
new	func	0.0110 seconds
new	grad	0.0001 seconds
old	func	0.0061 seconds	(old is twice as fast as the new one!)
old	grad	0.0364 seconds	(but grad is much slower

Numba backend

New method completely dominates the old for numba

version of adstock	grad or func	time	comments
new	func	0.0077 seconds
new	grad	0.0007 seconds
old	func	0.0120 seconds
old	grad	0.0863 seconds

Jax (CPU)

version of adstock	grad or func	time	comments
new	func	0.0006 seconds
new	grad	0.0001 seconds
old	func	0.0003 seconds	(also twice as fast lol but miniscule)
old	grad	0.0018 seconds

[daniel-s-tccc]

@cetagostini

[ricardoV94]

which version of PyTensor did you use? we speed up grad in the last release: pymc-devs/pytensor#1522

before that you could achieve the same by freezing the model dims or use nutpie/jax samplers that effectively do the same.

[ricardoV94]

ah you mean the forward is faster? In the C backend it's due to Blockwise using python mode, in JAX it's their thing. As I mentioned on the first message it's hard to beat a small unrolled loop, but the worst case scenario of the old approach was a disaster in run an compile times.

There's still work being done to speedup this approach including benchmarking conv2d to see if it's faster than batch1d (as it can use fft): pymc-devs/pytensor#1397

And further fussing forward and grad: pymc-devs/pytensor#1320

Also can you share your benchmark code?