test_random_matmul.py

import numpy as np
import tensorflow as tf


def test_random_embeddings():
    query_embeddings = tf.random.uniform((16, 48)).numpy()
    item_embeddings = tf.random.uniform((100, 48)).numpy()

    with tf.device("/gpu:0"):
        preds_gpu = tf.matmul(tf.constant(query_embeddings), tf.transpose(tf.constant(item_embeddings)))

    with tf.device("/cpu:0"):
        preds_cpu = tf.matmul(tf.constant(query_embeddings), tf.transpose(tf.constant(item_embeddings)))

    np.testing.assert_allclose(preds_gpu.numpy(), preds_cpu.numpy(), atol=1e-4)

if __name__ == "__main__":
    test_random_embeddings()