Jax array prod method (#19724)

Co-authored-by: hirwa-nshuti <[email protected]>
ivy-llc · Aug 16, 2023 · aeac992 · aeac992
1 parent 0e955a8
commit aeac992
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 2 deletions.
diff --git a/ivy/functional/frontends/jax/array.py b/ivy/functional/frontends/jax/array.py
@@ -127,6 +127,27 @@ def nonzero(self, *, size=None, fill_value=None):
             fill_value=fill_value,
         )
 
+    def prod(
+        self,
+        axis=None,
+        dtype=None,
+        keepdims=False,
+        initial=None,
+        where=None,
+        promote_integers=True,
+        out=None,
+    ):
+        return jax_frontend.numpy.product(
+            self,
+            axis=axis,
+            dtype=self.dtype,
+            keepdims=keepdims,
+            initial=initial,
+            where=where,
+            promote_integers=promote_integers,
+            out=out,
+        )
+
     def ravel(self, order="C"):
         return jax_frontend.numpy.ravel(
             self,

diff --git a/ivy/functional/frontends/jax/numpy/mathematical_functions.py b/ivy/functional/frontends/jax/numpy/mathematical_functions.py
@@ -698,8 +698,10 @@ def product(
     if ivy.is_array(where):
         a = ivy.where(where, a, ivy.default(out, ivy.ones_like(a)), out=out)
     if promote_integers:
-        if dtype is None:
-            dtype = a.dtype
+        if ivy.is_uint_dtype(a.dtype):
+            dtype = "uint64"
+        elif ivy.is_int_dtype(a.dtype):
+            dtype = "int64"
     if initial is not None:
         if axis is not None:
             s = ivy.to_list(ivy.shape(a, as_array=True))

diff --git a/ivy_tests/test_ivy/test_frontends/test_jax/test_array.py b/ivy_tests/test_ivy/test_frontends/test_jax/test_array.py
@@ -576,6 +576,48 @@ def test_jax_array_nonzero(
     )
 
 
+@handle_frontend_method(
+    class_tree=CLASS_TREE,
+    init_tree="jax.numpy.array",
+    method_name="prod",
+    dtype_x_axis=helpers.dtype_values_axis(
+        available_dtypes=helpers.get_dtypes("numeric"),
+        force_int_axis=True,
+        valid_axis=True,
+        min_dim_size=2,
+        max_dim_size=10,
+        min_num_dims=2,
+    ),
+)
+def test_jax_prod(
+    dtype_x_axis,
+    on_device,
+    frontend,
+    frontend_method_data,
+    init_flags,
+    method_flags,
+    backend_fw,
+):
+    input_dtype, x, axis = dtype_x_axis
+    helpers.test_frontend_method(
+        init_input_dtypes=input_dtype,
+        backend_to_test=backend_fw,
+        init_all_as_kwargs_np={
+            "object": x[0],
+        },
+        method_input_dtypes=input_dtype,
+        method_all_as_kwargs_np={
+            "axis": axis,
+        },
+        frontend=frontend,
+        frontend_method_data=frontend_method_data,
+        init_flags=init_flags,
+        method_flags=method_flags,
+        on_device=on_device,
+        atol_=1e-04,
+    )
+
+
 @handle_frontend_method(
     class_tree=CLASS_TREE,
     init_tree="jax.numpy.array",