How to merge two dimensions of a tensor without assuming the shape of it

user_866082 0
Viewed 529

Let's assume the following function:

from tensorflow.python.keras import backend as K

def broadcast_sum(a, b):
    a = K.expand_dims(a, 1)
    b = K.expand_dims(b, 2)
    c = a + b
    cs = K.shape(c)
    return K.reshape(c, (cs[0], -1, cs[-1]))

Given the two tensors of shapes (1, 3, 2) and (1, 4, 2), it correctly returns:

>>> broadcast_sum(K.placeholder((1, 3, 2)), K.placeholder((1, 4, 2)))
>>> <tf.Tensor 'Reshape_2:0' shape=(1, 12, 2) dtype=float32>

Right now, this function works only with 3D input (because of the reshape line). My question is, how can I make this work with any shape (using the same function) without knowing the shape? Of course, I'm assuming the inputs are of the same shape and at least a 3D. But how can I have a single function that works with 3D, 4D and so on?

And I'm assuming that it's always the second dimension (from left) that the function will broadcast and the rest of the dimensions are identical between the two inputs. Here are the shapes that I want to make the same function to work with:

>>> broadcast_sum(K.placeholder((1, 3, 5, 2)), K.placeholder((1, 4, 5, 2)))
>>> <tf.Tensor 'Reshape_3:0' shape=(1, 60, 2) dtype=float32>

Of course, the returned tensor is wrong right now. It should be of shape (1, 12, 5, 2).

[UPDATE]

Please also consider that the first dimension (the batch size) could be None. In fact, any of the dimensions except the rightmost one could be None.

tensor keras tensorflow
1 Answers

And I'm assuming that it's always the second dimension (from left) that the function will broadcast and the rest of the dimensions are identical between the two inputs.

Based on this, I reuse the shape information from one of the inputs.

from tensorflow.python.keras import backend as K
def broadcast_sum(a, b):
    final_shape = (a.shape[0], -1, *a.shape[2:])
    a = K.expand_dims(a, 1)
    b = K.expand_dims(b, 2)
    c = a + b
    return K.reshape(c, final_shape)


print(broadcast_sum(K.placeholder((1, 3, 2)), K.placeholder((1, 4, 2))))
print(broadcast_sum(K.placeholder((1, 3, 5, 2)), K.placeholder((1, 4, 5, 2))))

Tensor("Reshape:0", shape=(1, 4, 3, 2), dtype=float32)
Tensor("Reshape_1:0", shape=(1, 12, 5, 2), dtype=float32)
user_55977180
Related