Introduction

You can use a Keras layer directly inside ordinary TensorFlow code because a tf.keras.layers.Layer is just a callable object that takes tensors and returns tensors. That makes Keras layers useful not only in Sequential and functional models, but also in custom training loops, tf.Module classes, and lower-level TensorFlow programs.

A Keras Layer Is Still Just TensorFlow Code

At runtime, a Keras layer receives TensorFlow tensors, creates weights when needed, and applies TensorFlow operations in its call logic. That means you can instantiate a layer and invoke it directly:

1import tensorflow as tf
2
3dense = tf.keras.layers.Dense(4, activation="relu")
4
5x = tf.constant([[1.0, 2.0, 3.0]])
6y = dense(x)
7
8print(y)
9print(dense.weights)

The first time you call the layer, it builds its weights based on the input shape. After that, it behaves like a reusable TensorFlow component.

Use Layers in Custom TensorFlow Classes

One common pattern is to place Keras layers inside a tf.Module or a custom model class.

1import tensorflow as tf
2
3class MyNetwork(tf.Module):
4    def __init__(self):
5        super().__init__()
6        self.hidden = tf.keras.layers.Dense(8, activation="relu")
7        self.out = tf.keras.layers.Dense(1)
8
9    def __call__(self, x):
10        x = self.hidden(x)
11        return self.out(x)
12
13
14model = MyNetwork()
15inputs = tf.random.normal((5, 3))
16outputs = model(inputs)
17print(outputs.shape)

This lets you keep Keras convenience for layers while still controlling the outer TensorFlow structure yourself.

Combine Keras Layers With a Gradient Tape

Keras layers also work naturally inside custom training loops.

1import tensorflow as tf
2
3dense = tf.keras.layers.Dense(1)
4optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)
5
6x = tf.constant([[1.0], [2.0], [3.0]])
7y_true = tf.constant([[2.0], [4.0], [6.0]])
8
9for _ in range(5):
10    with tf.GradientTape() as tape:
11        y_pred = dense(x)
12        loss = tf.reduce_mean(tf.square(y_true - y_pred))
13
14    grads = tape.gradient(loss, dense.trainable_variables)
15    optimizer.apply_gradients(zip(grads, dense.trainable_variables))
16
17print("loss:", float(loss))

This is one of the most useful mixed-style patterns in TensorFlow: Keras layers for state and reusable building blocks, plain TensorFlow for the training loop.

Why This Works Well

Keras layers save you from manually creating variables and wiring common layer behaviors. TensorFlow still gives you direct control over gradients, graph tracing, and deployment paths.

That combination is especially useful when:

• you want custom training logic
• you need reusable trainable components
• a full high-level Keras model API feels too restrictive

Layers Also Work Inside tf.function

If you trace your computation with tf.function, you can still use the same layer object:

1import tensorflow as tf
2
3dense = tf.keras.layers.Dense(2)
4
5@tf.function
6def run_layer(x):
7    return dense(x)
8
9print(run_layer(tf.ones((1, 3))))

The important rule is to create the layer once outside the traced function, then reuse it. That keeps the variables stable and avoids recreating state on each call. It is one of the easiest mistakes to prevent. Just build once.

Common Pitfalls

• Layers do not create weights until they see input for the first time, so inspecting variables before the first call can be confusing.
• Recreating a layer inside a loop creates new weights repeatedly; instantiate it once and reuse it.
• Keras layers expect tensors with compatible shapes, so input rank mismatches still matter even in low-level TensorFlow code.
• Mixing eager-style debugging and graph-traced functions is fine, but be consistent about where stateful layers are created.

Summary

• A Keras layer can be used directly in ordinary TensorFlow code because it is a callable tensor-processing object.
• You can place tf.keras.layers inside custom classes, modules, and training loops.
• Keras layers work well with tf.GradientTape and TensorFlow optimizers.
• Instantiate layers once, reuse them, and let the first call build their weights.