jax.numpy.convolve frontend function and test. issue #10499

ivy/functional/frontends/jax/numpy/mathematical_functions.py

@@ -38,6 +38,37 @@ def arctan2(x1, x2):
     return ivy.atan2(x1, x2)
 
 
+@to_ivy_arrays_and_back
+def convolve(a, v, mode='full', *, precision=None):
+    a, v = promote_types_of_jax_inputs(a, v)
+    if ivy.get_num_dims(a) != 1:
+        raise ValueError(
+            "convolve() only support 1-dimensional inputs."
+        )
+    if len(a) == 0 or len(v) == 0:
+        raise ValueError(
+            f"convolve: inputs cannot be empty, got shapes {a.shape} and {v.shape}."
+        )
+    if len(a) < len(v):
+        a, v = v, a
+    v = ivy.flip(v)
+
+    out_order = slice(None)
+
+    if mode == 'valid':
+        padding = [(0, 0)]
+    elif mode == 'same':
+        padding = [(v.shape[0] // 2, v.shape[0] - v.shape[0] // 2 - 1)]
+    elif mode == 'full':
+        padding = [(v.shape[0] - 1, v.shape[0] - 1)]
+
+    result = ivy.conv_general_dilated(
+        a[None, None, :], v[:, None, None], (1,),
+        padding, dims=1, data_format='channel_first',
+    )
+    return result[0, 0, out_order]
+
+
 @to_ivy_arrays_and_back
 def cos(x):
     return ivy.cos(x)

ivy_tests/test_ivy/test_frontends/test_jax/test_jax_numpy_math.py

@@ -171,6 +171,51 @@ def test_jax_numpy_arctan2(
     )
 
 
+@st.composite
+def _get_pooling_mode(draw):
+    mode = draw(st.sampled_from(['full', 'valid', 'same']))
+    return mode
+
+
+# convolve
+@handle_frontend_test(
+    fn_tree="jax.numpy.convolve",
+    dtype_and_x=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("float"),
+        num_arrays=2,
+        min_num_dims=1,
+        max_num_dims=1,
+        min_value=-1e04,
+        max_value=1e04,
+    ),
+    mode=_get_pooling_mode(),
+)
+def test_jax_numpy_convolve(
+        *,
+        dtype_and_x,
+        on_device,
+        fn_tree,
+        frontend,
+        test_flags,
+        mode,
+):
+    input_dtype, x = dtype_and_x
+    assume("float16" not in input_dtype)
+    helpers.test_frontend_function(
+        input_dtypes=input_dtype,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        rtol=1e-4,
+        atol=1e-4,
+        on_device=on_device,
+        a=x[0],
+        v=x[1],
+        mode=mode,
+        precision=None,
+    )
+
+
 @handle_frontend_test(
     fn_tree="jax.numpy.cos",
     dtype_and_x=helpers.dtype_and_values(available_dtypes=helpers.get_dtypes("float")),

ivy_tests/test_ivy/test_frontends/test_torch/test_reduction_ops.py

@@ -8,9 +8,8 @@
 from ivy_tests.test_ivy.test_functional.test_core.test_statistical import (
     statistical_dtype_values,
 )
-from ivy_tests.test_ivy.test_functional.test_experimental.test_core.test_statistical import (
-    statistical_dtype_values as statistical_dtype_values_experimental,
-)
+from ivy_tests.test_ivy.test_functional.test_experimental.test_core.test_statistical \
+    import (statistical_dtype_values as statistical_dtype_values_experimental)
 
 
 @handle_frontend_test(