Added Wasserstein Loss, the Most Popular Loss Function for GANs

ivy/functional/ivy/losses.py

@@ -441,3 +441,59 @@ def ssim_loss(
     if ivy.exists(out):
         ret = ivy.inplace_update(out, ret)
     return ret
+
+
+@handle_exceptions
+@handle_nestable
+@handle_array_like_without_promotion
+@inputs_to_ivy_arrays
+@handle_array_function
+def wasserstein_loss_discriminator(
+    p_real: Union[ivy.Array, ivy.NativeArray],
+    p_fake: Union[ivy.Array, ivy.NativeArray],
+    out: Optional[ivy.Array] = None,
+) -> ivy.Array:
+    """Compute the Wasserstein loss for the discriminator (critic).
+
+    Parameters
+    ----------
+        p_real (`ivy.Array`): Predictions for real data.
+        p_fake (`ivy.Array`): Predictions for fake data.
+
+    Returns
+    -------
+        `ivy.Array`: Wasserstein loss for the discriminator.
+    """
+    r_loss = ivy.mean(p_real)
+    f_loss = ivy.mean(p_fake)
+    ret = f_loss - r_loss
+
+    if ivy.exists(out):
+        ret = ivy.inplace_update(out, ret)
+    return ret
+
+
+@handle_exceptions
+@handle_nestable
+@handle_array_like_without_promotion
+@inputs_to_ivy_arrays
+@handle_array_function
+def wasserstein_loss_generator(
+    pred_fake: Union[ivy.Array, ivy.NativeArray],
+    out: Optional[ivy.Array] = None,
+) -> ivy.Array:
+    """Compute the Wasserstein loss for the generator.
+
+    Parameters
+    ----------
+        pred_fake (ivy.Array): Predictions for fake data.
+
+    Returns
+    -------
+        ivy.Array: Wasserstein loss for the generator.
+    """
+    ret = -1 * ivy.mean(pred_fake)
+
+    if ivy.exists(out):
+        ret = ivy.inplace_update(out, ret)
+    return ret

ivy_tests/test_ivy/test_functional/test_nn/test_losses.py

@@ -241,3 +241,72 @@ def test_ssim_loss(
         rtol_=1e-02,
         atol_=1e-02,
     )
+
+
+# wasserstein_loss_discriminator
+@handle_test(
+    fn_tree="functional.ivy.wasserstein_loss_discriminator",
+    dtype_and_p_real=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("float"),
+        min_value=-1,
+        max_value=1,
+        allow_inf=False,
+        min_num_dims=1,
+        max_num_dims=1,
+        min_dim_size=2,
+    ),
+    dtype_and_p_fake=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("float"),
+        min_value=-1,
+        max_value=1,
+        allow_inf=False,
+        min_num_dims=1,
+        max_num_dims=1,
+        min_dim_size=2,
+    ),
+)
+def test_wasserstein_loss_discriminator(
+    dtype_and_p_real, dtype_and_p_fake, test_flags, backend_fw, fn_name, on_device
+):
+    dtype_p_real, p_real = dtype_and_p_real
+    dtype_p_fake, p_fake = dtype_and_p_fake
+
+    helpers.test_function(
+        input_dtypes=dtype_p_real + dtype_p_fake,
+        test_flags=test_flags,
+        backend_to_test=backend_fw,
+        fn_name=fn_name,
+        on_device=on_device,
+        p_real=p_real[0],
+        p_fake=p_fake[0],
+        rtol_=1e-02,
+        atol_=1e-02,
+    )
+
+
+# wasserstein_loss_generator
+@handle_test(
+    fn_tree="functional.ivy.wasserstein_loss_generator",
+    dtype_and_pred_fake=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("float"),
+        min_value=-1,
+        max_value=1,
+        allow_inf=False,
+        min_num_dims=1,
+        max_num_dims=1,
+        min_dim_size=2,
+    ),
+)
+def test_wasserstein_loss_generator(dtype_and_pred_fake, test_flags, backend_fw, fn_name, on_device):
+    dtype_pred_fake, pred_fake = dtype_and_pred_fake
+
+    helpers.test_function(
+        input_dtypes=dtype_pred_fake,
+        test_flags=test_flags,
+        backend_to_test=backend_fw,
+        fn_name=fn_name,
+        on_device=on_device,
+        pred_fake=pred_fake[0],
+        rtol_=1e-02,
+        atol_=1e-02,
+    )