Repository: ivy-llc/ivy

Issue:
multivariate_normal



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
@@ -352,3 +352,32 @@ def ball(key, d, p=2.0, shape=(), dtype="float64"):
     exp = -ivy.log(1 - uniform)
 
     return gn / (((ivy.abs(gn) ** p).sum(axis=-1) + exp) ** (1 / p))[..., None]
+
+
+@handle_jax_dtype
+@to_ivy_arrays_and_back
+@with_unsupported_dtypes(
+    {
+        "0.4.14 and below": (
+            "float16",
+            "bfloat16",
+        )
+    },
+    "jax",
+)
+def multivariate_normal(key, mean, cov, shape=None, dtype="float64", method="cholesky"):
+    if shape is None:
+        shape = ivy.broadcast_shapes(mean.shape[:-1], cov.shape[:-2])
+    if method == "cholesky":
+        cov_factor = ivy.cholesky(cov)
+    elif method == "eigh":
+        (w, v) = ivy.eigh(cov)
+        cov_factor = v * ivy.sqrt(w[..., None, :])
+    elif method == "svd":
+        (u, s, _) = ivy.svd(cov)
+        cov_factor = u * ivy.sqrt(s[..., None, :])
+
+    rand_normal = normal(key=key, shape=shape + mean.shape[-1:], dtype=dtype)
+    result = mean + ivy.einsum("...ij,...j->...i", cov_factor, rand_normal.ivy_array)
+
+    return result


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
@@ -2,6 +2,8 @@
 import pytest
 from hypothesis import strategies as st
 import ivy
+import numpy as np
+import sys
 
 # local
 import ivy_tests.test_ivy.helpers as helpers
@@ -1417,3 +1419,102 @@ def call():
     for u, v in zip(ret_np, ret_from_np):
         assert u.dtype == v.dtype
         assert u.shape == v.shape
+
+
+@st.composite
+def get_mean_cov_vector(draw):
+    input_dtype = draw(
+        st.shared(
+            st.sampled_from(draw(helpers.get_dtypes("float"))),
+            key="shared_dtype",
+        )
+    )
+    shared_size = draw(
+        st.shared(helpers.ints(min_value=2, max_value=4), key="shared_size")
+    )
+
+    # Generate shape for mean vector (..., n)
+    dtype_mean = draw(
+        helpers.array_values(
+            dtype=input_dtype,
+            shape=tuple([shared_size]),
+            min_value=2,
+            max_value=5,
+        )
+    )
+
+    # Generate shape for covariance matrix (..., n, n)
+    dtype_cov = draw(
+        helpers.array_values(
+            dtype=input_dtype,
+            shape=tuple([shared_size, shared_size]),
+            min_value=2,
+            max_value=5,
+        ).filter(lambda x: np.linalg.cond(x.tolist()) < 1 / sys.float_info.epsilon)
+    )
+
+    batch_shape = dtype_cov.shape[:-2]
+
+    return input_dtype, dtype_mean, dtype_cov, batch_shape
+
+
+@pytest.mark.xfail
+@handle_frontend_test(
+    fn_tree="jax.random.multivariate_normal",
+    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=helpers.get_dtypes("float", full=False),
+    mean_cov_vector=get_mean_cov_vector(),
+    method=st.sampled_from(["cholesky", "eigh", "svd"]),
+    test_with_out=st.just(False),
+)
+def test_jax_multivariate_normal(
+    *,
+    dtype_key,
+    mean_cov_vector,
+    dtype,
+    method,
+    frontend,
+    backend_fw,
+    test_flags,
+    fn_tree,
+):
+    input_dtype, key = dtype_key
+    shared_dtype, mean, cov, shape = mean_cov_vector
+
+    spd = np.matmul(cov.T, cov) + np.identity(cov.shape[0])
+
+    def call():
+        helpers.test_frontend_function(
+            input_dtypes=input_dtype + [shared_dtype],
+            frontend=frontend,
+            test_flags=test_flags,
+            backend_to_test=backend_fw,
+            fn_tree=fn_tree,
+            test_values=False,
+            key=key[0],
+            mean=mean,
+            cov=spd,
+            shape=shape,
+            dtype=dtype[0],
+            method=method,
+        )
+
+    ret = call()
+
+    if not ivy.exists(ret):
+        return
+
+    ret_np, ret_from_np = ret
+    ret_np = helpers.flatten_and_to_np(ret=ret_np, backend=backend_fw)
+    ret_from_np = helpers.flatten_and_to_np(ret=ret_from_np, backend=backend_fw)
+    for u, v in zip(ret_np, ret_from_np):
+        assert u.dtype == v.dtype
+        assert u.shape == v.shape