Once Layer.__call__'s CallSpec fast path has established that a call is a single leaf tensor (at most training=<bool>), later steps in the same call still re-derive that already-known structural fact using generic tree machinery. Specifically: the mask pipeline's tree.map_structure(get_keras_mask, ...) calls, an any_symbolic_tensors re-walk over the inputs, and a general-path tree.flatten inside _set_mask_metadata. Each of these re-walks a structure that the fast path already proved is a single leaf, so for the (very common) single-tensor-input case they're pure overhead paid on every call.
The fix hoists a single-tensor-call predicate and adds leaf fast branches ahead of the tree-based mask handling, so single-tensor calls skip the redundant walks entirely; the general nested-structure paths are untouched for multi-input/multi-output layers. Fixed by #23300.
Part of the per-call Python-dispatch overhead series in #22561.
Once
Layer.__call__'s CallSpec fast path has established that a call is a single leaf tensor (at mosttraining=<bool>), later steps in the same call still re-derive that already-known structural fact using generic tree machinery. Specifically: the mask pipeline'stree.map_structure(get_keras_mask, ...)calls, anany_symbolic_tensorsre-walk over the inputs, and a general-pathtree.flatteninside_set_mask_metadata. Each of these re-walks a structure that the fast path already proved is a single leaf, so for the (very common) single-tensor-input case they're pure overhead paid on every call.The fix hoists a single-tensor-call predicate and adds leaf fast branches ahead of the tree-based mask handling, so single-tensor calls skip the redundant walks entirely; the general nested-structure paths are untouched for multi-input/multi-output layers. Fixed by #23300.
Part of the per-call Python-dispatch overhead series in #22561.