ENH: add np.convolve

Author: ev-br

autogen/numpy_api_dump.py

@@ -176,14 +176,6 @@ def compress(condition, a, axis=None, out=None):
     raise NotImplementedError
 
 
-def convolve(a, v, mode="full"):
-    raise NotImplementedError
-
-
-def copyto(dst, src, casting="same_kind", where=True):
-    raise NotImplementedError
-
-
 def correlate(a, v, mode="valid"):
     raise NotImplementedError
 

torch_np/_funcs_impl.py

@@ -570,6 +570,28 @@ def cov(
     return result
 
 
+def convolve(a: ArrayLike, v: ArrayLike, mode="full"):
+    # NumPy: if v is longer than a, the arrays are swapped before computation
+    if a.shape[0] < v.shape[0]:
+        a, v = v, a
+
+    dt = _dtypes_impl.result_type_impl((a.dtype, v.dtype))
+    a = _util.cast_if_needed(a, dt)
+    v = _util.cast_if_needed(v, dt)
+
+    padding = v.shape[0] - 1 if mode == "full" else mode
+
+    # 1. NumPy only accepts 1D arrays; PyTorch requires 2D inputs and 3D weights
+    # 2. flip the weights since numpy does and torch does not
+    aa = a[None, :]
+    vv = torch.flip(v, (0,))[None, None, :]
+
+    result = torch.nn.functional.conv1d(aa, vv, padding=padding)
+
+    # torch returns a 2D result, numpy returns a 1D array
+    return result[0, :]
+
+
 # ### logic & element selection ###
 
 

torch_np/tests/numpy_tests/core/test_numeric.py

@@ -13,6 +13,7 @@
     assert_, assert_equal, assert_raises_regex,
     assert_array_equal, assert_almost_equal, assert_array_almost_equal,
     assert_warns, # assert_array_max_ulp, HAS_REFCOUNT, IS_WASM
+    assert_allclose,
     )
 from numpy.core._rational_tests import rational
 
@@ -2441,7 +2442,6 @@ def test_mode(self):
             np.correlate(d, k, mode=None)
 
 
-@pytest.mark.xfail(reason="TODO")
 class TestConvolve:
     def test_object(self):
         d = [1.] * 100
@@ -2455,6 +2455,7 @@ def test_no_overwrite(self):
         assert_array_equal(d, np.ones(100))
         assert_array_equal(k, np.ones(3))
 
+    @pytest.mark.skip(reason='do not implement deprecated behavior')
     def test_mode(self):
         d = np.ones(100)
         k = np.ones(3)
@@ -2470,6 +2471,16 @@ def test_mode(self):
         with assert_raises(TypeError):
             np.convolve(d, k, mode=None)
 
+    def test_numpy_doc_examples(self):
+        conv = np.convolve([1, 2, 3], [0, 1, 0.5])
+        assert_allclose(conv, [0., 1., 2.5, 4., 1.5], atol=1e-15)
+
+        conv = np.convolve([1, 2, 3], [0, 1, 0.5], 'same')
+        assert_allclose(conv, [1., 2.5, 4.], atol=1e-15)
+
+        conv = np.convolve([1, 2, 3], [0, 1, 0.5], 'valid')
+        assert_allclose(conv, [2.5], atol=1e-15)
+
 
 class TestDtypePositional: