I want to code a custom loss function that uses the Wasserstein distance from scipy:

from scipy.stats import wasserstein_distance

I coded this using tf.numpy_function since I need to pass numpy arrays to the wasserstein_distance function:

@tf.numpy_function(Tout=tf.float32)
def distance(y_true, y_pred):
    dist = np.abs(wasserstein_distance(y_true.ravel(), y_pred.ravel()))
    return dist.astype(np.float32)

@tf.function(input_signature=[tf.TensorSpec(None, tf.float32), tf.TensorSpec(None, tf.float32)])
def wasserstein_loss(y_true, y_pred):
    loss = tf.numpy_function(distance, [y_true, y_pred], tf.float32)
    return loss

When I evaluate distance eagerly on toy data:

y_true = np.random.random((10,1)).astype(np.float32)
y_pred = np.random.random((10,1)).astype(np.float32)

it evaluates just fine:

distance(y_true, y_pred).numpy()

0.070802614

Also, both functions return a tensor:

distance(y_true, y_pred)

<tf.Tensor: shape=(), dtype=float32, numpy=0.070802614>

wasserstein_loss(tf.Variable(y_true, tf.float32), tf.Variable(y_pred, tf.float32))

<tf.Tensor: shape=(), dtype=float32, numpy=0.070802614>

But when I use wasserstein_loss as a loss function inside model.compile() and then try to do model.fit(), I get the error:

ValueError: No gradients provided for any variable

Any help will be appreciated!