Feature #88: AdaptiveAvgPool2D Implementation

[0xrushi]

Hello all, This is my first PR here please guide me if something is missing.
Issue link: #88

I'm trying to add an adaptiveavg pooling which in pytorch works like below,

import numpy as np
from torch import nn
import torch
input = torch.from_numpy(np.array([[ 0.36900425, -0.46067554, -0.86509347],
             [ 1.2080882 ,  0.59699154, -0.87080586],
             [-0.3984998 , -0.6670093 ,  0.33689347]]))
input = input.reshape(1,3, 3)
m = nn.AdaptiveAvgPool2d((2, 2))
output = m(input)
output

I see in haiku/_src/pool_test.py there are unbatched tests like test_avg_pool_unbatched, will this also have an unbatched test case because I think this should only work for 2D images.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

For more information, open the CLA check for this pull request.