TensorFlow
Keras
build method
custom layers
neural networks

How does Tensorflow build work from tf.keras.layers.Layer

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

Introduction

In a custom Keras layer, build(input_shape) is where you create weights that depend on the input shape. You usually do not call it yourself. Keras calls it lazily the first time the layer sees input through __call__, once the input shape is known.

Why build() Exists

When you instantiate a layer, you often do not yet know the final input dimension. For example, a custom dense-style layer may need to know the last dimension of the incoming tensor before it can allocate a kernel matrix.

That is why Keras separates layer setup into two parts:

  • '__init__ for configuration that does not depend on input shape'
  • 'build(input_shape) for weights and state that do depend on input shape'

This keeps layer construction flexible and avoids forcing users to specify shapes manually up front.

The Lifecycle

The usual lifecycle is:

  1. create the layer object
  2. pass a tensor into the layer
  3. Keras sees the layer is not built yet
  4. Keras calls build(input_shape)
  5. weights are created
  6. Keras calls call(inputs)

After that, the layer is marked as built and build() is not run again for normal reuse.

A Minimal Example

python
1import tensorflow as tf
2
3
4class SimpleDense(tf.keras.layers.Layer):
5    def __init__(self, units):
6        super().__init__()
7        self.units = units
8
9    def build(self, input_shape):
10        input_dim = input_shape[-1]
11        self.kernel = self.add_weight(
12            shape=(input_dim, self.units),
13            initializer="glorot_uniform",
14            trainable=True,
15            name="kernel",
16        )
17        self.bias = self.add_weight(
18            shape=(self.units,),
19            initializer="zeros",
20            trainable=True,
21            name="bias",
22        )
23
24    def call(self, inputs):
25        return tf.matmul(inputs, self.kernel) + self.bias
26
27
28layer = SimpleDense(4)
29x = tf.ones((2, 3))
30y = layer(x)
31
32print(y.shape)
33print(layer.built)

When layer(x) runs for the first time, Keras infers input_shape=(2, 3) and calls build() before call().

What Belongs in __init__ Versus build()

Put shape-independent configuration in __init__:

  • number of units
  • activation choice
  • regularization settings

Put shape-dependent state in build():

  • weights whose dimensions depend on the input
  • lookup tables tied to input width
  • other variables created with add_weight

If you create shape-dependent weights in __init__, you usually end up hardcoding assumptions that make the layer less reusable.

Do You Need To Call super().build(...)?

For plain custom layers, calling super().build(input_shape) at the end is a good habit because it marks the layer as built in the standard way:

python
def build(self, input_shape):
    # create weights
    super().build(input_shape)

Keras often handles built-state correctly either way when using add_weight, but calling the superclass method makes your intent explicit and aligns with the expected lifecycle.

What Happens If Input Shape Changes Later

A built layer is generally expected to keep the same weight shapes. If you feed the same layer inputs with incompatible shapes later, Keras will usually raise a shape error rather than rebuilding the weights.

That is why build() is for one-time initialization, not for dynamically reshaping a layer on every call.

Common Pitfalls

The biggest mistake is doing all weight creation in call(). That can lead to duplicated variables and hard-to-debug tracing behavior.

Another mistake is putting input-shape-dependent logic into __init__, where the necessary shape information is not available yet.

A third issue is manually calling build() without a strong reason. In typical Keras usage, calling the layer on input is the correct trigger.

Summary

  • 'build(input_shape) creates weights that depend on the input shape.'
  • Keras calls it lazily the first time the layer receives input.
  • Use __init__ for configuration and build() for shape-dependent weights.
  • Create variables with self.add_weight(...), then implement forward logic in call(...).
  • A custom layer is generally built once and then reused with compatible input shapes.

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

All Rights Reserved.