What is the efficient way to write this custom operation in tensorflow?

I want to implement a new operation similar to matrix multiplication in tensorflow. For example, If there are two matrices input, W and C = Op(input,W), then, C[i,j] = tf.math.reduce_prod(input[i,:] - W[:,j]). For implementing this operation I am planning to use two "for-loops" iterating over rows & columns, however, I believe that the for-loop implementation will decrease the training performance. I tried looking into the implementation of tf.lingalg.matmul() but the function seems to call C++ binaries which are not editable. Is there any way to speed up this computation? The following is the function I have written to implement Op(input,w):

def difference(input,w):
    result = tf.Variable(tf.zeros((input.shape[0],w.shape[1]),dtype=tf.float16))
    for i in tf.range(0,input.shape[0]):
        for j in tf.range(0,w.shape[1]):
            result[i,j].assign(tf.math.reduce_prod(input[i,:] - W[:,j]))
    return result