Fix bug with correlate for tensorflow

keras/src/backend/tensorflow/numpy.py

@@ -3130,30 +3130,47 @@ def correlate(x1, x2, mode="valid"):
     x1 = tf.cast(x1, dtype)
     x2 = tf.cast(x2, dtype)
 
-    x1_len, x2_len = int(x1.shape[0]), int(x2.shape[0])
-
-    if mode == "full":
-        full_len = x1_len + x2_len - 1
-
-        x1_pad = (full_len - x1_len) / 2
-        x2_pad = (full_len - x2_len) / 2
-
-        x1 = tf.pad(
-            x1, paddings=[[tf.math.floor(x1_pad), tf.math.ceil(x1_pad)]]
-        )
-        x2 = tf.pad(
-            x2, paddings=[[tf.math.floor(x2_pad), tf.math.ceil(x2_pad)]]
+    def _pack(a, b):
+        # a: input [N] -> [1,N,1];
+        # b: filter [M] -> [M,1,1]
+        return (
+            tf.reshape(a, (1, tf.shape(a)[0], 1)),
+            tf.reshape(b, (tf.shape(b)[0], 1, 1)),
         )
 
-        x1 = tf.reshape(x1, (1, full_len, 1))
-        x2 = tf.reshape(x2, (full_len, 1, 1))
+    def _full_corr(x1, x2):
+        """Compute 'full' correlation result (length = n + m - 1)."""
+        m = tf.shape(x2)[0]
+        pad = tf.maximum(m - 1, 0)
+        x1 = tf.pad(x1, [[pad, pad]])  # pad input with zeros
+        x1, x2 = _pack(x1, x2)
+        out = tf.nn.conv1d(x1, x2, stride=1, padding="VALID")
+        return tf.squeeze(out)
 
-        return tf.squeeze(tf.nn.conv1d(x1, x2, stride=1, padding="SAME"))
+    n = tf.shape(x1)[0]
+    m = tf.shape(x2)[0]
 
-    x1 = tf.reshape(x1, (1, x1_len, 1))
-    x2 = tf.reshape(x2, (x2_len, 1, 1))
-
-    return tf.squeeze(tf.nn.conv1d(x1, x2, stride=1, padding=mode.upper()))
+    if mode == "full":
+        return _full_corr(x1, x2)
+    elif mode == "same":
+        # unfortunately we can't leverage 'SAME' padding directly like
+        # we can with "valid"
+        # it works fine for odd-length filters, but for even-length filters
+        # the output is off by 1 compared to numpy, due to how
+        # tf handles centering
+        full_corr = _full_corr(x1, x2)
+        full_len = n + m - 1
+        out_len = tf.maximum(n, m)
+        start = (full_len - out_len) // 2
+        return tf.slice(full_corr, [start], [out_len])
+    elif mode == "valid":
+        x1, x2 = _pack(x1, x2)
+        return tf.squeeze(tf.nn.conv1d(x1, x2, stride=1, padding="VALID"))
+    else:
+        raise ValueError(
+            f"Invalid mode: '{mode}'. Mode must be one of:"
+            f" 'full', 'same', 'valid'."
+        )
 
 
 def select(condlist, choicelist, default=0):

keras/src/ops/numpy_test.py

@@ -5133,6 +5133,26 @@ def test_correlate_different_size(self):
             knp.Correlate(mode="full")(x, y), np.correlate(x, y, mode="full")
         )
 
+    def test_correlate_bug(self):
+        # copied from https://github.com/keras-team/keras/issues/21617
+        x = np.array([1, 3, 5])
+        y = np.array([7, 9])
+        self.assertAllClose(knp.correlate(x, y), np.correlate(x, y))
+        self.assertAllClose(
+            knp.correlate(x, y, mode="same"), np.correlate(x, y, mode="same")
+        )
+        self.assertAllClose(
+            knp.correlate(x, y, mode="full"), np.correlate(x, y, mode="full")
+        )
+
+        self.assertAllClose(knp.Correlate()(x, y), np.correlate(x, y))
+        self.assertAllClose(
+            knp.Correlate(mode="same")(x, y), np.correlate(x, y, mode="same")
+        )
+        self.assertAllClose(
+            knp.Correlate(mode="full")(x, y), np.correlate(x, y, mode="full")
+        )
+
     def test_select(self):
         x = np.arange(6)
         condlist = [x < 3, x > 3]