More robust SGMCMC tests

[SamDuffield]

The error tolerances for the SGMCMC tests are quite large -> this should be reduced and robustified!

[KaelanDt]

could be as simple as checking that the error on the means and covariances decreases when collecting more samples, eg for sgld:

log_posts_1, log_posts_2 = [], []
n_steps_1, n_steps_2 = 1000, 10000

for _ in range(n_steps_1):
        sgld_state = sampler.update(sgld_state, batch, inplace=False)

        all_params_1 = tree_map(
            lambda x, y: torch.cat((x, y.unsqueeze(0))), all_params_1, sgld_state.params
        )

        log_posts_1.append(sgld_state.log_posterior.item())

for _ in range(n_steps_2):
        sgld_state = sampler.update(sgld_state, batch, inplace=False)

        all_params_2 = tree_map(
            lambda x, y: torch.cat((x, y.unsqueeze(0))), all_params_2, sgld_state.params
        )

        log_posts_2.append(sgld_state.log_posterior.item())

    burnin = 1000
    all_params_flat_1 = torch.vmap(lambda x: tree_ravel(x)[0])(all_params_1)
    sampled_mean_1 = all_params_flat_1[burnin:].mean(0)
    sampled_cov_1 = torch.cov(all_params_flat_1[burnin:].T)

    all_params_flat_2 = torch.vmap(lambda x: tree_ravel(x)[0])(all_params_2)
    sampled_mean_2 = all_params_flat_2[burnin:].mean(0)
    sampled_cov_2 = torch.cov(all_params_flat_2[burnin:].T)

    assert log_posts_2[-1] > log_posts_1[-1] > log_posts_1[0]
    assert torch.allclose(sampled_mean_1, target_mean_flat, atol=1e-0, rtol=1e-1)
    assert torch.allclose(sampled_cov_1, target_cov, atol=1e-0, rtol=1e-1)
    assert torch.linalg.norm(abs(sampled_mean_2 - target_mean_flat)) < torch.linalg.norm(abs(sampled_mean_1 - target_mean_flat))
    assert torch.linalg.norm(abs(sampled_cov_2 - target_cov)) < torch.linalg.norm(abs(sampled_cov_1 - target_cov))