Implement pooling using Haiku public API.
Implement pooling using Haiku public API.

PiperOrigin-RevId: 313389107
Change-Id: I78cd460e26bf2cca7dab3d1fc6c2a2c44b625b8e

deepmind/dm-haiku,deepmind/dm-haiku