Hi. Just one more thing. Even if qdd has shape n_sampesxn_inputs(neural network) does the gradient of 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? I assume in this case there is an independency but in the 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? Then of course since there are so many little differences the convergence will be different. I’m asking this because despite having everything the same I don’t get the same results as in jax. At the end of the first epochs the results are pretty similar, however after some epochs the results diverge, even though in pytorch I get convergence faster. So there might be some error propagation due to the non associativity of float numbers. But shouldn’t it be presented in jax too? I mean those very small values below 10⁻5/10-6 are noise and this noise is different according to different operations or in this case between frameworks? Also, in the JAX code they don’t perform the sum of L. However using a for loop to calculate the jacobian for each output of the neural network for each sample in the mini batch would be ridiculous computationally expensive in pytorch, taht’s why I’m using the sum since all outputs of the neural netowrk (L1,L2,L3…) are independent of each other since they depend on different samples.