Repository: ivy-llc/ivy

Issue:
 drichlet



Discussion:
This issue has been labelled as stale. If you would like to continue working on this issue, then please add another comment or create a PR . If a PR has already been created which refers to this issue, then you should explicitly mention this issue in the relevant PR. For more information please refer to this page in our documentation: https://lets-unify.ai/ivy/contributing/1_the_basics.html#todo-list-issues

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
@@ -45,3 +45,20 @@ def _get_seed(key):
 def beta(key, a, b, shape=None, dtype=None):
     seed = _get_seed(key)
     return ivy.beta(a, b, shape=shape, dtype=dtype, seed=seed)
+
+
+@handle_jax_dtype
+@to_ivy_arrays_and_back
+@with_unsupported_dtypes(
+    {
+        "0.3.14 and below": (
+            "float16",
+            "bfloat16",
+        )
+    },
+    "jax",
+)
+def dirichlet(key, alpha, shape=None, dtype="float32"):
+    seed = _get_seed(key)
+    alpha = ivy.astype(alpha, dtype)
+    return ivy.dirichlet(alpha, size=shape, dtype=dtype, seed=seed)


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
@@ -183,3 +183,70 @@ def call():
     for (u, v) in zip(ret_np, ret_from_np):
         assert u.dtype == v.dtype
         assert u.shape == v.shape
+
+
+@handle_frontend_test(
+    fn_tree="jax.random.dirichlet",
+    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,
+    ),
+    dtype_alpha=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("float", full=False),
+        shape=st.tuples(
+            st.integers(min_value=2, max_value=5),
+        ),
+        min_value=1.1,
+        max_value=100.0,
+        exclude_min=True,
+    ),
+    shape=helpers.get_shape(
+        min_num_dims=2, max_num_dims=2, min_dim_size=2, max_dim_size=5
+    ),
+    dtype=helpers.get_dtypes("float", full=False),
+    test_with_out=st.just(False),
+)
+def test_jax_dirichlet(
+    *,
+    dtype_key,
+    dtype_alpha,
+    shape,
+    dtype,
+    on_device,
+    fn_tree,
+    frontend,
+    test_flags,
+):
+    input_dtype, key = dtype_key
+    _, alpha = dtype_alpha
+
+    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],
+            alpha=alpha[0],
+            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