My initial attempt at extending the jax.experimental.stax library with a Spatial Pyramid Pooling feature, based on my rough understanding of this paper by Kaiming He, Xiangyu Zhang, Shaoqing Ren and Jian Sun.
There are most certainly errors and bugs, but if you let me know or send me a pull request, I'll try to fix it/use your suggestions!