Python
Keras
Neural Networks
Identity Layer
Machine Learning

Python Keras An layer output exactly the same thing as input

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

If you want a Keras layer whose output is exactly the same as its input, you want an identity mapping. The simplest way to do that is to return the incoming tensor unchanged with a Lambda layer or a small custom layer.

This is useful for debugging, building optional branches, or keeping a model graph consistent while one path does no transformation. The important part is choosing an implementation that is easy to understand and safe to save and reload.

Use a Lambda Identity Layer

For a quick pass-through layer, Lambda is the shortest solution.

python
1import tensorflow as tf
2from tensorflow.keras import layers, Model
3
4inputs = layers.Input(shape=(4,), name="features")
5identity = layers.Lambda(lambda x: x, name="identity")(inputs)
6model = Model(inputs, identity)
7
8sample = tf.constant([[1.0, 2.0, 3.0, 4.0]])
9print(model(sample).numpy())

The output is the same tensor content as the input. This works well when you need an explicit node in the graph but no computation.

Use tf.identity for Clarity

You can also make the intent more explicit by using tf.identity inside the lambda.

python
1import tensorflow as tf
2from tensorflow.keras import layers, Model
3
4inputs = layers.Input(shape=(2,), name="values")
5outputs = layers.Lambda(tf.identity, name="identity")(inputs)
6model = Model(inputs, outputs)
7
8print(model(tf.constant([[5.0, 6.0]])).numpy())

This reads more clearly than a bare lambda x: x, especially for readers who immediately recognize tf.identity as a pass-through operation.

Custom Layer for Better Serialization

Lambda is convenient, but custom layers are often easier to maintain in larger projects because their behavior is explicit and easier to serialize cleanly.

python
1import tensorflow as tf
2from tensorflow.keras import layers, Model
3
4class IdentityLayer(layers.Layer):
5    def call(self, inputs):
6        return inputs
7
8inputs = layers.Input(shape=(3,), name="features")
9outputs = IdentityLayer(name="identity_layer")(inputs)
10model = Model(inputs, outputs)
11
12sample = tf.constant([[7.0, 8.0, 9.0]])
13print(model(sample).numpy())

A custom layer is a good choice when the identity step is part of a reusable model component or when you care about long-term portability.

Useful Cases for Identity Layers

Identity mappings are not just toy examples. They are common in a few real situations:

  • debugging intermediate tensors in a model graph
  • conditionally swapping in a real transformation later
  • keeping model branches structurally aligned
  • implementing residual-style designs where one path is unchanged

For example, an unchanged branch can be combined with a transformed branch:

python
1import tensorflow as tf
2from tensorflow.keras import layers, Model
3
4inputs = layers.Input(shape=(4,))
5identity = layers.Lambda(tf.identity)(inputs)
6transformed = layers.Dense(4, activation="relu")(inputs)
7outputs = layers.Add()([identity, transformed])
8model = Model(inputs, outputs)
9
10print(model(tf.ones((1, 4))).shape)

Here the identity path preserves the original features while another branch learns an adjustment.

Remember That Some Layers Are Only Initially Identity-Like

A dense layer with linear activation is not automatically an identity mapping. It can learn one, but only if its weights become an identity matrix and its bias is zero. That is not the same as saying it already outputs exactly the input.

So if the requirement is "exactly the same thing as input," use a true pass-through layer rather than a trainable layer that might approximate identity under special weights.

Common Pitfalls

  • Using a trainable dense layer and assuming it is an identity mapping by default.
  • Hiding important behavior inside a complex Lambda instead of keeping the pass-through explicit.
  • Forgetting that custom or lambda-based layers may affect model save and load behavior.
  • Combining identity and transformed branches with mismatched shapes.
  • Calling something "identity" when it actually changes dtype, shape, or values upstream.

Summary

  • For exact pass-through behavior in Keras, use a Lambda layer or a tiny custom layer that returns inputs.
  • 'tf.identity makes the intent especially clear.'
  • Use a custom layer when you want better reuse and more explicit serialization behavior.
  • Do not confuse trainable linear layers with true identity mappings.
  • Check shape and graph design carefully when identity paths are combined with other branches.

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.