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