Repository: ivy-llc/ivy

Issue:
diag_indices_from

mask_indices



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
Resuming work on diag_indices_from
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
Resuming work on mask_indicues

Resolving patch:
diff --git a/ivy/functional/frontends/jax/numpy/indexing.py b/ivy/functional/frontends/jax/numpy/indexing.py
--- a/ivy/functional/frontends/jax/numpy/indexing.py
+++ b/ivy/functional/frontends/jax/numpy/indexing.py
@@ -1,3 +1,6 @@
+# global
+import inspect
+
 # local
 import ivy
 from ivy.functional.frontends.jax.func_wrapper import (
@@ -42,6 +45,7 @@ def triu_indices_from(arr, k=0):
     return ivy.triu_indices(arr.shape[-2], arr.shape[-1], k)
 
 
+@to_ivy_arrays_and_back
 def tril_indices_from(arr, k=0):
     return ivy.tril_indices(arr.shape[-2], arr.shape[-1], k)
 
@@ -51,3 +55,28 @@ def tril_indices_from(arr, k=0):
 def unravel_index(indices, shape):
     ret = [x.astype("int64") for x in ivy.unravel_index(indices, shape)]
     return tuple(ret)
+
+
+@to_ivy_arrays_and_back
+def mask_indices(n, mask_func, k=0):
+    mask_func_obj = inspect.unwrap(mask_func)
+    mask_func_name = mask_func_obj.__name__
+    try:
+        ivy_mask_func_obj = getattr(ivy.functional.frontends.jax.numpy, mask_func_name)
+        a = ivy.ones((n, n))
+        mask = ivy_mask_func_obj(a, k=k)
+        indices = ivy.argwhere(mask.ivy_array)
+        return indices[:, 0], indices[:, 1]
+    except AttributeError as e:
+        print(f"Attribute error: {e}")
+
+
+@to_ivy_arrays_and_back
+def diag_indices_from(arr):
+    print(arr)
+    n = arr.shape[0]
+    ndim = ivy.get_num_dims(arr)
+    if not all(arr.shape[i] == n for i in range(ndim)):
+        raise ValueError("All dimensions of input must be of equal length")
+    idx = ivy.arange(n, dtype=int)
+    return (idx,) * ndim


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_jax/test_jax_numpy_indexing.py b/ivy_tests/test_ivy/test_frontends/test_jax/test_jax_numpy_indexing.py
--- a/ivy_tests/test_ivy/test_frontends/test_jax/test_jax_numpy_indexing.py
+++ b/ivy_tests/test_ivy/test_frontends/test_jax/test_jax_numpy_indexing.py
@@ -1,6 +1,8 @@
 # global
 from hypothesis import strategies as st, assume
 import numpy as np
+from jax.numpy import tril, triu
+
 
 # local
 import ivy_tests.test_ivy.helpers as helpers
@@ -339,3 +341,68 @@ def test_jax_numpy_unravel_index(
         indices=x[0],
         shape=shape,
     )
+
+
+@handle_frontend_test(
+    fn_tree="jax.numpy.mask_indices",
+    n=helpers.ints(min_value=3, max_value=10),
+    mask_func=st.sampled_from([triu, tril]),
+    k=helpers.ints(min_value=-5, max_value=5),
+    input_dtype=helpers.get_dtypes("numeric"),
+    test_with_out=st.just(False),
+    number_positional_args=st.just(2),
+)
+def test_jax_numpy_mask_indices(
+    n,
+    mask_func,
+    k,
+    input_dtype,
+    test_flags,
+    frontend,
+    fn_tree,
+    on_device,
+):
+    helpers.test_frontend_function(
+        input_dtypes=input_dtype,
+        test_flags=test_flags,
+        frontend=frontend,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        n=n,
+        mask_func=mask_func,
+        k=k,
+    )
+
+
+@st.composite
+def _get_dtype_square_x(draw):
+    dim_size = draw(helpers.ints(min_value=2, max_value=5))
+    dtype_x = draw(
+        helpers.dtype_and_values(
+            available_dtypes=helpers.get_dtypes("numeric"), shape=(dim_size, dim_size)
+        )
+    )
+    return dtype_x
+
+
+@handle_frontend_test(
+    dtype_x=_get_dtype_square_x(),
+    fn_tree="jax.numpy.diag_indices_from",
+    test_with_out=st.just(False),
+)
+def test_jax_numpy_diag_indices_from(
+    dtype_x,
+    test_flags,
+    frontend,
+    fn_tree,
+    on_device,
+):
+    dtype, x = dtype_x
+    helpers.test_frontend_function(
+        input_dtypes=dtype,
+        test_flags=test_flags,
+        frontend=frontend,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        arr=x[0],
+    )