Custom Layers in Python

Every layer in tf.keras.layers — Dense, Conv2D, LSTM — ultimately boils down to the same shape: take some input tensors, apply a set of learnable weights, and produce an output tensor. A Dense layer specifically computes output = activation(dot(input, weights) + bias). That formula covers an enormous range of use cases, which is why it's the default building block for most networks.

But the built-in layers only cover the math someone already thought to implement. The moment you need something that isn't a standard affine transform followed by an activation — a custom attention mechanism, a physics-informed constraint, a layer that mixes two inputs with a learned gating function — the pre-built layers stop being enough. You can't bolt a new formula onto Dense; you have to build your own layer from scratch.

This is where Keras's subclassing API comes in. Instead of being limited to composing existing layers, you write a Python class that plugs directly into the same machinery Keras uses internally, so your custom math gets the same automatic differentiation, weight tracking, and model.fit() compatibility as any built-in layer.

To build a custom layer, you write a Python class that inherits from tf.keras.layers.Layer — not tf.keras.Model, which is reserved for composing whole networks, and not raw tf.Module, which lacks the Keras-specific bookkeeping (like automatic weight tracking and get_config serialization) that layers rely on.

The constructor, __init__, is where you store configuration values that don't depend on the shape of the incoming data — things like self.units, the number of output neurons this layer should produce. Critically, you must call super().__init__() first, because the parent Layer class sets up internal bookkeeping (like the list that will track your trainable weights) that your subclass depends on.

At this stage no actual weight tensors exist yet — __init__ only records intent. That distinction matters because Keras deliberately defers the creation of weight tensors to a separate method, build(), which only runs once the layer has seen the shape of its actual input.

It's worth pausing on why tf.keras.layers.Layer specifically is the required base class, since the three options look superficially similar. tf.keras.Model adds training loops, saving, and multi-layer orchestration on top of Layer — using it for a single mathematical operation is overkill and can confuse Keras about how the object should be treated when nested inside another model.

tf.Module is TensorFlow's lowest-level container for variables; it gives you variable tracking but none of the Keras-specific machinery — no build() lifecycle hook, no automatic handling of trainable/non_trainable weight lists that model.trainable_variables depends on, no compatibility with the Keras functional API.

tf.keras.layers.Layer sits in exactly the right spot: lightweight enough to represent a single operation, but wired into all the Keras conventions — build(), call(), weight tracking, serialization — that let your custom math drop into a Sequential or functional model exactly like Dense or Conv2D would.

Inside the class, you define a build(self, input_shape) method — this is where the layer's actual weight tensors get created. Keras calls build() automatically, exactly once, the first time data flows through the layer, and it passes in input_shape so you know the dimensions you're working with before allocating memory.

Inside build(), you call self.add_weight(...), specifying a shape, an initializer (like "random_normal" to start with small random values), and a trainable flag. add_weight doesn't just create a tf.Variable — it also registers that variable with the layer so Keras automatically includes it in layer.trainable_variables and in checkpoint saving.

Deferring weight creation to build() rather than __init__ is deliberate: at construction time you often don't know the shape of the data the layer will receive (e.g. how many features are in the input). By waiting until the first real batch arrives, input_shape[-1] is known, and the weight matrix can be sized correctly without the caller having to specify it manually.

A common early mistake is trying to call self.add_weight() inside __init__ instead of build(). It usually 'works' for a quick test because the shape happens to be known, but it breaks the layer's flexibility — the whole point of separating build() is that the same layer instance should be reusable across inputs of different shapes without you hardcoding a dimension.

build() also only runs once per layer instance, no matter how many times call() fires afterward. TensorFlow tracks whether the layer has already been built (via an internal self.built flag) and skips re-running build() on subsequent calls, so weight creation happens exactly one time even though the layer is invoked on every forward pass.

This single-invocation guarantee is also why weight shapes in a custom layer stay fixed after the first batch: if you later feed the layer an input with a different last-dimension size than what build() saw, TensorFlow will raise a shape-mismatch error rather than silently reallocating the weight matrix.

Finally, you define call(self, inputs). This is the forward pass — the method actually invoked every time data flows through the layer, whether that's once during a quick test or millions of times during training. In the simplest custom layer, call() might do nothing more than tf.matmul(inputs, self.w), a raw matrix multiplication against the weight tensor created in build().

Unlike build(), which fires only once, call() runs on every single forward pass — so whatever operations you write here directly determine both the layer's mathematical behavior and its runtime cost. Anything expensive placed in call() (like recreating a tensor from scratch) gets paid on every batch, which is why weight creation belongs in build() and only the actual math belongs in call().

Because call() is built from ordinary TensorFlow ops (tf.matmul, tf.nn.relu, elementwise arithmetic, and so on), it plugs directly into TensorFlow's automatic differentiation. You never write backpropagation by hand — as long as call() only uses differentiable TensorFlow operations, gradients with respect to self.w are computed automatically when you call .fit() or use a GradientTape.

It helps to keep the three methods straight by what question each one answers. __init__ answers 'what configuration does this layer need before it ever sees data?' — things like units, an activation name, or a dropout rate. build() answers 'now that I know the input's shape, what weight tensors do I need to allocate?'. call() answers 'given inputs and my existing weights, what output do I produce?'.

A useful way to internalize this split is to notice that __init__ and build() both run rarely (once each, at construction and at first use), while call() runs constantly — every forward pass, every training step, every inference request. That asymmetry is exactly why Keras separates configuration, state-creation, and computation into three distinct methods instead of one.

If you ever find yourself unsure which method a line of code belongs in, ask: does this line depend on the shape of the input? If yes, it belongs in build(). Does it produce the layer's actual output from inputs? If yes, it belongs in call(). Everything else — hyperparameters, flags, sub-layer construction that doesn't need input shape — belongs in __init__.

Before moving on, it's worth stress-testing your understanding of one of the most consequential decisions inside build(): whether a given weight should be trainable or not. This distinction is easy to gloss over, because both trainable and non-trainable weights are created the exact same way, with self.add_weight() — the only difference is a single boolean argument.

Get this wrong in a real model and the failure mode is quiet, not loud: nothing crashes, training just silently corrupts a value that was never meant to change, or a value that should have been learned stays frozen at its initial random value forever.

The next exercise walks through exactly this scenario: a custom layer that needs to track a running counter as data flows through it — a value that must survive across calls, but must never be treated as something gradient descent is allowed to optimize.

The trainable flag on self.add_weight() controls whether a variable is included in the set of tensors the Optimizer is allowed to touch during backpropagation. When trainable=True, the weight is added to layer.trainable_variables, gradients are computed with respect to it, and every call to optimizer.apply_gradients() nudges its value to reduce the loss.

When trainable=False, the variable still exists as a real tf.Variable — it still holds state, still persists across calls, still gets saved in checkpoints — but it's excluded from the gradient computation entirely. The Optimizer never sees it, so no matter how the loss changes, this variable's value only changes when you explicitly assign to it yourself (e.g. with .assign_add()).

This is precisely the mechanism Keras uses internally for things like BatchNormalization's running mean and variance: statistics that need to accumulate over the course of training, but that must never be treated as parameters the loss function gets to optimize directly.

Consider a custom layer that counts how many images have flowed through it — a plain integer counter implemented as self.count = self.add_weight(shape=(), initializer="zeros", trainable=False), incremented inside call() with something like self.count.assign_add(tf.cast(tf.shape(inputs)[0], tf.float32)).

If this weight were trainable=True by mistake, it would be pulled directly into the loss function's gradient computation. The Optimizer's entire job is to adjust every trainable variable in whatever direction reduces the loss — it has no concept of 'this number is just bookkeeping, please don't touch it.' It would treat the counter exactly like a weight matrix and drag its value around during backpropagation, corrupting the tally with values that have nothing to do with counting images.

Setting trainable=False is what tells Keras 'this variable holds state, not a parameter to be learned' — it stays out of trainable_variables, stays out of the gradient tape's default watch list, and only changes when your own code calls .assign() or .assign_add() on it directly.

At this point you have all three pieces of a working custom layer: __init__ to store configuration, build() to create weights sized to the actual input, and call() to define the forward pass using those weights. Together, that's enough to write any mathematical operation Keras doesn't already ship, and have it behave exactly like a built-in layer — droppable into Sequential, the functional API, or a subclassed Model.

The trainable flag you just worked through is the other half of the picture: it's what lets a custom layer hold state that persists and updates without being part of what the Optimizer learns, which is essential for anything beyond simple stateless math — running statistics, counters, moving averages, or cached lookups.

With build(), call(), and add_weight(trainable=...) under control, you're equipped to implement the custom layers that show up in real research code and production models, rather than being limited to whatever tf.keras.layers happens to already contain.

A standard Dense layer does simple math: output = activation(dot(input, weights) + bias). But what if you want to invent a completely new mathematical formula?

To do this, we create a Python Class that inherits from tf.keras.layers.Layer.

When creating a custom layer in TensorFlow, what base class MUST your new Python class inherit from?

• →tf.keras.Model
• →tf.keras.layers.Layer
• →tf.Module

Inside the class, you must define the build() method. This is where you tell TensorFlow to actually create the RAM variables (the weights and biases) for your layer.

What is the specific purpose of overriding the build(self, input_shape) method in a Custom Layer?

• →It defines the mathematical forward pass.
• →It is the exact place where you initialize and instantiate the layer's internal 'state' or 'weights' (e.g., using self.add_weight()) because it has access to the input_shape.
• →It calculates the loss.

Finally, you must define the call() method. This is the heart of the layer. It defines exactly what mathematical operations happen to the data when it passes through.

Which method holds the actual mathematical operations (the "forward pass") of your custom layer?

• →The __init__() method.
• →The build() method.
• →The call(self, inputs) method.

Now, prepare yourself. We are about to enter the ADA Defense Protocol. Ensure you understand training variables vs non-trainable variables.

When using self.add_weight(), there is a critical boolean flag: trainable. If it is True, the Optimizer will change this number. If False, the Optimizer ignores it.

ADA DEFENSE: You are building a custom layer that keeps a running tally of how many images have passed through it. You create a weight for this tally. Why MUST you set trainable=False for this specific weight?

• →Because you cannot train integers.
• →Because it is purely a tracking counter. If it were trainable=True, the Optimizer would ruthlessly alter the counter via backpropagation to minimize the Loss function, destroying your tally.
• →Because Keras only supports 1 trainable weight per layer.

Threat neutralized. Variable states confirmed. Custom logic securely integrated.

Build Real Custom Layer Weights. Finish build_custom_weights(): build() allocates the actual weight matrix.