Introduction

Combined tf.keras models become useful when one network should handle one kind of input and another network should handle a different kind. The important design choice is deciding whether the branch condition applies to the whole batch or to each example individually, because TensorFlow handles those two cases differently.

Build the Sub-Models First

A conditional model is easier to reason about when each branch is a normal Keras model with a clear input shape and output shape. Start with sub-models that can already run on their own.

1import tensorflow as tf
2from tensorflow import keras
3from tensorflow.keras import layers
4
5def build_branch(name: str) -> keras.Model:
6    inputs = keras.Input(shape=(4,), name=f"{name}_input")
7    x = layers.Dense(8, activation="relu")(inputs)
8    outputs = layers.Dense(1, name=f"{name}_output")(x)
9    return keras.Model(inputs, outputs, name=name)
10
11low_branch = build_branch("low_branch")
12high_branch = build_branch("high_branch")

Both branches return a single score, so they can be swapped in and out by a routing step without changing the downstream shape.

Use tf.cond for Whole-Batch Decisions

If the routing rule chooses one branch for the entire batch, use tf.cond inside a subclassed model. This is the closest match to ordinary conditional execution.

1class BatchConditionalModel(keras.Model):
2    def __init__(self):
3        super().__init__()
4        self.low_branch = build_branch("low")
5        self.high_branch = build_branch("high")
6
7    def call(self, inputs):
8        mean_feature = tf.reduce_mean(inputs[:, 0])
9        return tf.cond(
10            mean_feature < 0.0,
11            lambda: self.low_branch(inputs),
12            lambda: self.high_branch(inputs),
13        )
14
15model = BatchConditionalModel()
16batch = tf.constant(
17    [[-1.0, 0.2, 0.3, 0.4],
18     [-0.5, 0.1, 0.7, 0.9]],
19    dtype=tf.float32,
20)
21print(model(batch))

This works because there is one condition for the whole tensor. It is not appropriate when different rows in the same batch should go to different branches.

Route Individual Examples One Row at a Time

Per-example routing requires a different pattern. One practical option is to route each row through a branch with tf.map_fn, then stack the outputs back into a single tensor.

1class RoutedModel(keras.Model):
2    def __init__(self):
3        super().__init__()
4        self.low_branch = build_branch("masked_low")
5        self.high_branch = build_branch("masked_high")
6
7    def call(self, inputs):
8        def route_one(example):
9            example = tf.expand_dims(example, axis=0)
10            result = tf.cond(
11                example[0, 0] < 0.0,
12                lambda: self.low_branch(example),
13                lambda: self.high_branch(example),
14            )
15            return result[0]
16
17        return tf.map_fn(
18            route_one,
19            inputs,
20            fn_output_signature=tf.TensorSpec(shape=(1,), dtype=tf.float32),
21        )

The branch outputs keep the same final shape, while the internal routing remains flexible. This pattern is more verbose, but it models true per-example routing rather than pretending a batch-wide condition solves it.

Compile and Train Like a Normal Keras Model

Once the routing logic is inside call, the outer model can still use the normal Keras training flow.

1model = BatchConditionalModel()
2model.compile(optimizer="adam", loss="mse")
3
4x = tf.random.normal((32, 4))
5y = tf.random.normal((32, 1))
6model.fit(x, y, epochs=2, verbose=0)

That said, hard routing introduces optimization tradeoffs. If one branch is selected much more often than the other, the neglected branch may train poorly. In those cases, a learned gate or a softer mixture-of-experts design can behave better.

Common Pitfalls

• Using a Python if on a tensor inside call, which fails in graph execution because TensorFlow needs symbolic control flow.
• Sending branches with different output shapes into the same downstream pipeline, which makes the combined model impossible to use consistently.
• Assuming tf.cond gives per-example routing when it only selects one branch for the current tensor condition.
• Training with data that almost never activates one branch, leaving that sub-model effectively untrained.
• Forgetting that hard routing can make debugging harder because different inputs exercise different parameter sets.

Summary

• Build each branch as a normal Keras model first.
• Use tf.cond when one decision applies to the whole batch.
• Use masks and scatter updates when each example needs its own branch.
• Keep branch output shapes compatible so the outer model remains usable.
• Watch the training distribution so both branches receive enough signal.