stackMatricesGPU function

GPUTensor<Tensor3D> stackMatricesGPU(
  1. List<GPUTensor<Matrix>> matrices,
  2. CommandBuffer tape
)

Implementation

GPUTensor<Tensor3D> stackMatricesGPU(List<GPUTensor<Matrix>> matrices, CommandBuffer tape) {
  int count = matrices.length;
  int rows = matrices[0].shape[0];
  int cols = matrices[0].shape[1];

  List<int> shape = <int>[count, rows, cols];
  GPUTensor<Tensor3D> out = GPUTensor<Tensor3D>.empty(shape);

  tape.putInt(OP_STACK_ROWS);
  tape.putInt(count);
  for (int i = 0; i < count; i = i + 1) {
    tape.putString(matrices[i].id);
  }
  tape.putString(out.id);
  tape.putInt(0);

  out.creator = GPUNode(
    <GPUTensor>[...matrices],
        (CommandBuffer bTape) {
      bTape.putInt(OP_STACK_ROWS_BACKWARD);
      bTape.putString('${out.id}_grad');
      bTape.putInt(count);
      for (int i = 0; i < count; i = i + 1) {
        bTape.putString('${matrices[i].id}_grad');
      }
      bTape.putInt(0);
    },
    opName: 'stackMatricesGPU',
    cost: count * rows * cols,
  );

  return out;
}