maxPool1d function

Tensor<Vector> maxPool1d(Tensor<Vector> input, int poolSize, int stride)

Implementation

Tensor<Vector> maxPool1d(Tensor<Vector> input, int poolSize, int stride) {
  int inputSize = input.shape[0];
  int outputSize = (inputSize - poolSize) ~/ stride + 1;

  Vector outputValue = [];
  List<int> maxIndices = [];

  for (int i = 0; i < outputSize; i = i + 1) {
    double maxVal = -double.infinity;
    int maxIdx = -1;
    for (int p = 0; p < poolSize; p = p + 1) {
      int inIdx = i * stride + p;
      double val = input.data[inIdx];
      if (val > maxVal) {
        maxVal = val;
        maxIdx = inIdx;
      }
    }
    outputValue.add(maxVal);
    maxIndices.add(maxIdx);
  }

  Tensor<Vector> out = Tensor<Vector>(outputValue);

  out.creator = Node(
    [input],
        () {
      int outLength = out.data.length;
      for (int i = 0; i < outLength; i = i + 1) {
        int mIdx = maxIndices[i];
        input.grad[mIdx] = input.grad[mIdx] + out.grad[i];
      }
    },
    opName: 'max_pool_1d',
    cost: outputSize * poolSize,
  );

  return out;
}