matVecMul function

Tensor<Vector> matVecMul(Tensor<Matrix> M, Tensor<Vector> v)

Implementation

Tensor<Vector> matVecMul(Tensor<Matrix> M, Tensor<Vector> v) {
  int numRows = M.shape[0];
  int numCols = M.shape[1];

  Matrix mMat = M.value;
  Vector vVec = v.value;

  Vector outValue = [];
  for (int i = 0; i < numRows; i = i + 1) {
    double sum = 0.0;
    for (int j = 0; j < numCols; j = j + 1) {
      sum = sum + mMat[i][j] * vVec[j];
    }
    outValue.add(sum);
  }

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

  out.creator = Node(
    [M, v],
        () {
      for (int i = 0; i < numRows; i = i + 1) {
        double outGrad = out.grad[i];
        for (int j = 0; j < numCols; j = j + 1) {
          int mIdx = i * numCols + j;
          M.grad[mIdx] = M.grad[mIdx] + outGrad * v.data[j];
          v.grad[j] = v.grad[j] + M.data[mIdx] * outGrad;
        }
      }
    },
    opName: 'matVecMul',
    cost: 2 * numRows * numCols,
  );
  return out;
}