Solving minimization problem within custom loss function that supports tf gradient / jacobian for convergence

[cyan-at]

Hi, we are using tf v1 and trying to find Wasserstein distance in a custom loss function for a boundary condition (loss = sqrt(Wasserstein distance))

Finding wasserstein distance requires solving a minimization problem (minimize c @ x s.t. Ax = b, x >= 0, x < inf)

We can get the code to run using scipy.linprog (using tf.function to convert y_pred to numpy), but training with this loss function does not converge. (loss bounces around, does not go below 3)

We suspect this is because tf see our loss function as a black box and has no way to improve on an epoch. We are wondering if there is some way / tool to express our minimization problem so that deepxde / tf can watch for gradient / jacobian and better converge?

Thank you

[lululxvi]

I don't think you can use scipy with tensorflow/pytorch.