Repository: ivy-llc/ivy

Issue:
ball



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
@@ -1,6 +1,8 @@
+#global
+import operator
 # local
 import ivy
-from ivy.func_wrapper import with_unsupported_dtypes
+from ivy.func_wrapper import with_unsupported_dtypes, with_supported_dtypes
 from ivy.functional.frontends.jax.func_wrapper import (
     to_ivy_arrays_and_back,
     handle_jax_dtype,
@@ -303,3 +305,29 @@ def maxwell(key, shape=None, dtype="float64"):
     # applying inverse transform sampling
     x = (z**2) * ivy.exp(-(z**2) / 2)
     return x
+
+
+@handle_jax_dtype
+@to_ivy_arrays_and_back
+@with_supported_dtypes(
+    {
+        "0.4.13 and below": (
+            "float32",
+            "float64",
+        )
+    },
+    "jax",
+)
+def ball(key, d, p=2.0, shape=(), dtype="float64"):
+    seed = _get_seed(key)
+    d = operator.index(d)
+
+    g = ivy.gamma(1 / p, 1.0, shape=shape, dtype=dtype, seed=seed)
+    b = ivy.bernoulli(ivy.array([0.5]), shape=shape, dtype=dtype, seed=seed)
+    r = 2 * b - 1
+    gn =  r * g ** (1 / p)
+
+    uniform = ivy.random_uniform(seed=seed, shape=shape, dtype=dtype)
+    exp = -ivy.log(1 - uniform)
+
+    return gn / (((ivy.abs(gn) ** p).sum(axis=-1) + exp) ** (1 / p))[..., None]


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_jax/test_random.py b/ivy_tests/test_ivy/test_frontends/test_jax/test_random.py
--- a/ivy_tests/test_ivy/test_frontends/test_jax/test_random.py
+++ b/ivy_tests/test_ivy/test_frontends/test_jax/test_random.py
@@ -1245,3 +1245,65 @@ def call():
     for u, v in zip(ret_np, ret_from_np):
         assert u.dtype == v.dtype
         assert u.shape == v.shape
+
+
+@pytest.mark.xfail
+@handle_frontend_test(
+    fn_tree="jax.random.ball",
+    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,
+    ),
+    shape=helpers.get_shape(
+        min_num_dims=1, max_num_dims=6, min_dim_size=1, max_dim_size=6
+    ),
+    dtype=helpers.get_dtypes("float", full=False),
+    d=st.integers(min_value=1, max_value=100),
+    p=st.floats(min_value=1e-5, max_value=100, exclude_min=True),
+    test_with_out=st.just(False),
+)
+def test_jax_ball(
+    *,
+    dtype_key,
+    d,
+    p,
+    shape,
+    dtype,
+    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],
+            d=d,
+            p=p,
+            shape=shape,
+            dtype=dtype[0],
+        )
+    
+    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