Introduction

Keras models run on top of TensorFlow, so TensorFlow operations can be part of a model when you use them in the right place. The safest pattern is to wrap raw tf logic inside a Keras layer, loss, metric, or training step instead of scattering TensorFlow calls around the model definition.

Use TensorFlow Ops Inside a Custom Layer

When you need an operation that Keras does not provide as a built-in layer, define a small custom layer and place the TensorFlow code in call.

1import tensorflow as tf
2from tensorflow import keras
3
4class L2Normalize(keras.layers.Layer):
5    def call(self, inputs):
6        return tf.math.l2_normalize(inputs, axis=-1)
7
8inputs = keras.Input(shape=(8,))
9x = keras.layers.Dense(16, activation="relu")(inputs)
10x = L2Normalize()(x)
11outputs = keras.layers.Dense(1)(x)
12
13model = keras.Model(inputs, outputs)
14model.compile(optimizer="adam", loss="mse")

This approach works well because Keras tracks the layer as part of the model graph. You can save the model, inspect the layer stack, and keep custom logic isolated in one clear component.

Use TensorFlow Ops in a Lambda Layer for Small Cases

For very small one-off transformations, keras.layers.Lambda can be enough:

1inputs = keras.Input(shape=(4,))
2x = keras.layers.Lambda(lambda t: tf.math.log1p(tf.abs(t)))(inputs)
3outputs = keras.layers.Dense(1)(x)
4
5model = keras.Model(inputs, outputs)

This is convenient, but use it sparingly. A named custom layer is usually easier to debug and serialize, especially once the transformation becomes more complex than a single expression.

TensorFlow Ops Also Belong in Custom Losses and Metrics

Model code is not limited to forward-pass layers. Raw TensorFlow operations are often a good fit inside custom losses or metrics.

1def clipped_mae(y_true, y_pred):
2    error = tf.abs(y_true - y_pred)
3    return tf.reduce_mean(tf.minimum(error, 1.0))
4
5model.compile(
6    optimizer="adam",
7    loss=clipped_mae,
8    metrics=[keras.metrics.MeanAbsoluteError()],
9)

This works because Keras passes tensors into the loss during training, and TensorFlow handles the differentiation automatically for standard ops.

Know When to Prefer keras.ops

If your project is tightly tied to the TensorFlow backend, raw tf usage is fine. If you want backend portability in newer Keras versions, prefer keras.ops for generic mathematical operations. The model will then be easier to move across supported backends later.

A practical rule is:

• Use raw tf for TensorFlow-specific behavior, custom training loops, and low-level features.
• Use Keras layers or keras.ops when an equivalent high-level abstraction exists.

A Full Example with a Custom TensorFlow Operation

The example below clips activations with TensorFlow before the final prediction:

1import tensorflow as tf
2from tensorflow import keras
3
4class ClipValues(keras.layers.Layer):
5    def __init__(self, min_value, max_value):
6        super().__init__()
7        self.min_value = min_value
8        self.max_value = max_value
9
10    def call(self, inputs):
11        return tf.clip_by_value(inputs, self.min_value, self.max_value)
12
13inputs = keras.Input(shape=(6,))
14x = keras.layers.Dense(12, activation="relu")(inputs)
15x = ClipValues(-0.5, 0.5)(x)
16outputs = keras.layers.Dense(1)(x)
17
18model = keras.Model(inputs, outputs)
19model.compile(optimizer="adam", loss="mse")

This pattern keeps the model readable while still giving you access to TensorFlow functionality that does not map cleanly to a stock Keras layer.

Common Pitfalls

One common mistake is calling a TensorFlow function directly on a symbolic Keras tensor in random Python code and expecting Keras to track it correctly. Wrapping the logic in a layer avoids most of that confusion.

Another issue is relying on Lambda layers for large chunks of business logic. They are quick to write, but custom named layers are easier to test, reuse, and save.

Shape handling is another frequent source of bugs. Many TensorFlow ops reduce dimensions or change rank, so verify the tensor shape that leaves your custom layer before connecting it to the next layer.

Finally, avoid Python control flow based on tensor values inside call. Use TensorFlow control-flow ops such as tf.where when the decision depends on tensors.

Summary

• Put raw TensorFlow operations inside custom Keras layers, losses, metrics, or training steps.
• Use Lambda only for small transformations; prefer named custom layers for anything substantial.
• Keep an eye on tensor shapes when combining TensorFlow ops with standard Keras layers.
• Use keras.ops when backend portability matters and TensorFlow-specific behavior is not required.
• Wrapping TensorFlow logic cleanly makes models easier to save, debug, and maintain.