Repository: ivy-llc/ivy

Issue:
fromfunction
Working on (#8431-> #17806)


Discussion:
 frombuffer
 fromfile
 fromfunction
 fromiter
 fromstring
 from_dlpack

Resolving patch:
diff --git a/ivy/functional/frontends/jax/numpy/logic.py b/ivy/functional/frontends/jax/numpy/logic.py
--- a/ivy/functional/frontends/jax/numpy/logic.py
+++ b/ivy/functional/frontends/jax/numpy/logic.py
@@ -260,7 +260,9 @@ def right_shift(x1, x2, /):
 
 
 @to_ivy_arrays_and_back
-@with_unsupported_dtypes({"0.4.14 and below": ("bfloat16", "bool")}, "jax")
+@with_unsupported_dtypes(
+    {"0.4.14 and below": ("bfloat16", 'bool')}, "jax"
+)
 def setxor1d(ar1, ar2, assume_unique=False):
     common_dtype = ivy.promote_types(ivy.dtype(ar1), ivy.dtype(ar2))
     ar1 = ivy.asarray(ar1, dtype=common_dtype)
@@ -287,3 +289,27 @@ def setxor1d(ar1, ar2, assume_unique=False):
 
 alltrue = all
 sometrue = any
+
+@to_ivy_arrays_and_back
+@with_unsupported_dtypes({"0.4.14 and below": ("bfloat16",)}, "jax")
+def fromfunction(function, shape, *, dtype=float, **kwargs):
+    def canonicalize_shape(shape, context="shape argument"):
+        if isinstance(shape, int):
+            return (shape,)
+        elif isinstance(shape, list):
+            return tuple(shape)
+        elif isinstance(shape, tuple):
+            return shape
+        else:
+            msg = "{} must be an int, list, or tuple, but got {}."
+            raise TypeError(msg.format(context, type(shape)))
+
+    arr = ivy.zeros(shape, dtype=dtype)
+    shape = canonicalize_shape(shape)
+    # Iterate over the indices of the array
+    for indices in ivy.ndindex(shape):
+        f_indices = indices
+        ivy.set_nest_at_index(
+            arr, f_indices, ivy.asarray(function(*indices, **kwargs), dtype=dtype)
+        )
+    return arr


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_logic.py b/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_logic.py
--- a/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_logic.py
+++ b/ivy_tests/test_ivy/test_frontends/test_jax/test_numpy/test_logic.py
@@ -1148,3 +1148,57 @@ def test_jax_numpy_packbits(
         bitorder=bitorder,
         backend_to_test=backend_fw,
     )
+
+
+@st.composite
+def _func_and_shape_dtype_helper(draw):
+    # here assumption is that the input func will take the len(shape) no of parameters
+    def add_numbers(*args):
+        total = 0
+        for num in args:
+            total += num
+        return total
+
+    shape = draw(
+        helpers.get_shape(
+            allow_none=False,
+            min_num_dims=1,
+            max_num_dims=3,
+            min_dim_size=1,
+            max_dim_size=3,
+        )
+    )
+
+    dtype = draw(helpers.get_dtypes("valid"))
+
+    return add_numbers, shape, dtype[0]
+
+
+# fromfunction
+@handle_frontend_test(
+    fn_tree="jax.numpy.fromfunction",
+    input_dtype=helpers.get_dtypes("valid"),
+    function_and_shape_and_dtype=_func_and_shape_dtype_helper(),
+    test_with_out=st.just(False),
+)
+def test_jax_numpy_fromfunction(
+    input_dtype,
+    function_and_shape_and_dtype,
+    backend_fw,
+    frontend,
+    on_device,
+    fn_tree,
+    test_flags,
+):
+    function, shape, dtype = function_and_shape_and_dtype
+    helpers.test_frontend_function(
+        input_dtypes=input_dtype,
+        backend_to_test=backend_fw,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        function=function,
+        shape=shape,
+        dtype=dtype,
+    )