Repository: ivy-llc/ivy

Issue:
kl_div



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/torch/nn/functional/loss_functions.py b/ivy/functional/frontends/torch/nn/functional/loss_functions.py
--- a/ivy/functional/frontends/torch/nn/functional/loss_functions.py
+++ b/ivy/functional/frontends/torch/nn/functional/loss_functions.py
@@ -278,37 +278,19 @@ def huber_loss(
 
 
 @to_ivy_arrays_and_back
-@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, "torch")
+@with_supported_dtypes({"2.0.1 and below": ("float32", "float64")}, "torch")
 def kl_div(
     input, target, size_average=None, reduce=None, reduction="mean", log_target=False
 ):
-    size = ivy.shape(input)
-
-    if len(size) < 1:
-        size = [1]
-
-    def loss_fn():
-        if log_target:
-            return ivy.exp(target) * (target - input)
-        return target * (ivy.log(target) - input)
-
-    def batchmean(x):
-        if not reduce:
-            return x / size[0]
-
-        if size_average:
-            return ivy.mean(x) / size[0]
-
-        return ivy.sum(x) / size[0]
-
-    loss = ivy.nan_to_num(loss_fn())
-
-    if reduction == "batchmean":
-        reduction = batchmean
+    orig_red = reduction
+    if size_average is not None or reduce is not None:
+        reduction = _get_reduction_string(size_average, reduce)
     else:
-        reduction = _get_reduction(reduction, size_average, reduce)
-
-    return reduction(loss)
+        reduction = reduction if reduction != "batchmean" else "sum"
+    ret = ivy.kl_div(input, target, reduction=reduction, log_target=log_target)
+    if orig_red == "batchmean" and input.ndim != 0:
+        ret = ret / input.shape[0]
+    return ret
 
 
 @to_ivy_arrays_and_back


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_torch/test_nn/test_functional/test_loss_functions.py b/ivy_tests/test_ivy/test_frontends/test_torch/test_nn/test_functional/test_loss_functions.py
--- a/ivy_tests/test_ivy/test_frontends/test_torch/test_nn/test_functional/test_loss_functions.py
+++ b/ivy_tests/test_ivy/test_frontends/test_torch/test_nn/test_functional/test_loss_functions.py
@@ -1,4 +1,5 @@
 # global
+import numpy as np
 from hypothesis import strategies as st
 
 # local
@@ -450,27 +451,25 @@ def test_torch_huber_loss(
 # kl_div
 @handle_frontend_test(
     fn_tree="torch.nn.functional.kl_div",
-    dtype_and_inputs=helpers.dtype_and_values(
-        available_dtypes=helpers.get_dtypes("float"),
-        allow_inf=False,
-        shared_dtype=True,
-        min_value=0,
-        max_value=10,
-        min_num_dims=0,
-        max_num_dims=10,
-        min_dim_size=0,
-        max_dim_size=10,
-        num_arrays=2,
+    dtype_and_input=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("valid"),
+        min_value=1e-04,
+        max_value=1,
     ),
-    size_average=st.booleans(),
-    reduce=st.booleans(),
-    reduction=st.sampled_from(["none", "mean", "sum", "batchmean"]),
+    dtype_and_target=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("valid"),
+        min_value=1e-04,
+        max_value=1,
+    ),
+    reduction=st.sampled_from(["none", "sum", "batchmean", "mean"]),
     log_target=st.booleans(),
-    test_with_out=st.just(False),
+    size_average=st.one_of(st.just(None), st.booleans()),
+    reduce=st.one_of(st.just(None), st.booleans()),
 )
 def test_torch_kl_div(
     *,
-    dtype_and_inputs,
+    dtype_and_input,
+    dtype_and_target,
     size_average,
     reduce,
     reduction,
@@ -481,16 +480,20 @@ def test_torch_kl_div(
     backend_fw,
     on_device,
 ):
-    inputs_dtype, inputs = dtype_and_inputs
+    input_dtype, input = dtype_and_input
+    input[0] = np.array(np.log(input[0]))
+    target_dtype, target = dtype_and_target
+    if log_target:
+        target[0] = np.array(np.log(target[0]))
     helpers.test_frontend_function(
-        input_dtypes=inputs_dtype,
+        input_dtypes=input_dtype + target_dtype,
         backend_to_test=backend_fw,
         frontend=frontend,
         test_flags=test_flags,
         fn_tree=fn_tree,
         on_device=on_device,
-        input=inputs[0],
-        target=inputs[1],
+        input=input[0],
+        target=target[0],
         size_average=size_average,
         reduce=reduce,
         reduction=reduction,