Iainmon / ChAI

A Chapel library for Machine Learning that supports distributed inference, automatic differentiation, and CUDA/HIP utilization.
https://iainmon.github.io/ChAI
4 stars 5 forks source link

[Feature Request] Support rectangular dimensional pooling and stride in `maxPool` #8

Open ShreyasKhandekar opened 1 month ago

ShreyasKhandekar commented 1 month ago

maxPool currently only takes in a single integral parameter which creates a square kernel.

It should be simple enough to implement this to also support passing a two-tuple to support arbitrary kernel sizes.

See https://pytorch.org/docs/stable/generated/torch.nn.MaxPool1d.html#maxpool1d. kernel_size can be int or Tuple[int].

Ex:

    x = x.maxPool(3); // Square 2D pooling

instead of

    x = x.maxPool((3,2));

We should also add the other parameters that pytorch supports like stride, padding, dilation etc.

ShreyasKhandekar commented 3 weeks ago

I didn't add rectangular pooling yet, but added more features to maxPool that helped me like padding and dilation, which unblocked my progress for now.