Thank you so much. The equation I am trying to implement is the Moyal-Wigner equation.
\frac{\partial W(x,p,t)}{\partial t}  =- p* \frac{\partial W(x,p,t)}{\partial x} + \sum_{s=0}^{k}(-h^2)^{s}\frac{1}{(2s + 1!)}(1/2)^{2s} \frac{\partial^{2s+1} U(x)}{\partial x^{2s+1}} \frac{\partial^{2s+1} W(x,p,t)}{\partial p^{2s+1}}. The partial derivative w.r.t to p and t are easy but x is messy. But I can do the derivatives by hand. U(x) is a known function as well. For starters, k=0,1. But my goal is to train a NN to approximate U. So the above equation becomes part of the loss function. So I need to be able to backprop. That’s really where I am stuck.
Thank you so much for your answer