| 70 | } |
| 71 | |
| 72 | void denseConjugateGradient(void) { |
| 73 | array x = constant(0, b.dims(), f32); |
| 74 | array r = b - matmul(A, x); |
| 75 | array p = r; |
| 76 | |
| 77 | for (int i = 0; i < maxIter; ++i) { |
| 78 | array Ap = matmul(A, p); |
| 79 | array alpha_num = dot(r, r); |
| 80 | array alpha_den = dot(p, Ap); |
| 81 | array alpha = alpha_num / alpha_den; |
| 82 | r -= tile(alpha, Ap.dims()) * Ap; |
| 83 | x += tile(alpha, Ap.dims()) * p; |
| 84 | array beta_num = dot(r, r); |
| 85 | array beta = beta_num / alpha_num; |
| 86 | p = r + tile(beta, p.dims()) * p; |
| 87 | } |
| 88 | } |
| 89 | |
| 90 | void checkConjugateGradient(const af::array in) { |
| 91 | array x = constant(0, b.dims(), f32); |