Fix gh-18459: pre-build Sequential sublayers not reached through call()

Root cause: when a Sequential sublayer is only used inside a custom
`train_step` (not in `call()`), it is not built during `_symbolic_build`,
which only traces `call()`. For the TF backend without a distribute
strategy, `_maybe_symbolic_build()` previously returned early without
calling `_symbolic_build` at all. This caused the Sequential sublayer to
be built lazily during the first `tf.function` trace of `train_step`.
`Sequential.build()` creates KerasTensors and calls `compute_output_spec`,
which for TF creates a nested `FuncGraph` inside the already-active
`tf.function` context. When TF then tries to use a KerasTensor as a real
tensor, it hits `KerasTensor.__tf_tensor__` which raises:
  "A KerasTensor cannot be used as input to a TensorFlow function."

Fix (two parts):
1. `_symbolic_build` (trainer.py): after tracing `call()`, iterate over
   all sublayers and pre-build any that are still unbuilt by running
   `compute_output_spec` on the symbolic input `x`.
2. TF `_maybe_symbolic_build` (tensorflow/trainer.py): instead of always
   deferring when no distribute strategy is set, only defer when all
   layers are already built. If unbuilt sublayers exist, call
   `_symbolic_build` so they are built before `tf.function` traces
   `train_step`.

Also updates the regression test docstrings to accurately describe the
root cause (gh-18459).

Author: DeborahOlaboye

keras/src/backend/tensorflow/trainer.py

@@ -715,6 +715,15 @@ def _maybe_symbolic_build(self, iterator=None, data_batch=None):
             # When no distribution strategy is set, defer building
             # to when the train/test/predict function gets traced.
             # This maximizes backwards compatibility.
+            # Exception: if the model has unbuilt sublayers that won't be
+            # reached through `call()` (e.g. a Sequential sublayer only used
+            # in a custom `train_step`), we must pre-build them here.
+            # Otherwise `Sequential.build()` runs lazily inside `tf.function`
+            # tracing, where it creates a nested `FuncGraph` that is
+            # incompatible with the active `tf.function` context (gh-18459).
+            if all(layer.built for layer in self._flatten_layers()):
+                return
+            self._symbolic_build(iterator=iterator, data_batch=data_batch)
             return
 
         # Unlike jax/torch iterator, tf iterator returns an iterator instead

keras/src/trainers/trainer.py

@@ -1130,6 +1130,19 @@ def to_symbolic_input(v):
                     "Exception encountered:\n"
                     f"'{e}'"
                 )
+            # Pre-build any sublayers that were not reached through `call()`
+            # (e.g. a Sequential sublayer only used in a custom `train_step`).
+            # Without this, such layers would be built lazily during the first
+            # `tf.function` trace of `train_step`, where `Sequential.build()`
+            # creates a nested `FuncGraph` inside an active `tf.function`
+            # context, causing "A KerasTensor cannot be used as input to a
+            # TensorFlow function" (gh-18459).
+            for layer in self._flatten_layers():
+                if not layer.built:
+                    try:
+                        backend.compute_output_spec(layer, x)
+                    except Exception:
+                        pass
             if compile_metrics_unbuilt:
                 # Build all metric state with `backend.compute_output_spec`.
                 backend.compute_output_spec(

keras/src/trainers/trainer_test.py

@@ -82,11 +82,15 @@ def test_step(self, state, data):
 class SequentialSublayerInTrainStepModel(Trainer, layers.Layer):
     """Model that calls a Sequential sublayer only inside train_step.
 
-    Used as a regression test for GitHub issue #18459, which reported
-    that calling a Sequential model from a custom train_step during
-    model.fit() incorrectly received KerasTensors instead of real
-    tensors, causing: "A KerasTensor cannot be used as input to a
-    TensorFlow function."
+    Used as a regression test for GitHub issue #18459. The bug: a
+    Sequential sublayer used only in train_step (not in call()) is not
+    built during _symbolic_build, which only traces call(). For the TF
+    backend without a distribute strategy, _maybe_symbolic_build() used
+    to exit early entirely, so the sublayer was built lazily during the
+    first tf.function trace of train_step. Sequential.build() creates
+    KerasTensors and calls compute_output_spec() which for TF creates a
+    nested FuncGraph inside the active tf.function context, causing:
+    "A KerasTensor cannot be used as input to a TensorFlow function."
     """
 
     def __init__(self, units):
@@ -791,8 +795,10 @@ def test_fit_with_custom_train_step(self):
     def test_sequential_sublayer_in_custom_train_step(self):
         """Regression test for GitHub issue #18459.
 
-        Calling a Sequential sublayer from a custom train_step during
-        model.fit() must not raise a KerasTensor-related error.
+        A Sequential sublayer used only in train_step must be pre-built by
+        _symbolic_build before tf.function traces train_step. Without the
+        fix, Sequential.build() ran lazily inside the tf.function context,
+        creating a nested FuncGraph that caused a KerasTensor error.
         """
         if backend.backend() == "jax":
             self.skipTest(