fix pooling ops with explicit padding

[levskaya]

previous max_pool and avg_pool didn't work with explicit (e.g. ((1,1),(1,1)) padding specs) padding for the reduce window.
fixes #356

[levskaya]

Note we need to wait for JAX to cut a new release to test and use this modification.

[codecov-commenter]

Codecov Report

Merging #362 into master will increase coverage by 0.04%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master     #362      +/-   ##
==========================================
+ Coverage   78.87%   78.92%   +0.04%     
==========================================
  Files          34       34              
  Lines        2258     2263       +5     
==========================================
+ Hits         1781     1786       +5     
  Misses        477      477              

Impacted Files	Coverage Δ
flax/nn/pooling.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 6b58fc4...18a9ae5. Read the comment docs.

[levskaya]

Ok we can finally merge this!