Repository: ivy-llc/ivy

Issue:
[Bug]: Incomplete implementation of jax.numpy.where
### Bug Explanation

The function accepts only condition argument in the native function but the implementation of `ivy.frontends.jax.numpy` doesn't follow the same behaviour and only works when all three arguments are passed

### Steps to Reproduce Bug

```
import jax.numpy as jnp
jnp.where(jnp.array([True,False,True,True]))
#(DeviceArray([0, 2, 3], dtype=int32),)

import ivy.functional.frontends.jax.numpy as jnp
jnp.where(jnp.array([True,False,True,True]))
#raises error

```

### Environment

Ubuntu

### Ivy Version

Github

### Backend

- [ ] NumPy
- [ ] TensorFlow
- [ ] PyTorch
- [ ] JAX

### Device

_No response_


Discussion:

Resolving patch:
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
@@ -167,7 +167,9 @@ def unique(
 
 
 @to_ivy_arrays_and_back
-def where(condition, x=None, y=None, size=None, fill_value=0):
+def where(condition, x=None, y=None, *, size=None, fill_value=0):
+    if x is None and y is None:
+        return nonzero(condition, size=size, fill_value=fill_value)
     if x is not None and y is not None:
         return ivy.where(condition, x, y)
     else:


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
@@ -8,6 +8,9 @@
 from ivy_tests.test_ivy.test_functional.test_core.test_searching import (
     _broadcastable_trio,
 )
+from ...test_numpy.test_sorting_searching_counting.test_searching import (
+    _broadcastable_trio as _where_helper,
+)
 
 
 # argmax
@@ -420,39 +423,40 @@ def test_jax_searchsorted(
 # where
 @handle_frontend_test(
     fn_tree="jax.numpy.where",
-    dtype_and_x=helpers.dtype_and_values(
-        available_dtypes=helpers.get_dtypes("numeric"), min_num_dims=1
-    ),
-    test_with_out=st.just(False),
+    broadcastables=_where_helper(),
+    only_cond=st.booleans(),
+    size=st.integers(min_value=1, max_value=20),
+    fill_value=st.one_of(st.integers(-10, 10), st.floats(-10, 10), st.booleans()),
 )
 def test_jax_where(
     *,
-    dtype_and_x,
+    broadcastables,
+    only_cond,
+    size,
+    fill_value,
     frontend,
     backend_fw,
     fn_tree,
     on_device,
     test_flags,
 ):
-    input_dtype, x = dtype_and_x
-    x = x[0]
-    condition = x > 0.5
-    x1 = x * 2
-    x2 = x * -2
-
-    # Convert input_dtype from list to string
-    input_dtype = input_dtype[0]
-
+    cond, x1, x2, dtype = broadcastables
+    if only_cond:
+        x1, x2 = None, None
+    else:
+        size, fill_value = None, None
     helpers.test_frontend_function(
-        input_dtypes=["bool", input_dtype, input_dtype],
+        input_dtypes=["bool", dtype],
         fn_tree=fn_tree,
         on_device=on_device,
         test_flags=test_flags,
         frontend=frontend,
         backend_to_test=backend_fw,
-        condition=condition,
+        condition=cond,
         x=x1,
         y=x2,
+        size=size,
+        fill_value=fill_value,
     )