adding roots function for jax frontend

ivy-llc · Sep 13, 2023 · 3882b0a · 3882b0a
1 parent 4953e1b
commit 3882b0a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/ivy/functional/frontends/jax/numpy/mathematical_functions.py b/ivy/functional/frontends/jax/numpy/mathematical_functions.py
@@ -669,7 +669,7 @@ def roots(p):
         # finding the eigenvalue of comp_matrix which is also the roots
         roots = ivy.eigvals(comp_matrix)
 
-    return roots
+    return ivy.flip(roots)
 
 
 @to_ivy_arrays_and_back

diff --git a/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_mathematical_functions.py b/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_mathematical_functions.py
@@ -2826,7 +2826,7 @@ def test_jax_remainder(
         num_arrays=1,
         min_num_dims=1,
         max_num_dims=1,
-        min_dim_size=0,
+        min_dim_size=3,
     ),
 )
 def test_jax_roots(
@@ -2839,15 +2839,14 @@ def test_jax_roots(
     backend_fw,
 ):
     input_dtype, x = dtype_and_x
-    assume("float16" not in input_dtype)
     helpers.test_frontend_function(
         input_dtypes=input_dtype,
         test_flags=test_flags,
         frontend=frontend,
         backend_to_test=backend_fw,
         fn_tree=fn_tree,
         on_device=on_device,
-        p=x,
+        p=x[0],
         atol=1e-05,
         rtol=1e-03,
     )