Skip to content

Add max and average pooling, dispatched on every backend - #100

Merged
jessegrabowski merged 7 commits into
pymc-devs:mainfrom
jessegrabowski:pooling
Aug 20, 2026
Merged

Add max and average pooling, dispatched on every backend#100
jessegrabowski merged 7 commits into
pymc-devs:mainfrom
jessegrabowski:pooling

Conversation

@jessegrabowski

Copy link
Copy Markdown
Member

A conv stack had nothing to downsample with. MaxPool1D/MaxPool2D and AvgPool1D/AvgPool2D reduce the same windows a convolution correlates over, so they reuse the gather rather than repeating it, and each backend pools with its own primitive.

Max pooling deliberately departs from pytensor's max gradient. pt.max routes the full cotangent to every tap tied for the maximum, so a window returns more gradient than it received — and after a rectifier, whole windows of ties are routine rather than measure-zero. Selecting through argmax gives it to one tap, as jax and torch do. mlx splits it evenly instead, so the backend tests assert the gradient is conserved rather than where it lands.

@codecov-commenter

codecov-commenter commented Aug 20, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 95.75290% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.08%. Comparing base (af3297f) to head (537bb36).

Files with missing lines Patch % Lines
pytensor_ml/dispatch/pytorch/pooling.py 84.37% 5 Missing ⚠️
pytensor_ml/layers/conv.py 95.69% 4 Missing ⚠️
pytensor_ml/dispatch/mlx/pooling.py 92.00% 2 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (af3297f) and HEAD (537bb36). Click for more details.

HEAD has 1 upload less than BASE
Flag BASE (af3297f) HEAD (537bb36)
8 7
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #100      +/-   ##
==========================================
- Coverage   97.43%   91.08%   -6.35%     
==========================================
  Files          56       59       +3     
  Lines        2609     2850     +241     
==========================================
+ Hits         2542     2596      +54     
- Misses         67      254     +187     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@jessegrabowski
jessegrabowski merged commit 1c996ea into pymc-devs:main Aug 20, 2026
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants