Introduction

Yes, you can split a TensorFlow network into two pieces and still preserve gradient back-propagation, as long as the intermediate value remains an ordinary TensorFlow tensor in the same differentiable graph. The split itself is not the problem. Gradients break only when you convert the intermediate result to NumPy, wrap it as a fresh constant, or explicitly stop gradients.

The Core Rule

Gradient flow is preserved when:

• part one produces a tensor
• part two consumes that tensor
• the full forward pass stays inside TensorFlow ops tracked by tf.GradientTape or the Keras model graph

Gradient flow is broken when you do things such as:

• 'tensor.numpy() inside training'
• 'tf.stop_gradient(...)'
• rebuilding the intermediate output as tf.constant(...)
• moving values through non-differentiable external code

That distinction matters more than whether the model is stored in one object or two.

Simple Two-Part Model Example

1import tensorflow as tf
2
3part1 = tf.keras.Sequential([
4    tf.keras.layers.Dense(16, activation="relu"),
5    tf.keras.layers.Dense(8, activation="relu"),
6])
7
8part2 = tf.keras.Sequential([
9    tf.keras.layers.Dense(4, activation="relu"),
10    tf.keras.layers.Dense(1),
11])
12
13x = tf.random.normal((5, 10))
14y = tf.random.normal((5, 1))
15
16with tf.GradientTape() as tape:
17    z = part1(x)
18    y_pred = part2(z)
19    loss = tf.reduce_mean(tf.square(y - y_pred))
20
21variables = part1.trainable_variables + part2.trainable_variables
22grads = tape.gradient(loss, variables)
23
24print([g is not None for g in grads])

Gradients flow through z automatically because z is just an intermediate tensor in the same differentiable computation.

Functional API Version

If you want a clean model split for reuse, the Keras Functional API is often the best fit.

1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(10,))
4x = tf.keras.layers.Dense(16, activation="relu")(inputs)
5bridge = tf.keras.layers.Dense(8, activation="relu", name="bridge")(x)
6x = tf.keras.layers.Dense(4, activation="relu")(bridge)
7outputs = tf.keras.layers.Dense(1)(x)
8
9full_model = tf.keras.Model(inputs, outputs)
10encoder = tf.keras.Model(inputs, bridge)
11
12bridge_input = tf.keras.Input(shape=(8,))
13decoder_output = full_model.layers[-2](bridge_input)
14decoder_output = full_model.layers[-1](decoder_output)
15decoder = tf.keras.Model(bridge_input, decoder_output)

Now you have two reusable parts while still understanding that, during training, the actual bridge tensor must remain connected in one TensorFlow computation.

What Breaks The Gradient

This is the classic mistake:

1with tf.GradientTape() as tape:
2    z = part1(x)
3    z_numpy = z.numpy()
4    z_again = tf.constant(z_numpy)
5    y_pred = part2(z_again)
6    loss = tf.reduce_mean(tf.square(y - y_pred))

Why this breaks training:

• 'z.numpy() leaves TensorFlow's differentiable graph'
• 'tf.constant(z_numpy) creates a new leaf tensor unrelated to part1'
• the tape can no longer trace the path back into part1

If your gradients for the first half are None, this kind of graph break is the first thing to check.

Training The Parts With One Or Two Optimizers

You can still use separate optimizers if needed. The important part is that the loss is computed from the connected forward pass.

1opt1 = tf.keras.optimizers.Adam(1e-3)
2opt2 = tf.keras.optimizers.Adam(1e-3)
3
4with tf.GradientTape() as tape:
5    z = part1(x, training=True)
6    y_pred = part2(z, training=True)
7    loss = tf.reduce_mean(tf.square(y - y_pred))
8
9grads1 = tape.gradient(loss, part1.trainable_variables)
10grads2 = tape.gradient(loss, part2.trainable_variables)
11
12opt1.apply_gradients(zip(grads1, part1.trainable_variables))
13opt2.apply_gradients(zip(grads2, part2.trainable_variables))

In practice, many people compute all gradients in one call and split them later, but the principle is the same.

When tf.stop_gradient Is Useful

Sometimes you intentionally want to cut the backward path.

z = tf.stop_gradient(part1(x))
y_pred = part2(z)

This freezes learning into part1. That can be correct for staged training or frozen feature extractors, but it is the opposite of maintaining back-propagation across the split.

Common Pitfalls

• Converting the intermediate tensor to NumPy during training.
• Re-wrapping the intermediate output as a new constant or placeholder.
• Assuming separate model objects automatically break gradients. They do not if tensors stay connected.
• Using tf.stop_gradient without realizing it cuts the backward path intentionally.
• Debugging the optimizer first when the real problem is that the graph was broken between the two halves.

Summary

• Splitting a TensorFlow network into two parts does not break gradients by itself.
• Back-propagation is preserved as long as the intermediate output stays a connected TensorFlow tensor.
• 'tf.GradientTape and the Keras Functional API both support this pattern naturally.'
• Gradients break when you convert tensors to NumPy, recreate them as constants, or call tf.stop_gradient.
• If the first half gets no gradients, inspect the bridge tensor path before changing the optimizer or architecture.