Missing MaxPool2d

[KaleabTessera]

I am converting this simple tf example into JAX using Haiku. In process of doing this, I realized there is no MaxPool2d module. Any plans on adding this or is there a way to adapt the current MaxPool module?

More precisely, how would you write this using haiku?

layers.MaxPooling2D((2, 2))

[KaleabTessera]

I guess you could do this:

 hk.MaxPool(window_shape=(1, 2, 2, 1),
                      strides=(1, 2, 2, 1),
                      padding="VALID"),

Not sure if there is a better way though?