flatten weights

Author: mathsinn

art/metrics/metrics.py

@@ -377,6 +377,10 @@ def wasserstein_distance(
     u_values = u_values.flatten().reshape(u_values.shape[0], -1)
     v_values = v_values.flatten().reshape(v_values.shape[0], -1)
 
+    if u_weights is not None and v_weights is not None:
+        u_weights = u_weights.flatten().reshape(u_weights.shape[0], -1)
+        v_weights = v_weights.flatten().reshape(v_weights.shape[0], -1)
+
     wd = np.zeros(u_values.shape[0])
 
     for i in range(u_values.shape[0]):

tests/metrics/test_metrics.py

@@ -373,16 +373,20 @@ def test_wasserstein_distance(self):
 
         x_train = x_train[0:nb_train]
         x_test = x_test[0:nb_test]
+        weights = np.ones_like(x_train)
 
         wd_0 = wasserstein_distance(x_train[:batch_size], x_train[:batch_size])
         wd_1 = wasserstein_distance(x_train[:batch_size], x_test[:batch_size])
+        wd_2 = wasserstein_distance(x_train[:batch_size], x_train[:batch_size], weights[:batch_size],
+                                    weights[:batch_size])
 
         np.testing.assert_array_equal(wd_0, np.asarray([0.0, 0.0, 0.0]))
         np.testing.assert_array_almost_equal(wd_1, np.asarray([0.04564, 0.01235, 0.04787]), decimal=4)
 
         np.testing.assert_array_equal(x_train.shape, np.asarray([nb_train, 28, 28, 1]))
         np.testing.assert_array_equal(x_test.shape, np.asarray([nb_test, 28, 28, 1]))
 
+        np.testing.assert_array_equal(wd_2, np.asarray([0.0, 0.0, 0.0]))
 
 if __name__ == "__main__":
     unittest.main()