transpose function

Tensor<Matrix> transpose(Tensor<Matrix> a)

Implementation

Tensor<Matrix> transpose(Tensor<Matrix> a) {
  int M = a.shape[0];
  int N = a.shape[1];

  Matrix aMat = a.value;

  Matrix outValue = [];
  for (int i = 0; i < N; i = i + 1) {
    Vector row = [];
    for (int j = 0; j < M; j = j + 1) {
      row.add(aMat[j][i]);
    }
    outValue.add(row);
  }

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

  out.creator = Node(
    [a],
        () {
      for (int i = 0; i < N; i = i + 1) {
        for (int j = 0; j < M; j = j + 1) {
          int aIdx = j * N + i;
          int outIdx = i * M + j;
          a.grad[aIdx] = a.grad[aIdx] + out.grad[outIdx];
        }
      }
    },
    opName: 'transpose',
    cost: 0,
  );
  return out;
}