flax

549b0592 - Add support for single example pooling