Skip to content

Commit 907e5cc

Browse files
committed
fix(ci): repair adaptive pooling ZeroDivisionError and mypy numpy stubs
Two standing CI failures, both from dependency drift since the last green run (2026-06-19). Neither was caused by the recent _version.py move (that only broke the editable-install build, already fixed in #866). Continuous Integration (test failure): brainpy/dnn/pooling.py::_adaptive_pool1d divided by `size // target_size`, which is 0 whenever target_size > size (a spatial dim smaller than its target, e.g. AdaptiveMaxPool3d(target=[6,5,4]) on an input axis of size 2). A newer jax evaluates `arr.size % math.prod(...)` inside reshape, turning the resulting reshape(-1, 0) into ZeroDivisionError (4 failed / 4179 passed). Rewrite it with PyTorch-style adaptive bins -- output i = operation(x[floor(i*size/T) : ceil((i+1)*size/T)]) -- which never yields an empty window and handles both down- and up-sampling. Verified against the exact PyTorch formula (mean & max) and under vmap. Add a positive-target_size guard and regression tests. Type Checking (mypy failure): numpy 2.5.1 ships PEP 695 `type` aliases in its stubs, which mypy rejects as a syntax error unless the analysis target is Python 3.12+. Bump [tool.mypy] python_version 3.11 -> 3.14.
1 parent b75464a commit 907e5cc

3 files changed

Lines changed: 95 additions & 12 deletions

File tree

brainpy/dnn/pooling.py

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -703,30 +703,47 @@ def __init__(
703703
def _adaptive_pool1d(x, target_size: int, operation: Callable):
704704
"""Adaptive pool 1D.
705705
706+
Reduces a 1-D array to ``target_size`` values using PyTorch-style adaptive
707+
pooling bins: output ``i`` is ``operation`` applied to the input window
708+
``x[floor(i * size / target_size) : ceil((i + 1) * size / target_size)]``.
709+
This works for any ``size`` relative to ``target_size`` — including
710+
``target_size > size`` (upsampling, where windows shrink to a single element
711+
and are repeated across outputs).
712+
706713
Parameters
707714
----------
708715
x
709716
The input. Should be a JAX array of shape `(dim,)`.
710717
target_size : int
711718
The shape of the output after the pooling operation `(target_size,)`.
712719
operation : Callable
713-
The pooling operation to be performed on the input array.
720+
The pooling operation to be performed on the input array. It must reduce a
721+
1-D array to a scalar (e.g. ``jax.numpy.mean`` or ``jax.numpy.max``).
714722
715723
Returns
716724
-------
717725
A JAX array of shape `(target_size, )`.
726+
727+
Notes
728+
-----
729+
Bin boundaries are static (derived from ``x``'s shape and ``target_size``), so
730+
the per-bin comprehension is unrolled once at trace time rather than executed
731+
step-by-step at runtime. Every window is non-empty because
732+
``ceil((i + 1) * size / target_size) > floor(i * size / target_size)`` for
733+
``size >= 1``; the previous block-reshape implementation instead divided by
734+
``size // target_size`` and raised ``ZeroDivisionError`` whenever
735+
``target_size > size``.
718736
"""
737+
if target_size <= 0:
738+
raise ValueError(f"target_size must be a positive integer, got {target_size}.")
719739
x = bm.as_jax(x)
720-
size = jnp.size(x)
721-
num_head_arrays = size % target_size
722-
num_block = size // target_size
723-
if num_head_arrays != 0:
724-
head_end_index = num_head_arrays * (num_block + 1)
725-
heads = jax.vmap(operation)(x[:head_end_index].reshape(num_head_arrays, -1))
726-
tails = jax.vmap(operation)(x[head_end_index:].reshape(-1, num_block))
727-
outs = jnp.concatenate([heads, tails])
728-
else:
729-
outs = jax.vmap(operation)(x.reshape(-1, num_block))
740+
size = x.shape[0]
741+
# PyTorch adaptive-pooling bins: start = floor(i * size / T),
742+
# end = ceil((i + 1) * size / T) computed with integer arithmetic.
743+
outs = jnp.stack([
744+
operation(x[(i * size) // target_size: -((-((i + 1) * size)) // target_size)])
745+
for i in range(target_size)
746+
])
730747
return outs
731748

732749

brainpy/dnn/pooling_layers_test.py

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -274,5 +274,68 @@ def test_pool_leftmost_negative_channel_axis(self):
274274
self.assertEqual(out.shape, (6, 4, 4))
275275

276276

277+
def _adaptive_pool1d_reference(x, target_size, op):
278+
"""Reference PyTorch-style adaptive pooling of a 1-D array (numpy)."""
279+
x = np.asarray(x)
280+
size = x.shape[0]
281+
out = []
282+
for i in range(target_size):
283+
start = (i * size) // target_size
284+
end = -((-((i + 1) * size)) // target_size) # ceil((i + 1) * size / target_size)
285+
out.append(op(x[start:end]))
286+
return np.array(out)
287+
288+
289+
class TestAdaptivePool1d(parameterized.TestCase):
290+
"""Regression coverage for ``_adaptive_pool1d``.
291+
292+
Guards the fix for the ``ZeroDivisionError: integer modulo by zero`` that arose
293+
when ``target_size > size`` (a spatial dimension smaller than its target), which
294+
made the old block-reshape implementation build ``reshape(-1, 0)``.
295+
"""
296+
297+
@parameterized.product(
298+
size_target=((100, 6), (100, 7), (100, 10), (32, 4), (5, 5),
299+
(2, 6), (1, 4), (3, 8)),
300+
op=('mean', 'max'),
301+
)
302+
def test_matches_pytorch_formula(self, size_target, op):
303+
from brainpy.dnn.pooling import _adaptive_pool1d
304+
size, target = size_target
305+
jop, nop = (jnp.mean, np.mean) if op == 'mean' else (jnp.max, np.max)
306+
x = np.arange(size, dtype=np.float32) * 0.5 - 3.0
307+
got = np.asarray(_adaptive_pool1d(bm.as_jax(x), target, jop))
308+
expected = _adaptive_pool1d_reference(x, target, nop)
309+
self.assertEqual(got.shape, (target,))
310+
np.testing.assert_allclose(got, expected, atol=1e-5)
311+
312+
def test_upsampling_repeats_elements(self):
313+
# target_size (6) > size (2): the previously-crashing case. PyTorch adaptive
314+
# max pooling repeats each element across its bins.
315+
from brainpy.dnn.pooling import _adaptive_pool1d
316+
x = jnp.asarray([10.0, 20.0])
317+
out = np.asarray(_adaptive_pool1d(x, 6, jnp.max))
318+
np.testing.assert_array_equal(out, [10., 10., 10., 20., 20., 20.])
319+
320+
def test_rejects_nonpositive_target(self):
321+
from brainpy.dnn.pooling import _adaptive_pool1d
322+
with self.assertRaises(ValueError):
323+
_adaptive_pool1d(jnp.arange(4.0), 0, jnp.mean)
324+
with self.assertRaises(ValueError):
325+
_adaptive_pool1d(jnp.arange(4.0), -2, jnp.mean)
326+
327+
@parameterized.product(axis=(-1, 0, 1, 2, 3))
328+
def test_adaptivemaxpool3d_spatial_dim_smaller_than_target(self, axis):
329+
# A spatial dim of size 2 is pooled to target 6 for every channel_axis that
330+
# does not consume it; this raised ZeroDivisionError before the fix.
331+
bm.random.seed(123)
332+
inp = bm.random.randn(2, 128, 64, 32)
333+
net = bp.dnn.AdaptiveMaxPool3d(target_shape=[6, 5, 4], channel_axis=axis)
334+
out = net(inp)
335+
channel_size = inp.shape[axis]
336+
self.assertEqual(sorted(out.shape), sorted([channel_size, 6, 5, 4]))
337+
self.assertTrue(bool(jnp.all(jnp.isfinite(out))))
338+
339+
277340
if __name__ == '__main__':
278341
absltest.main()

pyproject.toml

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,10 @@ exclude_also = [
146146
# curated set lives in exactly one place. Grow ``files`` (and the override
147147
# ``module`` list) as more modules are fully typed.
148148
[tool.mypy]
149-
python_version = "3.11"
149+
# Target a modern Python for type-checking. numpy's published stubs use PEP 695
150+
# ``type`` aliases, which mypy rejects as a syntax error unless the analysis target
151+
# is Python 3.12+; pinning an older version breaks the check against current numpy.
152+
python_version = "3.14"
150153
ignore_missing_imports = true
151154
follow_imports = "silent"
152155
warn_unused_ignores = true

0 commit comments

Comments
 (0)