Custom loss function for tensorflow ML

I am a novice to Tensorflow but have significant experience in the scikit-learn environment. I am trying to shift some codes to the more flexible TF/Keras routines in Python. One modification I would like to do is to build a simple neural network that trains on an earth mover’s distance/Wasserstein distance between arrays representing densities, rather than on direct root-mean-square differences in the densities.

I have been able to modify a loss function, but do not know how to implement the EMD in TF. I use pyEMD to calculate the EMD between two of the arrays in Python, but is doesn’t translate to tensors in TF. It seems like I may have to “unstack” the tensors, evaluate by each row, then “restack” as a tensor. It doesn’t seem reasonable to do it this way, but I am willing to make it work. At this point, I am focused on making it work rather than be as-fast-as-possible.

Does anyone have advice on either a TF-friendly EMD approach or how I may use pyEMD in a custom loss function? I appreciate any and all suggestions.