Practice Now
Keras
Dot layer
broadcasting
neural networks
TensorFlow

Keras -- no Dot layer with broadcasting?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
TensorFlow & Keras

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
Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.