internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: ivy-llc/ivy

Issue:
binary_cross_entropy_with_logits



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
@@ -19,7 +19,7 @@ def _apply_reduction(reduction, size_average, reduce, to_reduce):
 
 def _get_reduction(reduction, size_average=None, reduce=None):
     if size_average is not None or reduce is not None:
-        return _get_reduction_func(_legacy_get_string(size_average, reduce))
+        return _get_reduction_func(_get_reduction_string(size_average, reduce))
     else:
         return _get_reduction_func(reduction)
 
@@ -66,20 +66,6 @@ def _get_reduction_string(size_average, reduce):
     return ret
 
 
-def _legacy_get_string(size_average, reduce):
-    if size_average is None:
-        size_average = True
-    if reduce is None:
-        reduce = True
-    if size_average and reduce:
-        ret = "mean"
-    elif reduce:
-        ret = "sum"
-    else:
-        ret = "none"
-    return ret
-
-
 # --- Main --- #
 # ------------ #
 
@@ -89,12 +75,11 @@ def _legacy_get_string(size_average, reduce):
 def binary_cross_entropy(
     input, target, weight=None, size_average=None, reduce=None, reduction="mean"
 ):
-    reduction = _get_reduction(reduction, size_average, reduce)
-    result = ivy.binary_cross_entropy(target, input, epsilon=0.0)
+    result = ivy.binary_cross_entropy(target, input, epsilon=0.0, reduction=reduction)
 
     if weight is not None:
         result = ivy.multiply(weight, result)
-    result = reduction(result)
+
     return result
 
 
@@ -111,14 +96,14 @@ def binary_cross_entropy_with_logits(
     result = ivy.binary_cross_entropy(
         target,
         input,
-        reduction="none",
+        reduction=reduction,
         from_logits=True,
         pos_weight=pos_weight,
     )
-    reduction = _get_reduction(reduction, size_average, reduce)
+
     if weight is not None:
         result = ivy.multiply(weight, result)
-    result = reduction(result).astype(target.dtype)
+
     return result
 
 
@@ -212,7 +197,17 @@ def cross_entropy(
     reduction="mean",
     label_smoothing=0.0,
 ):
-    return ivy.cross_entropy(target, input, epsilon=label_smoothing)
+    loss = ivy.cross_entropy(target, input, epsilon=label_smoothing)
+
+    if ignore_index != -100:
+        mask = ivy.not_equal(target, ignore_index)
+        loss = ivy.where(mask, loss, ivy.zeros_like(loss))
+
+    if weight is not None:
+        result = ivy.multiply(weight, loss)
+
+    reduction = _get_reduction(reduction, size_average, reduce)
+    return reduction(result).astype(target.dtype)
 
 
 @to_ivy_arrays_and_back


Test patch:
Repository: ivy-llc/ivy

Issue:
Comparison with a callable in `ivy\functional\ivy\experimental\layers.py`
In the file [`ivy\functional\ivy\experimental\layers.py`](https://github.com/unifyai/ivy/blob/main/ivy/functional/ivy/experimental/layers.py) In the following lines comparisons with a callable were made. I think parenthesis were omitted which might lead to unwanted behaviour.

https://github.com/unifyai/ivy/blob/11058e8d025ce4aefa3528a9a2f8a65b86f8ac73/ivy/functional/ivy/experimental/layers.py#L2535

https://github.com/unifyai/ivy/blob/11058e8d025ce4aefa3528a9a2f8a65b86f8ac73/ivy/functional/ivy/experimental/layers.py#L2544


Discussion:

Resolving patch:
diff --git a/ivy/functional/ivy/experimental/layers.py b/ivy/functional/ivy/experimental/layers.py
--- a/ivy/functional/ivy/experimental/layers.py
+++ b/ivy/functional/ivy/experimental/layers.py
@@ -2532,7 +2532,7 @@ def sliding_window(
             padding=padding,
         )
 
-    if ivy.current_backend_str == "tensorflow":
+    if ivy.current_backend_str() == "tensorflow":
         return ivy.current_backend(input).sliding_window(
             input,
             kernel_size,
@@ -2541,7 +2541,7 @@ def sliding_window(
             padding=padding,
         )
 
-    if ivy.current_backend_str == "paddle":
+    if ivy.current_backend_str() == "paddle":
         return ivy.current_backend(input).sliding_window(
             input,
             kernel_size,


Test patch:
Repository: ivy-llc/ivy

Issue:
add_n
add_n function


Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/paddle/tensor/tensor.py b/ivy/functional/frontends/paddle/tensor/tensor.py
--- a/ivy/functional/frontends/paddle/tensor/tensor.py
+++ b/ivy/functional/frontends/paddle/tensor/tensor.py
@@ -110,6 +110,11 @@ def abs(self):
     def acosh(self, name=None):
         return paddle_frontend.Tensor(ivy.acosh(self._ivy_array))
 
+    @with_unsupported_dtypes({"2.5.1 and below": ("float16", "bfloat16")}, "paddle")
+    def add_n(self, inputs, name=None):
+        inputs = ivy.array(inputs)
+        return ivy.sum(inputs, dtype=inputs.dtype, axis=0)
+
     @with_unsupported_dtypes({"2.5.1 and below": ("float16", "bfloat16")}, "paddle")
     def ceil(self):
         return paddle_frontend.ceil(self)


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py b/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
--- a/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
+++ b/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
@@ -568,6 +568,40 @@ def test_paddle_tensor_add_(
     )
 
 
+@handle_frontend_method(
+    class_tree=CLASS_TREE,
+    init_tree="paddle.to_tensor",
+    method_name="add_n",
+    dtype_and_x=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("float"),
+        num_arrays=helpers.ints(min_value=1, max_value=5),
+        shared_dtype=True,
+    ),
+)
+def test_paddle_tensor_add_n(
+    dtype_and_x,
+    frontend_method_data,
+    init_flags,
+    method_flags,
+    frontend,
+    on_device,
+    backend_fw,
+):
+    input_dtype, x = dtype_and_x
+    helpers.test_frontend_method(
+        init_input_dtypes=input_dtype,
+        backend_to_test=backend_fw,
+        init_all_as_kwargs_np={"inputs": x},
+        method_input_dtypes=input_dtype,
+        method_all_as_kwargs_np={"inputs": x},
+        frontend_method_data=frontend_method_data,
+        init_flags=init_flags,
+        method_flags=method_flags,
+        frontend=frontend,
+        on_device=on_device,
+    )
+
+
 # all
 @handle_frontend_method(
     class_tree=CLASS_TREE,
Repository: ivy-llc/ivy

Issue:
less_equal



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/paddle/tensor/tensor.py b/ivy/functional/frontends/paddle/tensor/tensor.py
--- a/ivy/functional/frontends/paddle/tensor/tensor.py
+++ b/ivy/functional/frontends/paddle/tensor/tensor.py
@@ -799,3 +799,9 @@ def inner(self, y, name=None):
     @with_supported_dtypes({"2.5.1 and below": ("float32", "float64")}, "paddle")
     def mean(self, axis=None, keepdim=False, name=None):
         return paddle_frontend.mean(self._ivy_array, axis=axis, keepdim=keepdim)
+
+    @with_supported_dtypes(
+        {"2.5.1 and below": ("float32", "float64", "int32", "int64")}, "paddle"
+    )
+    def less_equal(self, y, name=None):
+        return paddle_frontend.less_equal(self._ivy_array, y)


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py b/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
--- a/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
+++ b/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
@@ -2397,6 +2397,45 @@ def test_paddle_tensor_lerp_(
     )
 
 
+# less_equal
+@handle_frontend_method(
+    class_tree=CLASS_TREE,
+    init_tree="paddle.to_tensor",
+    method_name="less_equal",
+    dtype_and_x=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("valid"),
+        num_arrays=2,
+        allow_inf=False,
+        large_abs_safety_factor=2,
+        small_abs_safety_factor=2,
+        safety_factor_scale="log",
+        shared_dtype=True,
+    ),
+)
+def test_paddle_tensor_less_equal(
+    dtype_and_x,
+    frontend_method_data,
+    init_flags,
+    method_flags,
+    frontend,
+    on_device,
+    backend_fw,
+):
+    input_dtype, x = dtype_and_x
+    helpers.test_frontend_method(
+        init_input_dtypes=input_dtype,
+        backend_to_test=backend_fw,
+        init_all_as_kwargs_np={"data": x[0]},
+        method_input_dtypes=input_dtype,
+        method_all_as_kwargs_np={"y": x[1]},
+        frontend_method_data=frontend_method_data,
+        init_flags=init_flags,
+        method_flags=method_flags,
+        frontend=frontend,
+        on_device=on_device,
+    )
+
+
 #  less_than
 @handle_frontend_method(
     class_tree=CLASS_TREE,
Repository: ivy-llc/ivy

Issue:
acos



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/paddle/tensor/tensor.py b/ivy/functional/frontends/paddle/tensor/tensor.py
--- a/ivy/functional/frontends/paddle/tensor/tensor.py
+++ b/ivy/functional/frontends/paddle/tensor/tensor.py
@@ -957,6 +957,10 @@ def nonzero(self):
     def inner(self, y, name=None):
         return paddle_frontend.inner(self, y, name)
 
+    @with_supported_dtypes({"2.5.2 and below": ("float32", "float64")}, "paddle")
+    def acos(self, name=None):
+        return paddle_frontend.Tensor(ivy.acos(self._ivy_array))
+
     @with_supported_dtypes({"2.5.2 and below": ("float32", "float64")}, "paddle")
     def mean(self, axis=None, keepdim=False, name=None):
         return paddle_frontend.mean(self, axis=axis, keepdim=keepdim)


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py b/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
--- a/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
+++ b/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
@@ -1144,6 +1144,40 @@ def test_paddle_abs(
     )
 
 
+@handle_frontend_method(
+    class_tree=CLASS_TREE,
+    init_tree="paddle.to_tensor",
+    method_name="acos",
+    dtype_and_x=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("float"),
+    ),
+)
+def test_paddle_tensor_acos(
+    dtype_and_x,
+    frontend_method_data,
+    init_flags,
+    method_flags,
+    frontend,
+    on_device,
+    backend_fw,
+):
+    input_dtype, x = dtype_and_x
+    helpers.test_frontend_method(
+        init_input_dtypes=input_dtype,
+        backend_to_test=backend_fw,
+        init_all_as_kwargs_np={
+            "data": x[0],
+        },
+        method_input_dtypes=input_dtype,
+        method_all_as_kwargs_np={},
+        frontend_method_data=frontend_method_data,
+        init_flags=init_flags,
+        method_flags=method_flags,
+        frontend=frontend,
+        on_device=on_device,
+    )
+
+
 # acosh
 @handle_frontend_method(
     class_tree=CLASS_TREE,
Repository: ivy-llc/ivy

Issue:
as_complex



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/paddle/tensor/tensor.py b/ivy/functional/frontends/paddle/tensor/tensor.py
--- a/ivy/functional/frontends/paddle/tensor/tensor.py
+++ b/ivy/functional/frontends/paddle/tensor/tensor.py
@@ -772,6 +772,20 @@ def inner(self, y, name=None):
     def mean(self, axis=None, keepdim=False, name=None):
         return paddle_frontend.mean(self, axis=axis, keepdim=keepdim)
 
+    @with_supported_dtypes({"2.5.1 and below": ("float32", "float64")}, "paddle")
+    def as_complex(self, name=None):
+        if self.ivy_array.shape[-1] != 2:
+            raise ivy.exceptions.IvyError(
+                "The size of the last dimension of tensor does not equals 2"
+            )
+        dtype = (
+            ivy.complex64 if ivy.dtype(self.ivy_array) == "float32" else ivy.complex128
+        )
+        re_part = self.ivy_array[..., 0]
+        im_part = ivy.multiply(1j, self.ivy_array[..., 1])
+        value = paddle_frontend.Tensor(ivy.add(re_part, im_part).astype(dtype))
+        return value
+
     @with_supported_dtypes(
         {"2.5.1 and below": ("float32", "float64", "int32", "int64")}, "paddle"
     )


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py b/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
--- a/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
+++ b/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
@@ -36,6 +36,26 @@ def _filter_query(query):
     )
 
 
+# as_complex
+@st.composite
+def _get_as_complex_inputs_(draw):
+    shape = draw(
+        helpers.get_shape(
+            min_num_dims=2, max_num_dims=5, min_dim_size=2, max_dim_size=10
+        )
+    )
+
+    x_dtype, x = draw(
+        helpers.dtype_and_values(
+            available_dtypes=helpers.get_dtypes("valid"),
+            shape=(*shape, 2),
+            min_value=0,
+            max_value=50,
+        )
+    )
+    return x_dtype, x
+
+
 # clip
 @st.composite
 def _get_clip_inputs(draw):
@@ -762,6 +782,37 @@ def test_paddle_tensor_argsort(
     )
 
 
+# as_complex
+@handle_frontend_method(
+    class_tree=CLASS_TREE,
+    init_tree="paddle.to_tensor",
+    method_name="as_complex",
+    dtypes_and_x=_get_as_complex_inputs_(),
+)
+def test_paddle_tensor_as_complex(
+    dtypes_and_x,
+    frontend_method_data,
+    init_flags,
+    method_flags,
+    frontend,
+    on_device,
+    backend_fw,
+):
+    input_dtype, x = dtypes_and_x
+    helpers.test_frontend_method(
+        init_input_dtypes=input_dtype,
+        backend_to_test=backend_fw,
+        init_all_as_kwargs_np={"data": x[0]},
+        method_input_dtypes=input_dtype,
+        method_all_as_kwargs_np={},
+        frontend_method_data=frontend_method_data,
+        init_flags=init_flags,
+        method_flags=method_flags,
+        frontend=frontend,
+        on_device=on_device,
+    )
+
+
 # as_real
 @handle_frontend_method(
     class_tree=CLASS_TREE,
Repository: ivy-llc/ivy

Issue:
real



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/paddle/tensor/tensor.py b/ivy/functional/frontends/paddle/tensor/tensor.py
--- a/ivy/functional/frontends/paddle/tensor/tensor.py
+++ b/ivy/functional/frontends/paddle/tensor/tensor.py
@@ -756,3 +756,7 @@ def mean(self, axis=None, keepdim=False, name=None):
     )
     def less_equal(self, y, name=None):
         return paddle_frontend.less_equal(self._ivy_array, y)
+
+    @with_supported_dtypes({"2.5.1 and below": ("complex64", "complex128")}, "paddle")
+    def real(self, name=None):
+        return paddle_frontend.real(self._ivy_array)


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py b/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
--- a/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
+++ b/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
@@ -3064,6 +3064,43 @@ def test_paddle_tensor_rad2deg(
     )
 
 
+@handle_frontend_method(
+    class_tree=CLASS_TREE,
+    init_tree="paddle.to_tensor",
+    method_name="real",
+    dtype_and_x=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes(kind="valid"),
+        num_arrays=2,
+        min_num_dims=1,
+        allow_inf=True,
+    ),
+)
+def test_paddle_tensor_real(
+    dtype_and_x,
+    frontend_method_data,
+    init_flags,
+    method_flags,
+    frontend,
+    on_device,
+    backend_fw,
+):
+    input_dtype, x = dtype_and_x
+    helpers.test_frontend_method(
+        init_input_dtypes=input_dtype,
+        backend_to_test=backend_fw,
+        init_all_as_kwargs_np={
+            "data": x[0],
+        },
+        method_input_dtypes=input_dtype,
+        method_all_as_kwargs_np={},
+        frontend_method_data=frontend_method_data,
+        init_flags=init_flags,
+        method_flags=method_flags,
+        frontend=frontend,
+        on_device=on_device,
+    )
+
+
 # reciprocal
 @handle_frontend_method(
     class_tree=CLASS_TREE,
Repository: ivy-llc/ivy

Issue:
floor_mod



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/paddle/tensor/tensor.py b/ivy/functional/frontends/paddle/tensor/tensor.py
--- a/ivy/functional/frontends/paddle/tensor/tensor.py
+++ b/ivy/functional/frontends/paddle/tensor/tensor.py
@@ -781,6 +781,12 @@ def floor_divide(self, y, name=None):
     def mod(self, y, name=None):
         return paddle_frontend.Tensor(ivy.fmod(self._ivy_array, _to_ivy_array(y)))
 
+    @with_supported_dtypes(
+        {"2.5.2 and below": ("float32", "float64", "int32", "int64")}, "paddle"
+    )
+    def floor_mod(self, y, name=None):
+        return paddle_frontend.remainder(self, y)
+
     # cond
     @with_supported_dtypes({"2.5.2 and below": ("float32", "float64")}, "paddle")
     def cond(self, p=None, name=None):


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py b/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
--- a/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
+++ b/ivy_tests/test_ivy/test_frontends/test_paddle/test_tensor/test_tensor.py
@@ -5392,6 +5392,42 @@ def test_paddle_tensor_expand(
     )
 
 
+# floor_mod
+@handle_frontend_method(
+    class_tree=CLASS_TREE,
+    init_tree="paddle.to_tensor",
+    method_name="floor_mod",
+    dtypes_and_x=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("valid"),
+        num_arrays=2,
+        min_value=2,
+        shared_dtype=True,
+    ),
+)
+def test_paddle_tensor_floor_mod(
+    dtypes_and_x,
+    frontend_method_data,
+    init_flags,
+    method_flags,
+    frontend,
+    on_device,
+    backend_fw,
+):
+    input_dtype, x = dtypes_and_x
+    helpers.test_frontend_method(
+        init_input_dtypes=input_dtype,
+        backend_to_test=backend_fw,
+        init_all_as_kwargs_np={"data": x[0]},
+        method_input_dtypes=input_dtype,
+        method_all_as_kwargs_np={"y": x[1]},
+        frontend_method_data=frontend_method_data,
+        init_flags=init_flags,
+        method_flags=method_flags,
+        frontend=frontend,
+        on_device=on_device,
+    )
+
+
 @handle_frontend_method(
     class_tree=CLASS_TREE,
     init_tree="paddle.to_tensor",
Repository: ivy-llc/ivy

Issue:
Class method `from_CPTensor()` should use `cls` as the first parameter.
The class method `from_CPTensor()` in the `Parafac2Tensor` class currently uses `self` as the first parameter. However, it should use `cls` as the first parameter, as it is a class method and using `cls` is generally recommended.

https://github.com/unifyai/ivy/blob/f45b3c232e97fb53f661b633f7590f50d53b4196/ivy/data_classes/factorized_tensor/parafac2_tensor.py#L92




Discussion:

Resolving patch:
diff --git a/ivy/data_classes/factorized_tensor/parafac2_tensor.py b/ivy/data_classes/factorized_tensor/parafac2_tensor.py
--- a/ivy/data_classes/factorized_tensor/parafac2_tensor.py
+++ b/ivy/data_classes/factorized_tensor/parafac2_tensor.py
@@ -89,7 +89,7 @@ def n_param(self):
             return factors_params
 
     @classmethod
-    def from_CPTensor(self, cp_tensor, parafac2_tensor_ok=False):
+    def from_CPTensor(cls, cp_tensor, parafac2_tensor_ok=False):
         """
         Create a Parafac2Tensor from a CPTensor.
 


Test patch:
Repository: ivy-llc/ivy

Issue:
dev_util



Discussion:

Resolving patch:
diff --git a/ivy/functional/ivy/device.py b/ivy/functional/ivy/device.py
--- a/ivy/functional/ivy/device.py
+++ b/ivy/functional/ivy/device.py
@@ -639,7 +639,10 @@ def percent_used_mem_on_dev(
 
 
 @handle_exceptions
-def dev_util(device: Union[ivy.Device, ivy.NativeDevice], /) -> float:
+def dev_util(
+    device: Union[ivy.Device, ivy.NativeDevice],
+    /,
+) -> float:
     """
     Get the current utilization (%) for a given device.
 


Test patch:
Repository: ivy-llc/ivy

Issue:
count_nonzero



Discussion:

Resolving patch:
diff --git a/ivy/functional/backends/torch/experimental/elementwise.py b/ivy/functional/backends/torch/experimental/elementwise.py
--- a/ivy/functional/backends/torch/experimental/elementwise.py
+++ b/ivy/functional/backends/torch/experimental/elementwise.py
@@ -95,14 +95,14 @@ def count_nonzero(
         return x
     if isinstance(axis, int):
         if axis == -1:
-            temp = x.dim() - 2
+            temp = x.dim() - 1
             if temp < -1:
                 temp = 0
             return x.unsqueeze(temp)
-        return x.unsqueeze(axis - 1)
+        return x.unsqueeze(axis)
     elif axis is not None:
         for d in sorted(axis):
-            x = x.unsqueeze(d - 1)
+            x = x.unsqueeze(d)
         return x
     return x
 
diff --git a/ivy/functional/frontends/jax/numpy/searching_sorting.py b/ivy/functional/frontends/jax/numpy/searching_sorting.py
--- a/ivy/functional/frontends/jax/numpy/searching_sorting.py
+++ b/ivy/functional/frontends/jax/numpy/searching_sorting.py
@@ -56,6 +56,21 @@ def argwhere(a, /, *, size=None, fill_value=None):
     return result.reshape(result.shape[0], num_of_dimensions)
 
 
+@with_unsupported_dtypes(
+    {
+        "0.4.16 and below": (
+            "uint8",
+            "int8",
+            "bool",
+        )
+    },
+    "jax",
+)
+@to_ivy_arrays_and_back
+def count_nonzero(a, axis=None, keepdims=False):
+    return ivy.astype(ivy.count_nonzero(a, axis=axis, keepdims=keepdims), "int64")
+
+
 @to_ivy_arrays_and_back
 def extract(condition, arr):
     if condition.dtype is not bool:


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_searching_sorting.py b/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_searching_sorting.py
--- a/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_searching_sorting.py
+++ b/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_searching_sorting.py
@@ -169,6 +169,42 @@ def test_jax_argwhere(
     )
 
 
+# count_nonzero
+@handle_frontend_test(
+    fn_tree="jax.numpy.count_nonzero",
+    dtype_input_axis=helpers.dtype_values_axis(
+        available_dtypes=helpers.get_dtypes("valid"),
+        min_num_dims=1,
+        force_int_axis=True,
+        valid_axis=True,
+        allow_neg_axes=True,
+    ),
+    keepdims=st.booleans(),
+    test_with_out=st.just(False),
+)
+def test_jax_count_nonzero(
+    dtype_input_axis,
+    keepdims,
+    frontend,
+    test_flags,
+    fn_tree,
+    backend_fw,
+    on_device,
+):
+    input_dtype, x, axis = dtype_input_axis
+    helpers.test_frontend_function(
+        input_dtypes=input_dtype,
+        backend_to_test=backend_fw,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        a=x[0],
+        axis=axis,
+        keepdims=keepdims,
+    )
+
+
 # extract
 @handle_frontend_test(
     fn_tree="jax.numpy.extract",