stackMatricesTo3D function

Tensor<Tensor3D> stackMatricesTo3D(List<Tensor<Matrix>> matrices)

Implementation

Tensor<Tensor3D> stackMatricesTo3D(List<Tensor<Matrix>> matrices) {
  int depth = matrices.length;
  int height = matrices[0].shape[0];
  int width = matrices[0].shape[1];

  Tensor3D outValue = [];
  for (int d = 0; d < depth; d = d + 1) {
    outValue.add(matrices[d].value);
  }

  Tensor<Tensor3D> out = Tensor<Tensor3D>(outValue);

  out.creator = Node(
    matrices,
        () {
      for (int d = 0; d < depth; d = d + 1) {
        Tensor<Matrix> m = matrices[d];
        int offset = d * height * width;
        for (int i = 0; i < height * width; i = i + 1) {
          m.grad[i] = m.grad[i] + out.grad[offset + i];
        }
      }
    },
    opName: 'stack_to_3d',
    cost: 0,
  );

  return out;
}