keras/src/tree/optree_impl.py/map_structure_up_to
def map_structure_up_to(shallow_structure, func, *structures):
    return _map_structure_with_path_up_to(
        shallow_structure,
        lambda _, *args: func(*args),  # Discards path.
        *structures,
    )