Implement a Convolve Op

Author: Smit-create

aesara/tensor/basic.py

@@ -3811,7 +3811,7 @@ def make_node(self, a, choices):
         static_out_shape = ()
         for s in out_shape:
             try:
-                s_val = aesara.get_scalar_constant_value(s)
+                s_val = get_scalar_constant_value(s)
             except (NotScalarConstantError, AttributeError):
                 s_val = None
 

aesara/tensor/math.py

@@ -1,6 +1,6 @@
 import builtins
 import warnings
-from typing import TYPE_CHECKING, List, Optional, Tuple
+from typing import TYPE_CHECKING, List, Literal, Optional, Tuple
 
 import numpy as np
 
@@ -22,10 +22,12 @@
     cast,
     concatenate,
     constant,
+    get_scalar_constant_value,
     stack,
     switch,
 )
 from aesara.tensor.elemwise import CAReduce, DimShuffle, Elemwise, scalar_elemwise
+from aesara.tensor.exceptions import NotScalarConstantError
 from aesara.tensor.shape import shape, specify_broadcastable
 from aesara.tensor.type import (
     DenseTensorType,
@@ -46,6 +48,8 @@
 if TYPE_CHECKING:
     from numpy.typing import ArrayLike, DTypeLike
 
+    from aesara.tensor.var import TensorVariable
+
 # We capture the builtins that we are going to replace to follow the numpy API
 _abs = builtins.abs
 
@@ -2998,6 +3002,108 @@ def matmul(x1: "ArrayLike", x2: "ArrayLike", dtype: Optional["DTypeLike"] = None
     return MatMul(dtype=dtype)(x1, x2)
 
 
+class Convolve(Op):
+    __props__ = ("mode",)
+
+    def __init__(self, mode="full"):
+        self.mode = mode
+
+    @classmethod
+    def _get_output_shape(cls, a, v, shapes, mode, validate=False):
+        a_shape, v_shape = shapes
+        from aesara.tensor.math import maximum, minimum
+
+        if a.ndim == 1 and v.ndim == 1:
+            m, n = a_shape[0], v_shape[0]
+            if n is None or m is None:
+                return (None,)
+            if mode == "full":
+                return (m + n - 1,)
+            elif mode == "same":
+                return (maximum(m, n),)
+            elif mode == "valid":
+                return (maximum(m, n) - minimum(m, n) + 1,)
+            if validate:
+                raise ValueError("Invalid mode - must be full, valid or same")
+            return ()
+        else:
+            if validate:
+                raise ValueError("`a` and `v` must be 1-dim.")
+            return ()
+
+    def make_node(self, a, v):
+        a = as_tensor_variable(a)
+        v = as_tensor_variable(v)
+
+        if a.ndim != 1 or v.ndim != 1:
+            raise ValueError("inputs to `convolve` must be 1-dim.")
+
+        out_shape = self._get_output_shape(
+            a, v, (a.type.shape, v.type.shape), self.mode, validate=True
+        )
+
+        static_out_shape = ()
+        for s in out_shape:
+            try:
+                s_val = get_scalar_constant_value(s)
+            except (NotScalarConstantError, AttributeError):
+                s_val = None
+
+            if s_val:
+                static_out_shape += (s_val,)
+            else:
+                static_out_shape += (None,)
+
+        out = TensorType(
+            aes.upcast(a.type.dtype, v.type.dtype), shape=static_out_shape
+        )()
+        return Apply(self, [a, v], [out])
+
+    def perform(self, node, inputs, outputs):
+        a, v = inputs
+        outputs[0][0] = np.convolve(a, v, mode=self.mode)
+
+    def infer_shape(self, fgraph, node, shapes):
+        a, v = node.inputs
+        return [self._get_output_shape(a, v, shapes, self.mode)]
+
+
+def convolve(
+    a: "ArrayLike", v: "ArrayLike", mode: Literal["full", "same", "valid"] = "full"
+) -> "TensorVariable":
+    """Compute the discrete, linear convolution of two one-dimensional sequences.
+
+    Parameters
+    ----------
+    a, v
+        Input arrays, both should be one dimensional.
+    mode
+       'full':
+            By default, mode is 'full'. This returns the convolution
+            at each point of overlap.
+
+        'same':
+            Mode 'same'. Boundary effects are still visible.
+
+        'valid':
+            The convolution product is only given for points
+            where the signals overlap completely.
+            Values outside the signal boundary have no effect.
+
+    Returns
+    -------
+    out
+        Discrete, linear convolution of a and v.
+
+    Raises
+    ------
+    ValueError
+        If the a and v are not one-dimensional.
+
+    """
+    return Convolve(mode=mode)(a, v)
+
+
 __all__ = [
     "max_and_argmax",
     "max",
@@ -3126,6 +3232,7 @@ def matmul(x1: "ArrayLike", x2: "ArrayLike", dtype: Optional["DTypeLike"] = None
     "logsumexp",
     "hyp2f1",
     "hyp2f1_der",
+    "convolve",
 ]
 
 DEPRECATED_NAMES: List[Tuple[str, str, object]] = [

tests/tensor/test_math.py

@@ -33,6 +33,7 @@
 from aesara.tensor.elemwise import CAReduce, Elemwise
 from aesara.tensor.math import (
     Argmax,
+    Convolve,
     Dot,
     MatMul,
     MaxAndArgmax,
@@ -59,6 +60,7 @@
     clip,
     complex_from_polar,
     conj,
+    convolve,
     cos,
     cosh,
     cov,
@@ -3576,3 +3578,93 @@ def test_deprecations():
 
     with pytest.deprecated_call():
         from aesara.tensor.math import sqr  # noqa: F401 F811
+
+
+class TestConvolve(utt.InferShapeTester):
+    @pytest.mark.parametrize(
+        "a, v, mode",
+        [
+            (
+                np.arange(3).astype(config.floatX),
+                np.arange(6).astype(config.floatX),
+                "full",
+            ),
+            (
+                np.arange(5).astype(config.floatX),
+                np.arange(3).astype(config.floatX),
+                "valid",
+            ),
+            (
+                np.arange(8).astype(config.floatX),
+                np.arange(9).astype(config.floatX),
+                "same",
+            ),
+            (np.arange(3).astype(int), np.arange(6).astype(int), "full"),
+            (
+                np.random.normal(size=4).astype(config.floatX),
+                np.random.uniform(size=8).astype(int),
+                "valid",
+            ),
+        ],
+    )
+    def test_op(self, a, v, mode):
+        aesara_sol = convolve(a, v, mode=mode).eval()
+        numpy_sol = np.convolve(a, v, mode=mode)
+        assert np.allclose(numpy_sol, aesara_sol)
+
+    def test_scalar_error(self):
+        with pytest.raises(ValueError, match="must be 1-dim"):
+            convolve(4, [4, 1])
+
+    @pytest.mark.parametrize(
+        "a_shape, v_shape, mode, error_regex",
+        [
+            ((2,), (3, 1), "full", "must be 1-dim"),
+            ((2,), (3,), "val", "must be full, valid or same"),
+        ],
+    )
+    def test_get_output_shape_error(self, a_shape, v_shape, mode, error_regex):
+        a = tensor(dtype=np.float64, shape=a_shape)
+        v = tensor(dtype=np.float64, shape=v_shape)
+
+        with pytest.raises(ValueError, match=error_regex):
+            Convolve._get_output_shape(a, v, (a_shape, v_shape), mode, validate=True)
+        assert (
+            Convolve._get_output_shape(a, v, (a_shape, v_shape), mode, validate=False)
+            == ()
+        )
+
+    @pytest.mark.parametrize(
+        "a_shape, v_shape, a1_shape, v1_shape",
+        [
+            ((5,), (5,), (5,), (5,)),
+            ((1,), (5,), (1,), (5,)),
+            ((1,), (1,), (1,), (1,)),
+            ((3,), (1,), (3,), (1,)),
+            ((None,), (3,), (4,), (3,)),
+            ((None,), (None,), (4,), (2,)),
+            ((2,), (None,), (2,), (3,)),
+        ],
+    )
+    @pytest.mark.parametrize(
+        "mode",
+        [
+            "full",
+            "valid",
+            "same",
+        ],
+    )
+    def test_infer_shape(self, a_shape, v_shape, a1_shape, v1_shape, mode):
+        a = tensor(dtype=config.floatX, shape=a_shape)
+        v = tensor(dtype=config.floatX, shape=v_shape)
+
+        rng = np.random.default_rng(utt.fetch_seed())
+        a1 = rng.random(a1_shape).astype(config.floatX)
+        v1 = rng.random(v1_shape).astype(config.floatX)
+
+        self._compile_and_check(
+            [a, v],
+            [convolve(a, v, mode)],
+            [a1, v1],
+            Convolve,
+        )