concatenateMatricesByColumnGPU function

GPUTensor<Matrix> concatenateMatricesByColumnGPU(
  1. List<GPUTensor<Matrix>> matrices,
  2. CommandBuffer tape, {
  3. GPUTensor<Matrix>? outTensor,
})

Appends a list of matrices matrices together in column order. The operation is appended to tape for execution.

Implementation

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

  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;
    bool isLast = i == matrices.length - 1;

    if (isLast == true && outTensor != null) {
      out = outTensor;
    } else {
      out = GPUTensor<Matrix>.empty(<int>[rows, totalCols]);
      if (isLast == false) {
        intermediates.add(out);
      }
    }

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

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

    result = out;
  }

  for (int i = 0; i < intermediates.length; i = i + 1) {
    result.subMap['intermediate_$i'] = intermediates[i];
  }

  return result;
}