Repository: ivy-llc/ivy

Issue:
product



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/jax/numpy/mathematical_functions.py b/ivy/functional/frontends/jax/numpy/mathematical_functions.py
--- a/ivy/functional/frontends/jax/numpy/mathematical_functions.py
+++ b/ivy/functional/frontends/jax/numpy/mathematical_functions.py
@@ -620,3 +620,31 @@ def polymul(a1, a2, *, trim_leading_zeros=False):
 def signbit(x, /):
     x = ivy.array(x)
     return ivy.signbit(x)
+
+
+@to_ivy_arrays_and_back
+def product(
+    a,
+    *,
+    axis=None,
+    dtype=None,
+    keepdims=False,
+    initial=None,
+    where=None,
+    promote_integers=True,
+    out=None,
+):
+    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 initial is not None:
+        if axis is not None:
+            s = ivy.to_list(ivy.shape(a, as_array=True))
+            s[axis] = 1
+            header = ivy.full(ivy.Shape(tuple(s)), initial)
+            a = ivy.concat([header, a], axis=axis)
+        else:
+            a[0] *= initial
+    return ivy.prod(a, axis=axis, dtype=dtype, keepdims=keepdims, out=out)


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_jax/test_jax_numpy_math.py b/ivy_tests/test_ivy/test_frontends/test_jax/test_jax_numpy_math.py
--- a/ivy_tests/test_ivy/test_frontends/test_jax/test_jax_numpy_math.py
+++ b/ivy_tests/test_ivy/test_frontends/test_jax/test_jax_numpy_math.py
@@ -17,6 +17,31 @@
 )
 
 
+@st.composite
+def _get_castable_dtypes_values(draw, *, allow_nan=False, use_where=False):
+    available_dtypes = helpers.get_dtypes("numeric")
+    shape = draw(helpers.get_shape(min_num_dims=1, max_num_dims=4, max_dim_size=6))
+    dtype, values = draw(
+        helpers.dtype_and_values(
+            available_dtypes=available_dtypes,
+            num_arrays=1,
+            large_abs_safety_factor=24,
+            small_abs_safety_factor=24,
+            safety_factor_scale="log",
+            shape=shape,
+            allow_nan=allow_nan,
+        )
+    )
+    axis = draw(helpers.get_axis(shape=shape, force_int=True))
+    dtype1, values, dtype2 = draw(
+        helpers.get_castable_dtype(draw(available_dtypes), dtype[0], values[0])
+    )
+    if use_where:
+        where = draw(np_frontend_helpers.where(shape=shape))
+        return [dtype1], [values], axis, dtype2, where
+    return [dtype1], [values], axis, dtype2
+
+
 # sign
 @handle_frontend_test(
     fn_tree="jax.numpy.sign",
@@ -2771,3 +2796,44 @@ def test_jax_numpy_signbit(
         on_device=on_device,
         x=x[0],
     )
+
+
+@handle_frontend_test(
+    fn_tree="jax.numpy.product",
+    dtype_x_axis_dtype_where=_get_castable_dtypes_values(use_where=True),
+    keepdims=st.booleans(),
+    initial=st.one_of(st.floats(min_value=-100, max_value=100)),
+    promote_integers=st.booleans(),
+)
+def test_jax_numpy_product(
+    dtype_x_axis_dtype_where,
+    keepdims,
+    initial,
+    promote_integers,
+    frontend,
+    test_flags,
+    fn_tree,
+    on_device,
+):
+    input_dtypes, x, axis, dtype, where = dtype_x_axis_dtype_where
+    if ivy.current_backend_str() == "torch":
+        assume(not test_flags.as_variable[0])
+    where, input_dtypes, test_flags = np_frontend_helpers.handle_where_and_array_bools(
+        where=where,
+        input_dtype=input_dtypes,
+        test_flags=test_flags,
+    )
+    helpers.test_frontend_function(
+        input_dtypes=input_dtypes,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        a=x[0],
+        axis=axis,
+        dtype=dtype,
+        keepdims=keepdims,
+        initial=initial,
+        where=where,
+        promote_integers=promote_integers,
+    )