Accumulate low-precision average pooling in float32

[sylvesterkaczmarek]

Summary

Accumulate float16/bfloat16 window sums and SAME-padding counts in float32, then return the original dtype after division. This prevents premature half-precision sum overflow when the pooled mean is representable and improves bfloat16 averaging.

Float32/float64 inputs, max pooling, shape inference and padding semantics are unchanged. The change is separate from the shape-handling proposal in #824. Float32 accumulator overflow and accelerator performance are outside this correction.

Reproduction

import jax.numpy as jnp
import haiku as hk
x = jnp.full((1, 16, 16, 1), 256., dtype=jnp.float16)
print(hk.avg_pool(x, (1, 16, 16, 1), 1, 'VALID'))

Unchanged main returns infinity instead of 256.

Validation

python -m pytest -q haiku/_src/pool_test.py haiku/_src/pool_precision_test.py

All 20 existing pooling tests and 16 new cases pass. Coverage includes independent NumPy window means, positive/negative constant activations, nonuniform values, both layouts and padding modes, float16/bfloat16/float32, eager/JIT execution, vmap, and a real AvgPool module with a finite loss and gradients. Eight new tests fail on unchanged upstream and eight controls pass.

Eight existing shape-inference deprecation warnings were emitted. No trained model, image dataset or GPU/TPU execution was used; no performance claim is made.

Tested locally on macOS CPU using real module imports. New test formatting, scoped static checks, syntax checks and git diff --check pass. No dependency or workflow changes.