tfp.substrates.jax.math.cholesky_update

Returns cholesky of chol @ chol.T + multiplier * u @ u.T.

tfp.substrates.jax.math.cholesky_update(
    chol, update_vector, multiplier=1.0, name=None
)

Given a (batch of) lower triangular cholesky factor(s) chol, along with a (batch of) vector(s) update_vector, compute the lower triangular cholesky factor of the rank-1 update chol @ chol.T + multiplier * u @ u.T, where multiplier is a (batch of) scalar(s).

If chol has shape [L, L], this has complexity O(L^2) compared to the naive algorithm which has complexity O(L^3).

chol Floating-point Tensor with shape [B1, ..., Bn, L, L]. Cholesky decomposition of mat = chol @ chol.T. Batch dimensions must be broadcastable with update_vector and multiplier.
update_vector Floating-point Tensor with shape [B1, ... Bn, L]. Vector defining rank-one update. Batch dimensions must be broadcastable with chol and multiplier.
multiplier Floating-point Tensor with shape [B1, ..., Bn]. Scalar multiplier to rank-one update. Batch dimensions must be broadcastable withcholandupdate_vector. Note that updates wheremultiplieris positive are numerically stable, while whenmultiplieris negative (downdating), the update will only work if the new resulting matrix is still positive definite. </td> </tr><tr> <td>name` Optional name for this op.

References

[1] Oswin Krause. Christian Igel. A More Efficient Rank-one Covariance Matrix Update for Evolution Strategies. 2015 ACM Conference. https://www.researchgate.net/publication/300581419_A_More_Efficient_Rank-one_Covariance_Matrix_Update_for_Evolution_Strategies