Broadcasting dynamic dimension in Tensorflow

Viewed 185

I am using Tensorflow from python. I have two tensors I wish to concatenate (it could also be another operation, I don't think the exact operation matters to this question). These tensors have their shapes defined as (N1 != N2 are positive integers):

a: (None, N1)
b: (1   , N2)

Since I will be concatenating along the last axis, it seems like this operation could be performed. But tensorflow refues. The code

from tensorflow import keras
from tensorflow.keras import layers

N1 = 2
N2 = 3
D1 = None

a = keras.Input(shape=(D1, N1))
b = keras.Input(shape=(1, N2))

c = layers.Concatenate(axis=-1)([a, b])

fails with

ValueError: A `Concatenate` layer requires inputs with matching shapes except for the concat axis. Got inputs shapes: [(None, None, 2), (None, 1, 3)]

The closest I have come to making this work is by using RepeatVector as below, but unfortunately, this only works with static dimensions, not dynamic ones:

N1 = 2
N2 = 3
D1 = 7

a = keras.Input(shape=(D1, N1))
b = keras.Input(shape=(N2))
b_repeated = layers.RepeatVector(D1)(b)

c = layers.Concatenate()([a, b_repeated])

Any suggestions of how to concatenate -- i.e. do the right broadcasting or repeating -- with such None dimensions would be much appreciated!

1 Answers

Here is a way to do that with a lambda layer:

import keras
from keras import layers
import keras.backend as K

N1 = 2
N2 = 3
D1 = None

a = keras.Input(shape=(D1, N1))
b = keras.Input(shape=(N2,))
c = layers.Lambda(lambda ab: K.concatenate([ab[0], K.repeat(ab[1], K.shape(ab[0])[1])],
                                           axis=-1))([a, b])
print(c)
# Tensor("lambda_1/concat:0", shape=(?, ?, 5), dtype=float32)
Related