Got numerical difference between two implementations

Question

I've been working around RetNet (Paper: https://arxiv.org/pdf/2307.08621, PyTorch implementation: https://github.com/Jamie-Stirling/RetNet/). I rewrote the some of the code with TensorFlow:

import tensorflow as tf
from keras import layers
from xpos import XPOS


@tf.keras.utils.register_keras_serializable('RetNet')
class SimpleRetention(layers.Layer):
    def __init__(self, hidden_size, gamma, head_size=None, double_v_dim=False):
        super(SimpleRetention, self).__init__()

        self.hidden_size = hidden_size
        head_size = head_size if head_size is not None else hidden_size
        self.head_size = head_size
        self.v_dim = head_size * 2 if double_v_dim else head_size
        self.gamma = gamma

        self.W_Q = layers.Dense(head_size, use_bias=False)
        self.W_K = layers.Dense(head_size, use_bias=False)
        self.W_V = layers.Dense(self.v_dim, use_bias=False)

        self.xpos = XPOS(head_size)

    def call(self, X):
        Q = self.W_Q(X)
        K = self.W_K(X)
        V = self.W_V(X)

        Q = self.xpos(Q)
        K = self.xpos(K, downscale=True)

        ret = tf.matmul(Q, K, transpose_b=True) * self._get_D(tf.shape(X)[1])
        return tf.matmul(ret, V)

    def call_recurrent(self, x_n, s_n_1, n):
        Q = self.W_Q(x_n)
        K = self.W_K(x_n)
        V = self.W_V(x_n)

        Q = self.xpos(Q, offset=n+1)
        K = self.xpos(K, offset=n+1, downscale=True)

        s_n = self.gamma * s_n_1 + tf.matmul(K, V, transpose_a=True)
        return tf.matmul(Q, s_n), s_n

    def _get_D(self, sequence_length):
        n = tf.range(sequence_length)[:, tf.newaxis]
        m = tf.range(sequence_length)[tf.newaxis, :]
        D = tf.pow(self.gamma, tf.cast(n - m, dtype=tf.float32))
        mask = tf.cast(n >= m, dtype=tf.float32)  # causal mask
        return D * mask

Now I've written another test case to validate the outputs from call and call_recurrent are identical, following the test case:

import tensorflow as tf
from retentive import SimpleRetention
import numpy as np
import matplotlib.pyplot as plt


def test_simple():
    """
    Verify that the three implementations of SimpleRetention are identical
    """
    batch_size = 4
    sequence_length = 12
    hidden_size = 6

    gamma = 0.9

    X = tf.random.uniform((batch_size, sequence_length, hidden_size))
    sr = SimpleRetention(hidden_size, gamma, double_v_dim=True)

    Y_parallel = sr(X)

    s_n_1 = tf.zeros((batch_size, hidden_size, sr.v_dim))
    Y_recurrent = []
    for i in range(sequence_length):
        y_n, s_n = sr.call_recurrent(X[:, i:i+1, :], s_n_1, i)
        Y_recurrent.append(y_n)
        s_n_1 = s_n

    Y_recurrent = tf.concat(Y_recurrent, axis=1)

    Yp = Y_parallel.numpy()
    Yr = Y_recurrent.numpy()
    print(Yp.shape, Yr.shape)

    plt.figure()
    for i in range(batch_size):
        plt.subplot(batch_size // 2, 2, i+1)
        error_map = Yp[i] - Yr[i]
        plt.imshow(error_map)
        plt.title(f"sq error={np.sum(error_map * error_map)}")
    plt.show()

    error = np.sum(np.abs(Yp - Yr))
    print(f"Error: {error}")


test_simple()

The error should be really close to 0.0, but I got relatively large (absolute) error like 1.27 overall. The error map was like

I'm not sure if this is a normal case, and if not, do anyone know what might be the cause? Will the numerical differences hurt my model's performance?