[Question] odd behaviour when solving SDE for differential initial conditions

Hello,
I am trying to use diffrax as an sde solver to reverse a diffusion process in the context of diffusion models for machine learning.
In particular, given the SDE of the forward diffusion process, I want to solve the reverse SDE to be able to sample new data.
In every implementation I've seen it seems that they are reinventing the wheel and reimplementing an SDE solver, which I would like to avoid.

Coming to my problem, when I implement "manually" the Euler algorithm, I obtain reasonable results (I obtain samples close to every point of the training dataset with the same probability). On the other hand, when using diffrax to solve the reverse SDE, not every solution seems equiprobable, which is very odd, and problematic for this application.

I attach a colab notebook with a minimal example of the problem for a very simple dataset:

https://colab.research.google.com/drive/1V1nU3vn9hkZvWcJnaYYWOwnsOvBB9iSd?usp=sharing

In this case, I don't implement the NN which would approximate the score, and compute the score directly for the dataset, as it is very simple, to test whether the SDE solver behaves as expected.

If you could help me figure out this problem I would really appreciate it. I am not too familiar with both diffusion models and this library, so I don't have a good intuition, but could it be a problem of the rng somehow?

Thank you very much for your help and time!

Start with the linked Colab notebook and reproduce the minimal reverse-SDE example, comparing its manual Euler implementation with the diffrax solver. Inspect the solver configuration and random-number handling described there, then determine whether the unequal sampling is a reproducible solver issue or an expected difference. Done means the cause is identified and the example either produces the expected probabilities or documents the limitation.