//  [python]
def test_sde_time_diffusion_score(Diffusion, predict_gaussian_noise):
    np.random.seed(132)
    torch.manual_seed(123)

    N = 10_000
    T = 10
    t1 = 1

    x = torch.randn(1, T, 1).repeat(N, 1, 1).square().to(device) + 1
    t = torch.linspace(0, 1, T).view(1, -1, 1).repeat(N, 1, 1).to(device)

    class SDE(torch.nn.Module):
        sde_type = 'ito'
        noise_type = 'general'

        def __init__(self, beta_fn, cov):
            super().__init__()
            self.beta_fn = beta_fn
            self.L = torch.linalg.cholesky(cov)

        def f(self, t, x):
            return -0.5 * self.beta_fn(t) * x

        def g(self, t, x):
            return self.beta_fn(t).sqrt() * self.L


    for ratio in [0.01, 0.3, 0.9]:
        i = torch.Tensor([ratio * t1]).view(1, 1, 1).repeat(N, T, 1).to(device)

        beta_fn = BetaLinear(0.1, 10)
        sdediff = Diffusion(
            dim=1,
            t1=t1,
            beta_fn=beta_fn,
            predict_gaussian_noise=predict_gaussian_noise,
        ).to(device)

        y, diff_noise, diff_mean, diff_std, diff_cov = sdediff(x, t=t, i=i, _return_all=True)
        L = torch.linalg.cholesky(diff_cov[0])

        y = y.requires_grad_(True)

        assert len(diff_mean.unique()) == 10
        assert torch.all(diff_cov[0] == diff_cov[1])

        diff_mean = diff_mean[0,:,0]
        diff_cov_tilde = (diff_cov * diff_std**2)[0]

        def model(*args, noise=None, **kwargs):
            # "Fake" model that perfectly predicts the noise
            # In case `predict_gaussian_noise=True`, undo covariance
            if predict_gaussian_noise:
                noise = torch.linalg.inv(L) @ noise
            return noise

        model_score = sdediff._get_score(model, x, i=i, t=t, L=L, noise=diff_noise)

        # Statistics from diffusion function vs. empirical covariance of diffused values
        empirical_diff_mean = y.mean([0, 2])
        empirical_diff_cov = torch.cov(y.squeeze(-1).T)
        assert torch.allclose(empirical_diff_mean, diff_mean, atol=0.05)
        assert torch.allclose(empirical_diff_cov, diff_cov_tilde, atol=0.05)

        # Empirical score vs. score from diffusion function
        diff_dist = torch.distributions.MultivariateNormal(diff_mean, diff_cov_tilde)
        log_prob = diff_dist.log_prob(y.squeeze(-1))
        empirical_diff_score = torch.autograd.grad(log_prob.sum(), y)[0]
        assert torch.allclose(empirical_diff_score, model_score, atol=0.05, rtol=0.05)

        # True theoretical covariance vs. calculated in diffusion function
        time_cov = sdediff.noise.covariance(t)
        assert torch.allclose(time_cov[0], diff_cov, atol=0.05)
        true_cov = (1 - torch.exp(-beta_fn.integral(i[0][0]))) * time_cov[0]
        assert torch.allclose(true_cov, diff_cov_tilde, atol=0.05)

        # SDE sequential computation vs. direct covariance computation
        sde = SDE(beta_fn=beta_fn, cov=time_cov)
        times = torch.Tensor([0, ratio * t1]).to(y)
        with torch.no_grad():
            true_y = torchsde.sdeint(sde, x.squeeze(-1), times, dt=5e-4)[-1]
        sde_cov = torch.cov(true_y.T)
        assert torch.allclose(true_cov, sde_cov, atol=0.05)