tf_agents.bandits.policies.linalg.conjugate_gradient

Returns X such that A * X = B.

tf_agents.bandits.policies.linalg.conjugate_gradient(
    a_mat: tf_agents.typing.types.Tensor,
    b_mat: tf_agents.typing.types.Tensor,
    tol: float = 1e-10
) -> tf_agents.typing.types.Float

Implements the Conjugate Gradient method. https://en.wikipedia.org/wiki/Conjugate_gradient_method

Args
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.

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

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