keras/src/tree/optree_impl.py/_multiyield_flat_up_to
def _multiyield_flat_up_to(shallow_tree, *input_trees):
    """Same as `_yield_flat_up_to`, but takes multiple input trees."""
    zipped_iterators = zip(
        *[
            _yield_flat_up_to(shallow_tree, input_tree)
            for input_tree in input_trees
        ]
    )
    try:
        for paths_and_values in zipped_iterators:
            paths, values = zip(*paths_and_values)
            yield paths[:1] + values
    except KeyError as e:
        paths = locals().get("paths", ((),))
        raise ValueError(
            f"Could not find key '{e.args[0]}' in some `input_trees`. "
            "Please ensure the structure of all `input_trees` are "
            "compatible with `shallow_tree`. The last valid path "
            f"yielded was {paths[0]}."
        ) from e