Conjugate gradient algorithm

conj.py

import tensorflow as tf
from src.utils.dtypes import tf_float_type


def conj_gradient(A, b, iters):
    # todo can be optimized

    x = tf.zeros_like(b)
    r = b - tf.matmul(A, x)
    p = r

    for _ in range(iters):
        alpha = tf.matmul(tf.transpose(r), r) / (tf.matmul(tf.matmul(tf.transpose(p), A), p) + 1e-7)
        x = x + alpha * p

        r_new = r - alpha * tf.matmul(A, p)

        beta = tf.matmul(tf.transpose(r_new), r_new) / tf.matmul(tf.transpose(r), r)

        p = r_new + beta * p
        r = r_new

    return x


A = tf.constant([[4, 1], [1, 3]], dtype=tf_float_type)
b = tf.transpose(tf.constant([[1, 2]], dtype=tf_float_type))

print(conj_gradient(A, b, iters=10))