Keras -- no Dot layer with broadcasting?

Keras

Dot layer

broadcasting

neural networks

TensorFlow

Keras -- no Dot layer with broadcasting?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Start Practicing Learn More

Introduction

Keras's Dot layer computes dot products along a specified axis but does not support broadcasting — both inputs must have the same batch size and compatible shapes. When you need a dot product with broadcasting (e.g., multiplying a batch of vectors by a single weight vector, or computing attention scores between tensors of different ranks), use tf.keras.layers.Lambda with tf.tensordot, tf.einsum, or the Multiply layer combined with tf.reduce_sum. These alternatives give you full control over which axes are contracted and how broadcasting is applied.

The Problem with Dot Layer

python

1import tensorflow as tf
2
3# Dot layer requires matching dimensions
4a = tf.keras.Input(shape=(10, 64))   # (batch, seq, features)
5b = tf.keras.Input(shape=(64,))      # (batch, features) — different rank
6
7# This fails — Dot cannot broadcast across different ranks
8try:
9    output = tf.keras.layers.Dot(axes=-1)([a, b])
10except ValueError as e:
11    print(e)
12# "inputs should have the same number of dimensions"

The Dot layer requires both inputs to have the same number of dimensions. It cannot broadcast a 2D tensor across a 3D tensor.

Solution 1: Lambda Layer with tf.einsum

tf.einsum is the most flexible and readable option:

python

1a = tf.keras.Input(shape=(10, 64))   # (batch, seq, features)
2b = tf.keras.Input(shape=(64,))      # (batch, features)
3
4# Einstein summation — contract over the features dimension
5output = tf.keras.layers.Lambda(
6    lambda x: tf.einsum('bsf,bf->bs', x[0], x[1])
7)([a, b])
8# output shape: (batch, 10)
9
10model = tf.keras.Model(inputs=[a, b], outputs=output)
11model.summary()

einsum('bsf,bf->bs') means: for each batch b, multiply the (seq, features) matrix by the (features,) vector, producing a (seq,) result.

Common einsum Patterns for Broadcasting Dot Products

python

1# (batch, m, k) dot (batch, k) -> (batch, m)
2tf.einsum('bmk,bk->bm', x, y)
3
4# (batch, m, k) dot (k, n) -> (batch, m, n)  — weight matrix, no batch dim
5tf.einsum('bmk,kn->bmn', x, W)
6
7# (batch, heads, seq, dim) dot (batch, heads, dim, seq) -> attention scores
8tf.einsum('bhsd,bhds->bhss', queries, keys_transposed)

Solution 2: Multiply + Reduce Sum

Element-wise multiplication with broadcasting followed by sum reduction:

python

1a = tf.keras.Input(shape=(10, 64))   # (batch, seq, features)
2b = tf.keras.Input(shape=(64,))      # (batch, features)
3
4# Expand b to (batch, 1, 64) for broadcasting
5b_expanded = tf.keras.layers.Reshape((1, 64))(b)
6
7# Element-wise multiply: (batch, 10, 64) * (batch, 1, 64) -> (batch, 10, 64)
8multiplied = tf.keras.layers.Multiply()([a, b_expanded])
9
10# Sum over features axis: (batch, 10, 64) -> (batch, 10)
11output = tf.keras.layers.Lambda(lambda x: tf.reduce_sum(x, axis=-1))(multiplied)
12
13model = tf.keras.Model(inputs=[a, b], outputs=output)

This is equivalent to the dot product but uses broadcasting explicitly.

Solution 3: Custom Layer

For reuse across models, create a custom layer:

python

1class BroadcastDot(tf.keras.layers.Layer):
2    """Dot product with broadcasting support."""
3
4    def __init__(self, axes=-1, **kwargs):
5        super().__init__(**kwargs)
6        self.axes = axes
7
8    def call(self, inputs):
9        a, b = inputs
10        return tf.einsum('...i,i->...', a, b) if len(b.shape) == 1 else \
11               tf.einsum('...i,...i->...', a, b)
12
13    def get_config(self):
14        config = super().get_config()
15        config.update({"axes": self.axes})
16        return config
17
18# Usage
19a = tf.keras.Input(shape=(10, 64))
20b = tf.keras.Input(shape=(64,))
21output = BroadcastDot()([a, b])  # (batch, 10)

Solution 4: tf.tensordot in a Lambda

python

1a = tf.keras.Input(shape=(10, 64))
2b = tf.keras.Input(shape=(64,))
3
4# tensordot contracts specified axes
5output = tf.keras.layers.Lambda(
6    lambda x: tf.tensordot(x[0], x[1], axes=[[-1], [-1]])
7)([a, b])
8# Careful: tensordot does not preserve batch dimension automatically
9# Use einsum instead for batched operations

tf.tensordot is powerful but does not handle the batch dimension implicitly. For batched operations, tf.einsum is usually a better choice.

Attention Score Example

A common use case is computing attention scores between queries and keys of different shapes:

python

1# Scaled dot-product attention
2seq_len = 50
3d_model = 64
4
5queries = tf.keras.Input(shape=(seq_len, d_model))
6keys = tf.keras.Input(shape=(seq_len, d_model))
7
8# Dot product: (batch, seq, dim) x (batch, dim, seq) -> (batch, seq, seq)
9scores = tf.keras.layers.Lambda(
10    lambda x: tf.einsum('bsd,btd->bst', x[0], x[1]) / (d_model ** 0.5)
11)([queries, keys])
12
13# Apply softmax
14attention_weights = tf.keras.layers.Softmax()(scores)
15
16model = tf.keras.Model(inputs=[queries, keys], outputs=attention_weights)
17print(model.output_shape)  # (None, 50, 50)

When to Use the Built-in Dot Layer

The Dot layer works fine when both inputs have the same rank:

python

1a = tf.keras.Input(shape=(10, 64))
2b = tf.keras.Input(shape=(10, 64))
3
4# Same rank — Dot works
5output = tf.keras.layers.Dot(axes=-1)([a, b])
6# output shape: (batch, 10, 10)
7
8# Or contract along axis 1 (sequence)
9output = tf.keras.layers.Dot(axes=1)([a, b])
10# output shape: (batch, 64, 64)

Common Pitfalls

Assuming Dot supports broadcasting: The Dot layer requires both inputs to have the same number of dimensions. For different ranks, use einsum, Multiply + reduce_sum, or a custom layer.
Losing the batch dimension with tf.tensordot: tf.tensordot contracts all specified axes including batch if you are not careful. Prefer tf.einsum for batched operations because it explicitly names all dimensions.
Wrong axis in einsum string: A typo in the einsum subscript silently produces wrong shapes. Print output.shape after each einsum operation to verify correctness.
Not using get_config in custom layers: Custom layers without get_config cannot be serialized (saved/loaded). Always implement get_config in custom Keras layers.
Using Lambda layers for complex logic: Lambda layers are hard to serialize and debug. For anything beyond a simple one-liner, create a proper custom layer subclassing tf.keras.layers.Layer.

Summary

Keras Dot layer does not support broadcasting between tensors of different ranks
Use tf.einsum in a Lambda layer for the most readable and flexible broadcasting dot products
Use Multiply + Reshape + reduce_sum for explicit element-wise broadcasting
Create a custom BroadcastDot layer for reusable broadcasting dot products
Prefer tf.einsum over tf.tensordot for batched operations to preserve the batch dimension