def test_cholesky_update(chol_batch_shape, vec_batch_shape, dim, coef):
A = random.normal(random.PRNGKey(0), chol_batch_shape + (dim, dim))
A = A @ jnp.swapaxes(A, -2, -1) + jnp.eye(dim)
x = random.normal(random.PRNGKey(0), vec_batch_shape + (dim,)) * 0.1
xxt = x[..., None] @ x[..., None, :]
expected = jnp.linalg.cholesky(A + coef * xxt)
actual = cholesky_update(jnp.linalg.cholesky(A), x, coef)
assert_allclose(actual, expected, atol=1e-4, rtol=1e-4)
def _get_proposal_loc_and_scale(samples, loc, scale, new_sample):
# get loc/scale of q_{-n} (Algorithm 1, line 5 of ref [1]) for n from 1 -> N
# these loc/scale will be stacked to the first dim; so
# proposal_loc.shape[0] = proposal_loc.shape[0] = N
# Here, we use the numerical stability procedure in Appendix 6 of [1].
weight = 1 / samples.shape[0]
if scale.ndim > loc.ndim:
new_scale = cholesky_update(scale, new_sample - loc, weight)
proposal_scale = cholesky_update(new_scale, samples - loc, -weight)
proposal_scale = cholesky_update(proposal_scale, new_sample - samples, - (weight ** 2))
else:
var = jnp.square(scale) + weight * jnp.square(new_sample - loc)
proposal_var = var - weight * jnp.square(samples - loc)
proposal_var = proposal_var - weight ** 2 * jnp.square(new_sample - samples)
proposal_scale = jnp.sqrt(proposal_var)
proposal_loc = loc + weight * (new_sample - samples)
return proposal_loc, proposal_scale