From 895cd3ae1d93a6e2770788fbd60632e88a1c5952 Mon Sep 17 00:00:00 2001
From: Filippo Miatto <miatto@gmail.com>
Date: Fri, 1 Nov 2024 11:10:30 -0700
Subject: [PATCH 1/2] convolution added to numpy backend

---
 mrmustard/math/backend_manager.py |  3 +-
 mrmustard/math/backend_numpy.py   | 54 +++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/mrmustard/math/backend_manager.py b/mrmustard/math/backend_manager.py
index cdaafda02..7f20529d5 100644
--- a/mrmustard/math/backend_manager.py
+++ b/mrmustard/math/backend_manager.py
@@ -1527,8 +1527,7 @@ def all_diagonals(self, rho: Tensor, real: bool) -> Tensor:
 
     def poisson(self, max_k: int, rate: Tensor) -> Tensor:
         """Poisson distribution up to ``max_k``."""
-        k = self.arange(max_k)
-        rate = self.cast(rate, k.dtype)
+        k = self.arange(max_k, dtype=rate.dtype)
         return self.exp(k * self.log(rate + 1e-9) - rate - self.lgamma(k + 1.0))
 
     def binomial_conditional_prob(self, success_prob: Tensor, dim_out: int, dim_in: int):
diff --git a/mrmustard/math/backend_numpy.py b/mrmustard/math/backend_numpy.py
index 18bffa73b..b2ab76694 100644
--- a/mrmustard/math/backend_numpy.py
+++ b/mrmustard/math/backend_numpy.py
@@ -23,6 +23,7 @@
 
 import numpy as np
 import scipy as sp
+from scipy.signal import convolve2d as scipy_convolve2d
 from scipy.linalg import expm as scipy_expm
 from scipy.linalg import sqrtm as scipy_sqrtm
 from scipy.special import xlogy as scipy_xlogy
@@ -136,6 +137,59 @@ def concat(self, values: list[np.ndarray], axis: int) -> np.ndarray:
     def conj(self, array: np.ndarray) -> np.ndarray:
         return np.conj(array)
 
+    def convolution(
+        self,
+        array: np.ndarray,  # shape: [width, in_channels]
+        filters: np.ndarray,  # shape: [kernel_width, in_channels, out_channels]
+        padding: str = "VALID",
+        data_format: str | None = None,
+    ) -> np.ndarray:  # returns: [width, out_channels]
+        """Performs 2D convolution operation similar to tf.nn.convolution using numpy.
+
+        Args:
+            array: Input array of shape (batch, height, width, channels)
+            filters: Filter kernel of shape (kernel_height, kernel_width, in_channels, out_channels)
+            padding: String indicating the padding type ('VALID' or 'SAME')
+            data_format: Unused, kept for API compatibility
+
+        Returns:
+            np.ndarray: Result of the convolution operation with shape (batch, new_height, new_width, out_channels)
+        """
+        # Extract shapes
+        batch, in_height, in_width, in_channels = array.shape
+        kernel_h, kernel_w, _, out_channels = filters.shape
+
+        # Reshape filter to 2D for convolution
+        filter_2d = filters[:, :, 0, 0]
+
+        # For SAME padding, calculate padding sizes
+        if padding == "SAME":
+            pad_h = (kernel_h - 1) // 2
+            pad_w = (kernel_w - 1) // 2
+            array = np.pad(
+                array[:, :, :, 0], ((0, 0), (pad_h, pad_h), (pad_w, pad_w)), mode="constant"
+            )
+        else:
+            array = array[:, :, :, 0]
+
+        # Calculate output dimensions
+        out_height = array.shape[1] - kernel_h + 1
+        out_width = array.shape[2] - kernel_w + 1
+
+        # Initialize output array
+        output = np.zeros((batch, out_height, out_width, out_channels))
+
+        # Perform convolution for each batch
+        for b in range(batch):
+            # Convolve using scipy's convolve2d which is more efficient than np.convolve for 2D
+            output[b, :, :, 0] = scipy_convolve2d(
+                array[b],
+                np.flip(np.flip(filter_2d, 0), 1),  # Flip kernel for proper convolution
+                mode="valid",
+            )
+
+        return output
+
     def cos(self, array: np.ndarray) -> np.ndarray:
         return np.cos(array)
 

From b264136a09983736f4ee4ade4b0f19e4d65734bd Mon Sep 17 00:00:00 2001
From: Filippo Miatto <miatto@gmail.com>
Date: Tue, 5 Nov 2024 10:02:15 -0800
Subject: [PATCH 2/2] fixing codefactor complaints

---
 mrmustard/math/backend_numpy.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mrmustard/math/backend_numpy.py b/mrmustard/math/backend_numpy.py
index b2ab76694..68b695432 100644
--- a/mrmustard/math/backend_numpy.py
+++ b/mrmustard/math/backend_numpy.py
@@ -142,7 +142,7 @@ def convolution(
         array: np.ndarray,  # shape: [width, in_channels]
         filters: np.ndarray,  # shape: [kernel_width, in_channels, out_channels]
         padding: str = "VALID",
-        data_format: str | None = None,
+        data_format: str | None = None,  # pylint: disable=unused-argument
     ) -> np.ndarray:  # returns: [width, out_channels]
         """Performs 2D convolution operation similar to tf.nn.convolution using numpy.
 
@@ -156,7 +156,7 @@ def convolution(
             np.ndarray: Result of the convolution operation with shape (batch, new_height, new_width, out_channels)
         """
         # Extract shapes
-        batch, in_height, in_width, in_channels = array.shape
+        batch, _, _, _ = array.shape
         kernel_h, kernel_w, _, out_channels = filters.shape
 
         # Reshape filter to 2D for convolution