W2 gradient

[tmasthay]

Wasserstein gradient gives nonsensical results still.

It is quite possible that what is happening is that I am not passing the right indices of the spline coefficient evaluation onto the correct rank.

For example say that I have 100 splines and two ranks.

On rank 0, I have loop splines[:50] and obs_data[:50] as I should. However, on rank 1, I have splines[:50] but obs_data[50:]! This would obviously lead to completely nonsensical results, which is what we observe.

[tmasthay]

The original description above seems to check out, at least partially. When only one rank is used, the results are still weird, but they are more sensible. random_traces.gif shows that some traces seem to converge in the eyeball norm, but other traces give totally nonsensical predictions with totally disjoint support.

vp.gif gives some regions of some agreement and again with the weird upswinging artifacts.

[tmasthay]

[tmasthay]

Consider using torchinterp1d instead. It has a simpler interface which makes it slightly less bug prone. But more importantly, there might be something weird with multiple ranks for torchcubicspline. Finally, while the splines work fine for the interpolation, they are overkill; linear interpolation should do just fine for what we are doing here.