CODEWITHSUNDEEP
tensorflow
kPow989
kPow989

Reputation: 426

Custom Merge Function for Different Size Tensors in Tensorflow

I have two tensors of different sizes and want to write a custom merge function 

a = tf.constant([[1,2,3]])
b = tf.constant([[1,1,2,2,3,3]])

I want to take the dot product of each point in tensor a with two points in tensor b. So in the example above element 1 in a is multiplied with the first two elements in b and so on. I'm unsure of how to do loops in tensorflow:

def customMergeFunct(x):
    # not sure how to write a loop over a tensor

The output should be:

c = Lambda(customMergeFunct)([a,b])

with tf.Session() as sess: 
    print(c.eval())

=> [[2,8,18]]

Upvotes: 0

Views: 251

Answers (2)

Kaushik Roy
Kaushik Roy

Reputation: 1685

You can do something following:

a = tf.constant([[1,2,3]])  # Shape: (1, 3)
b = tf.constant([[1,1,2,2,3,3]])  # Shape: (1, 6)


def customMergeFunct(x):

    # a_ = tf.tile(x[0], [2, 1]) # Duplicating 2 times (Original)  # Update: No need of doing this as tf.multiply will use Broadcasting
    b_ = tf.transpose(tf.reshape(x[1], [-1, 2]))  # reshaping followed by transpose to make a shape of (2, 3) to serve the purpose + multiplication rule
    return tf.reduce_sum(tf.multiply(x[0], b_), axis=0)  # Element-wise multiplication followed by sum

# Using function
c = Lambda(customMergeFunct)([a,b])
# OR in a reduced form
c = Lambda(lambda x: tf.reduce_sum(tf.multiply(x[0], tf.transpose(tf.reshape(x[1], [-1, 2]))), axis=0))([a,b])

Output:

with tf.Session() as sess:
    print(c.eval()) # Output: [2 8 18]

# OR in eager mode
print(c.numpy()) # Output: [2 8 18]

Updated solution is computationally efficient than the original solution as we don't actually need to apply tile on x[0]

Upvotes: 1

thushv89
thushv89

Reputation: 11333

I'm not exactly sure why you call this a merge function. You don't really need to define a custom function. You can do this with a simple lambda function. Here's my solution.

import tensorflow as tf
from tensorflow.keras.layers import Lambda
import tensorflow.keras.backend as K

a = tf.constant([[1,2,3]])
b = tf.constant([[1,1,2,2,3,3]])

a_res = tf.reshape(a,[-1,1]) # make a.shape [3,1]
b_res = tf.reshape(b,[-1,2]) # make b.shape [3,2]

layer = Lambda(lambda x: K.sum(x[0]*x[1],axis=1))

res = layer([a_res,b_res])

with tf.Session() as sess: 
    print(res.eval())

Upvotes: 2

Related Questions