| 52 | } |
| 53 | |
| 54 | void sparseConjugateGradient(void) { |
| 55 | array x = constant(0, b.dims(), f32); |
| 56 | array r = b - matmul(spA, x); |
| 57 | array p = r; |
| 58 | |
| 59 | for (int i = 0; i < maxIter; ++i) { |
| 60 | array Ap = matmul(spA, p); |
| 61 | array alpha_num = dot(r, r); |
| 62 | array alpha_den = dot(p, Ap); |
| 63 | array alpha = alpha_num / alpha_den; |
| 64 | r -= tile(alpha, Ap.dims()) * Ap; |
| 65 | x += tile(alpha, Ap.dims()) * p; |
| 66 | array beta_num = dot(r, r); |
| 67 | array beta = beta_num / alpha_num; |
| 68 | p = r + tile(beta, p.dims()) * p; |
| 69 | } |
| 70 | } |
| 71 | |
| 72 | void denseConjugateGradient(void) { |
| 73 | array x = constant(0, b.dims(), f32); |