Repository: ivy-llc/ivy

Issue:
fold_in



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/jax/random.py b/ivy/functional/frontends/jax/random.py
--- a/ivy/functional/frontends/jax/random.py
+++ b/ivy/functional/frontends/jax/random.py
@@ -172,6 +172,17 @@ def randint(key, shape, minval, maxval, dtype="int64"):
     seed = _get_seed(key)
     return ivy.randint(minval, maxval, shape=shape, dtype=dtype, seed=seed)
 
+@to_ivy_arrays_and_back
+def bernoulli(key, p=0.5, shape=None):
+    seed = _get_seed(key)
+    return ivy.bernoulli(p, shape=shape, seed=seed)
+
+@to_ivy_arrays_and_back
+def fold_in(key, data):
+    s = ivy.bitwise_left_shift(
+        ivy.asarray(data, dtype=ivy.uint32), ivy.array(32, dtype=ivy.uint32)
+    )
+    return ivy.bitwise_xor(key, s)
 
 @to_ivy_arrays_and_back
 def permutation(key, x, axis=0, independent=False):
@@ -184,5 +195,4 @@ def permutation(key, x, axis=0, independent=False):
         return ivy.shuffle(x, axis, seed=seed)
     rand = ivy.arange(x.shape[axis])
     ind = ivy.shuffle(rand, 0, seed=seed)
-
     return ivy.gather(x, ind, axis=axis)


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_jax/test_jax_random.py b/ivy_tests/test_ivy/test_frontends/test_jax/test_jax_random.py
--- a/ivy_tests/test_ivy/test_frontends/test_jax/test_jax_random.py
+++ b/ivy_tests/test_ivy/test_frontends/test_jax/test_jax_random.py
@@ -732,6 +732,124 @@ def call():
         assert u.shape == v.shape
 
 
+@st.composite
+def dtype_p_shape(draw):
+    dtype = draw(helpers.array_dtypes(available_dtypes=("float32", "float64")))
+    shape = draw(helpers.get_shape(allow_none=False, min_num_dims=1, max_num_dims=3))
+
+    dtype_and_probs = draw(
+        helpers.dtype_and_values(
+            available_dtypes=dtype, min_value=0, max_value=1, shape=shape
+        )
+    )
+    return dtype_and_probs, shape
+
+
+# TODO Update the test by fixing the uint32 unsupported problem
+@pytest.mark.xfail
+@handle_frontend_test(
+    fn_tree="jax.random.bernoulli",
+    dtype_key=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("integer", full=False),
+        min_value=0,
+        max_value=2000,
+        min_num_dims=1,
+        max_num_dims=1,
+        min_dim_size=2,
+        max_dim_size=2,
+    ),
+    dtype_p_shape_=dtype_p_shape(),
+)
+def test_jax_bernoulli(
+    *,
+    dtype_key,
+    dtype_p_shape_,
+    on_device,
+    fn_tree,
+    frontend,
+    test_flags,
+):
+    input_dtype, key = dtype_key
+    dtype_p, shape = dtype_p_shape_
+    dtype, p = dtype_p
+
+    def call():
+        return helpers.test_frontend_function(
+            input_dtypes=input_dtype + dtype,
+            frontend=frontend,
+            test_flags=test_flags,
+            fn_tree=fn_tree,
+            on_device=on_device,
+            test_values=False,
+            key=key[0],
+            p=p[0],
+            shape=shape,
+        )
+
+    ret = call()
+
+    if not ivy.exists(ret):
+        return
+
+    ret_np, ret_from_np = ret
+    ret_np = helpers.flatten_and_to_np(ret=ret_np)
+    ret_from_np = helpers.flatten_and_to_np(ret=ret_from_np)
+    for u, v in zip(ret_np, ret_from_np):
+        assert u.dtype == v.dtype
+        assert u.shape == v.shape
+
+
+# TODO Update the test by fixing the uint32 unsupported problem
+@pytest.mark.xfail
+@handle_frontend_test(
+    fn_tree="jax.random.fold_in",
+    dtype_key=helpers.dtype_and_values(
+        available_dtypes=["uint32"],
+        min_value=0,
+        max_value=2000,
+        min_num_dims=1,
+        max_num_dims=1,
+        min_dim_size=2,
+        max_dim_size=2,
+    ),
+    data=helpers.ints(),
+)
+def test_jax_fold_in(
+    *,
+    dtype_key,
+    data,
+    on_device,
+    fn_tree,
+    frontend,
+    test_flags,
+):
+    input_dtype, key = dtype_key
+
+    def call():
+        return helpers.test_frontend_function(
+            input_dtypes=input_dtype,
+            frontend=frontend,
+            test_flags=test_flags,
+            fn_tree=fn_tree,
+            on_device=on_device,
+            test_values=False,
+            key=key[0],
+            data=data,
+        )
+
+    ret = call()
+
+    if not ivy.exists(ret):
+        return
+
+    ret_np, ret_from_np = ret
+    ret_np = helpers.flatten_and_to_np(ret=ret_np)
+    ret_from_np = helpers.flatten_and_to_np(ret=ret_from_np)
+    for u, v in zip(ret_np, ret_from_np):
+        assert u.dtype == v.dtype
+        assert u.shape == v.shape
+
+  
 # TODO Update the test by fixing the uint32 unsupported problem
 @pytest.mark.xfail
 @handle_frontend_test(