saul_santos:

the sum of L which has shape n_samples gives the same gradient of a for loop over L0/qqd_sample0, L1/qqd_sample1 and so on?


You need the independence of the computation for this to work. If you have it it will be fine.

MSE loss function the differences can be of order of magnitude of 10^0. I suppose the results diverge even more when n_samples(mini batch) increases. Is this due to error propagation?

Note that the 1e-6 error happen in one op. But then when you keep doing computation, this error is usually amplified by every operation that you perform. So for a deep network or after few iterations, this error can grow quite a lot.

At the end of the first epochs the results are pretty similar, however after some epochs the results diverge

This is totally expected when you use neural nets: they will end up converging to different places in the space of parameters.
But most neural network loss functions are well behaved enough that all of these optima that are found give similar loss value (even though you can have very different parameter values).
You can see this as having the same impact as sampling a different set of initial weights.

But shouldn’t it be presented in jax too?

This is present in jax as well. So you will have the same behavior of very small differences in one op that grow over time.