Skip to content

Restore fast convolution Ops, rewrites, and docs #548

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Draft
wants to merge 7 commits into
base: main
Choose a base branch
from
Draft

Restore fast convolution Ops, rewrites, and docs #548

wants to merge 7 commits into from

Conversation

jessegrabowski
Copy link
Member

@jessegrabowski jessegrabowski commented Dec 11, 2023 •
edited by ricardoV94
Loading

Description

When the nnet sub-module was depreciated, the old theano convolution functions went into the trashcan, along with the associated docs. This was partially reverted, but the docs and the "efficient" versions of the Ops are still missing.

There is interest in convolutions on the discourse from time-to-time, so it's worth talking about these.

This PR is a draft because I just restored everything. To be frank this part of the library seems somewhat bloated, and would probably be better served by an overhaul than by just restoring this code. For example, there is commentary about CUDA kernels in the docstrings -- I assume this is from the days when thenao was trying to target GPUs with pycuda. We use JAX these days, and I wager these Ops don't compile to JAX at all (though I haven't tried), or numba for that matter...

It's marked as draft for 2 reasons:

  1. It might be better to just re-implement convolution, be it 1d or 2d, in a manner that more closely matches the numpy and scipy signatures, just using pure Pytensor Ops derived from scan. For batches of convolutional filters, I'm pretty sure we can just use Blockwise.
  2. Even if we think this existing code is great for CPU optimization (which probably remains the most common use case for our users), I definitely restored too much stuff. Do we really want all 3d convolutions? Also there's tons of references to "use this one not that one" -- I haven't taken the time to figure out which on is the good one. If we want to keep that, we should just rip all the other ones out

Comments requested.

Related Issue

Checklist

Type of change

  • New feature / enhancement
  • Bug fix
  • Documentation
  • Maintenance
  • Other (please specify):

@ricardoV94
Copy link
Member

+1 Follow numpy/scipy API would be the best. Part of the reason we usually do that is it saves us from having to think about it, and users from having to learn something new.

So Scipy signal seems to have convolve/convolve2d and correlate (I saw we have corr something?) and they are already batched (at least convolve is)? JAX also followed their API: https://jax.readthedocs.io/en/latest/notebooks/convolutions.html

I would suggest trying to implement an Op like the abstract ones, perhaps way simpler that corresponds to this API, and then specialize via rewrites to the fancy ones. Here if you think one of them is way less relevant feel free to kick it out.

I would focus first on mapping to JAX (and numba if they have anything) and if some of these cases have a nice correspondence to the old C Ops all the better, if not one more reason to kick them out.

Also I see they had old logic about shape inference for the C code, we should update it to use new static shape stuff

This was referenced May 3, 2024
Open
Merged
@rom1v rom1v mentioned this pull request Feb 22, 2025
Open
@ricardoV94 ricardoV94 mentioned this pull request Mar 24, 2025
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
docs feature request NumPy compatibility Op implementation request discussion
Projects
PyTensor hackathon issues
Status: No status
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

Efficient Conv Ops were removed accidentally Missing conv documentation
2 participants
@jessegrabowski @ricardoV94