Returns X such that A * X = B.

Implements the Conjugate Gradient method.

a_mat a Symmetric Positive Definite matrix, represented as a Tensor of shape [n, n].
b_mat a Tensor of shape [n, k].
tol (float) desired tolerance on the residual.

X Tensor X of shape [n, k] such that A * X = B.

ValueError if a_mat is not square or a_mat and b_mat have incompatible shapes.