concatenateMatricesByColumnGPU function

GPUTensor<Matrix> concatenateMatricesByColumnGPU(List<GPUTensor<Matrix>> matrices, CommandBuffer tape)

Implementation

GPUTensor<Matrix> concatenateMatricesByColumnGPU(List<GPUTensor<Matrix>> matrices, CommandBuffer tape) {
  GPUTensor<Matrix> result = matrices[0];

  for (int i = 1; i < matrices.length; i = i + 1) {
    GPUTensor<Matrix> a = result;
    GPUTensor<Matrix> b = matrices[i];

    int rows = a.shape[0];
    int colsA = a.shape[1];
    int colsB = b.shape[1];
    int totalCols = colsA + colsB;

    GPUTensor<Matrix> out = GPUTensor<Matrix>.empty([rows, totalCols]);

    tape.putInt(OP_CONCATENATE);
    tape.putString(a.id);
    tape.putString(b.id);
    tape.putString(out.id);
    tape.putInt(1);

    out.creator = GPUNode(
      [a, b],
          (CommandBuffer bTape) {
        bTape.putInt(OP_CONCATENATE_BACKWARD);
        bTape.putString('${out.id}_grad');
        bTape.putString('${a.id}_grad');
        bTape.putString('${b.id}_grad');
        bTape.putInt(1);
        bTape.putInt(colsA);
      },
      opName: 'concat_matrix_colGPU',
      cost: 0,
    );

    result = out;
  }

  return result;
}