| 88 | } |
| 89 | |
| 90 | void checkConjugateGradient(const af::array in) { |
| 91 | array x = constant(0, b.dims(), f32); |
| 92 | array r = b - matmul(in, x); |
| 93 | array p = r; |
| 94 | |
| 95 | for (int i = 0; i < maxIter; ++i) { |
| 96 | array Ap = matmul(in, p); |
| 97 | array alpha_num = dot(r, r); |
| 98 | array alpha_den = dot(p, Ap); |
| 99 | array alpha = alpha_num / alpha_den; |
| 100 | r -= tile(alpha, Ap.dims()) * Ap; |
| 101 | x += tile(alpha, Ap.dims()) * p; |
| 102 | array beta_num = dot(r, r); |
| 103 | array beta = beta_num / alpha_num; |
| 104 | p = r + tile(beta, p.dims()) * p; |
| 105 | } |
| 106 | array res = x0 - x; |
| 107 | |
| 108 | std::cout << "Final difference in solutions:\n"; |
| 109 | af_print(dot(res, res)); |
| 110 | } |
| 111 | |
| 112 | int main(int, char **) { |
| 113 | af::info(); |