How to properly connect Dense layer to Lambda layer in Keras?

Question

I'm using Keras. In the following code, model takes [a0, a1], [b0, b1, b2] as inputs and gives [a0*b0, a0*b1, a0*b2, a1*b0, a1*b1, a1*b2] as an output:

from keras import backend as K
from keras.models import Model
from keras.models import Input
from keras.layers import Dense

def mix(ts):
    t0 = K.expand_dims(ts[0], axis=-1)
    t1 = K.expand_dims(ts[1], axis=1)
    return K.batch_flatten(t0 * t1)

a = Input(shape=(2,))
b = Input(shape=(3,))
c = Lambda(mix)([a, b])

model = Model(inputs=[a,b], outputs=c)

Here's the test:

u = np.array([1,2]).reshape(1,2)
v = np.array([3,4,5]).reshape(1,3)
print(model.predict([u,v]))

[[ 3. 4. 5. 6. 8. 10.]]

But if I try to connect Dense layer to Lambda layer, I get an error:

from keras import backend as K
from keras.models import Model
from keras.models import Input
from keras.layers import Dense

def mix(ts):
    t0 = K.expand_dims(ts[0], axis=-1)
    t1 = K.expand_dims(ts[1], axis=1)
    return K.batch_flatten(t0 * t1)

a = Input(shape=(2,))
b = Input(shape=(3,))
c = Lambda(mix)([a, b])
d = Dense(2)(c)

model = Model(inputs=[a,b], outputs=d)

Here's the error I get:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-6-0f7f977a1e79> in <module>()
      7 b = Input(shape=(3,))
      8 c = Lambda(mix)([a, b])
----> 9 d = Dense(2)(c)
     10 
     11 model = Model(inputs=[a,b], outputs=d)

~\Anaconda3\envs\mind\lib\site-packages\keras\engine\base_layer.py in __call__(self, inputs, **kwargs)
    429                                          'You can build it manually via: '
    430                                          '`layer.build(batch_input_shape)`')
--> 431                 self.build(unpack_singleton(input_shapes))
    432                 self.built = True
    433 

~\Anaconda3\envs\mind\lib\site-packages\keras\layers\core.py in build(self, input_shape)
    864                                       name='kernel',
    865                                       regularizer=self.kernel_regularizer,
--> 866                                       constraint=self.kernel_constraint)
    867         if self.use_bias:
    868             self.bias = self.add_weight(shape=(self.units,),

~\Anaconda3\envs\mind\lib\site-packages\keras\legacy\interfaces.py in wrapper(*args, **kwargs)
     89                 warnings.warn('Update your `' + object_name + '` call to the ' +
     90                               'Keras 2 API: ' + signature, stacklevel=2)
---> 91             return func(*args, **kwargs)
     92         wrapper._original_function = func
     93         return wrapper

~\Anaconda3\envs\mind\lib\site-packages\keras\engine\base_layer.py in add_weight(self, name, shape, dtype, initializer, regularizer, trainable, constraint)
    247         if dtype is None:
    248             dtype = K.floatx()
--> 249         weight = K.variable(initializer(shape),
    250                             dtype=dtype,
    251                             name=name,

~\Anaconda3\envs\mind\lib\site-packages\keras\initializers.py in __call__(self, shape, dtype)
    207             scale /= max(1., fan_out)
    208         else:
--> 209             scale /= max(1., float(fan_in + fan_out) / 2)
    210         if self.distribution == 'normal':
    211             # 0.879... = scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)

TypeError: unsupported operand type(s) for +: 'NoneType' and 'int'

How can I properly connect Dense layer to Lambda layer?

Answer 1

You need to set the output shape of the Lambda layer in this case, since it cannot be inferred automatically. Either pass the output_shape manually:

c = Lambda(mix, output_shape=(6,))([a, b])

Or better than that, pass a function to compute the output shape based on the shapes of input tensors of the layer:

def mix_output_shape(input_shape):
    # input_shape[0] is the shape of first input tensor
    # input_shape[1] is the shape of second input tensor
    return (input_shape[0][0], input_shape[0][1] * input_shape[1][1])

# ...
c = Lambda(mix, mix_output_shape)([a, b])