Skip to content

Add Flatten, and take batch norm's statistics per channel - #101

Merged
jessegrabowski merged 5 commits into
pymc-devs:mainfrom
jessegrabowski:flatten
Aug 21, 2026
Merged

Add Flatten, and take batch norm's statistics per channel#101
jessegrabowski merged 5 commits into
pymc-devs:mainfrom
jessegrabowski:flatten

Conversation

@jessegrabowski

Copy link
Copy Markdown
Member

A conv stack had no way to reach a dense head, and batch norm placed after a convolution was computing the wrong thing. Flatten collapses everything after the batch axis. Batch norm now reduces over every axis but the last, so a channel gets one mean and variance instead of every pixel position getting its own against the batch. It also rejects rank-1 input now, which has no batch axis to take statistics over and used to normalize across features instead.

BatchNorm2D is renamed to BatchNorm, since the reduction rule is identical at every rank. That's what flax and mlx each do with a single class; torch's 1d/2d/3d split is three rank validators over one shared rule rather than three semantics.

@review-notebook-app

Copy link
Copy Markdown

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@codecov-commenter

codecov-commenter commented Aug 21, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.89%. Comparing base (1c996ea) to head (5a46947).

Additional details and impacted files
@@           Coverage Diff           @@
##             main     #101   +/-   ##
=======================================
  Coverage   97.89%   97.89%           
=======================================
  Files          59       59           
  Lines        2850     2856    +6     
=======================================
+ Hits         2790     2796    +6     
  Misses         60       60           

☔ 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 e867086 into pymc-devs:main Aug 21, 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