padMatrixGPU function

GPUTensor<Matrix> padMatrixGPU(
  1. GPUTensor<Matrix> input,
  2. int padSize,
  3. CommandBuffer tape
)

Implementation

GPUTensor<Matrix> padMatrixGPU(GPUTensor<Matrix> input, int padSize, CommandBuffer tape) {
  int inHeight = input.shape[0];
  int inWidth = input.shape[1];
  int outHeight = inHeight + 2 * padSize;
  int outWidth = inWidth + 2 * padSize;

  GPUTensor<Matrix> out = GPUTensor<Matrix>.empty([outHeight, outWidth]);

  tape.putInt(OP_PAD2D);
  tape.putString(input.id);
  tape.putString(out.id);
  tape.putInt(padSize);
  tape.putInt(padSize);
  tape.putInt(padSize);
  tape.putInt(padSize);

  out.creator = GPUNode(
    [input],
        (CommandBuffer bTape) {
      bTape.putInt(OP_PAD2D_BACKWARD);
      bTape.putString('${out.id}_grad');
      bTape.putString('${input.id}_grad');
      bTape.putInt(padSize);
      bTape.putInt(padSize);
      bTape.putInt(padSize);
      bTape.putInt(padSize);
    },
    opName: 'padMatrixGPU',
    cost: outHeight * outWidth,
  );

  return out;
}